Continuous Thought Machines (CTM)

The models/ directory contains the core implementation of the Continuous Thought Machine. The models utilize a modified ResNet architecture (defined in resnet.py) designed to encourage information-gathering behavior.

To prevent the model from finding a 'path of least resistance' via SGD that bypasses intelligent behavior, the ResNet structure includes modifications to constrain the receptive field of the yielded features. This constraint forces the CTM or baseline methods to learn an active process of gathering information.

The repository is organized into the following main components:

• tasks/: Contains task-specific code (training, analysis, plotting, and scripts) for image classification, mazes, sorting, parity, QAMNIST, and RL.
• models/: Core model implementations including ctm.py (main model), baseline models (ff.py, lstm.py), and helper modules (modules.py).
• data/: Storage for downloaded and custom datasets.
• utils/: Shared utility functions for housekeeping, loss functions, and learning rate schedulers.
• checkpoints/: Directory for storing model checkpoints.

Training code is organized by task within the tasks/ directory. Each task contains its own train.py script. To facilitate running multiple high-level training scripts from the top-level directory, training scripts should be executed as modules.

While argparsers provide reasonable defaults, scripts to replicate the exact setups used in the paper are located in the scripts/ subdirectories of each task folder.

Example for image classification:

python -m tasks.image_classification.train

To set up the environment using conda, create a new environment with Python 3.12, activate it, and install the required dependencies via requirements.txt.

If you encounter issues with PyTorch versions, you may need to manually uninstall the existing torch and install the version compatible with CUDA 12.1.

Note: ffmpeg is required for generating .mp4 files from analysis scripts. You can install it via conda-forge.

conda create --name=ctm python=3.12
conda activate ctm
pip install -r requirements.txt

# If PyTorch issues occur:
pip uninstall torch
pip install torch --index-url https://download.pytorch.org/whl/cu121

# Install ffmpeg for analysis video generation:
conda install -c conda-forge ffmpeg

To replicate the RL training used in the paper, execute the provided bash scripts from the root level of the repository. For example, to train a 2-iteration Continuous Thought Machine (CTM) on the Acrobot task, use the specific script for that task.

Training progress can be monitored using TensorBoard. Ensure TensorBoard is installed in your environment.

The image classification tasks use the ILSVRC/imagenet-1k dataset via Hugging Face. To enable automatic downloading and loading during experiments, follow these steps:

1. Create a Hugging Face account and agree to the dataset's Terms and Conditions.
2. Generate a new Hugging Face access token.
3. Install the huggingface_hub library:

   pip install huggingface_hub

4. Authenticate your machine using the CLI:

   huggingface-cli login

   (Paste your token when prompted).
5. Run your ImageNet experiment; the code will handle the download automatically.

To perform analysis on the Q&A MNIST task, use the tasks.qamnist.analysis.run module. This requires checkpoints to be present in the directory specified by the --log_dir argument.

Checkpoints can be generated by running the training scripts or downloaded externally. Ensure the path provided to --log_dir contains the necessary checkpoint files.

python -m tasks.qamnist.analysis.run --log_dir <PATH_TO_LOG_DIR>

To perform analysis on trained models, you must provide a path to a directory containing saved checkpoints. This directory is specified via the --log_dir argument.

Checkpoints can be generated by running the training scripts or downloaded manually from the project's provided storage link. Run the analysis module using the following command:

python -m tasks.rl.analysis.run --log_dir <PATH_TO_LOG_DIR>

The repository includes analysis and plotting code to replicate results from the paper. Analysis entry points are located in the analysis/ subdirectory of each task (e.g., tasks/image_classification/analysis/run_imagenet_analysis.py or tasks/mazes/analysis/run.py).

Requirement: ffmpeg must be installed to generate .mp4 files from these scripts.

To replicate the Q&A MNIST training used in the paper, execute the provided bash scripts from the root level of the repository. For example, to train a Continuous Thought Machine (CTM) with 10 iterations, use the train_ctm_10.sh script.

bash tasks/qamnist/scripts/train_ctm_10.sh

To use pre-trained models or specific datasets, download them from the provided external links and place them in the appropriate local directories:

• Checkpoints: Place in the checkpoints/ folder (e.g., checkpoints/mazes/...).
• Maze Data: Download and place in the data/ directory.

External Links:

• Checkpoints: Google Drive Link
• Maze Data: Google Drive Link