π PAPER: Technical Report | π Blog | πΉοΈ Interactive Website | βοΈ Tutorial
We present the Continuous Thought Machine (CTM), a model designed to unfold and then leverage neural activity as the underlying mechanism for observation and action. Our contributions are:
-
An internal temporal axis, decoupled from any input data, that enables neuron activity to unfold.
-
Neuron-level temporal processing, where each neuron uses unique weight parameters to process a history of incoming signals, enabling fine-grained temporal dynamics.
-
Neural synchronisation, employed as a direct latent representation for modulating data and producing outputs, thus directly encoding information in the timing of neural activity.
We demonstrate the CTM's strong performance and versatility across a range of challenging tasks, including ImageNet classification, solving 2D mazes, sorting, parity computation, question-answering, and RL tasks.
We provide all necessary code to reproduce our results and invite others to build upon and use CTMs in their own work.
Please see our Interactive Website for a maze-solving demo, many demonstrative videos of the method, results, and other findings.
βββ tasks
βΒ Β βββ image_classification
βΒ Β βΒ Β βββ train.py # Training code for image classification (cifar, imagenet)
βΒ Β βΒ Β βββ imagenet_classes.py # Helper for imagenet class names
βΒ Β βΒ Β βββ plotting.py # Plotting utils specific to this task
βΒ Β βΒ Β βββ analysis
βΒ Β βΒ Β βββrun_imagenet_analysis.py # ImageNet eval and visualisation code
βΒ Β βΒ Β Β Β βββoutputs/ # Folder for outputs of analysis
βΒ Β βββ mazes
βΒ Β βΒ Β βββ train.py # Training code for solving 2D mazes (by way of a route; see paper)
βΒ Β βΒ Β βββ plotting.py # Plotting utils specific to this task
βΒ Β βΒ Β βββ analysis
βΒ Β βΒ Β βββrun.py # Maze analysis code
βΒ Β βΒ Β Β Β βββoutputs/ # Folder for outputs of analysis
βΒ Β βββ sort
βΒ Β βΒ Β βββ train.py # Training code for sorting
βΒ Β βΒ Β βββ utils.py # Sort specific utils (e.g., CTC decode)
βΒ Β βββ parity
βΒ Β βΒ Β βββ train.py # Training code for parity task
βΒ Β βΒ Β βββ utils.py # Parity-specific helper functions
βΒ Β βΒ Β βββ plotting.py # Plotting utils specific to this task
βΒ Β βΒ Β βββ scripts/
βΒ Β βΒ Β βΒ Β βββ *.sh # Training scripts for different experimental setups
βΒ Β βΒ Β βββ analysis/
βΒ Β βΒ Β βββ run.py # Entry point for parity analysis
βΒ Β βββ qamnist
βΒ Β βΒ Β βββ train.py # Training code for QAMNIST task (quantized MNIST)
βΒ Β βΒ Β βββ utils.py # QAMNIST-specific helper functions
βΒ Β βΒ Β βββ plotting.py # Plotting utils specific to this task
βΒ Β βΒ Β βββ scripts/
βΒ Β βΒ Β βΒ Β βββ *.sh # Training scripts for different experimental setups
βΒ Β βΒ Β βββ analysis/
βΒ Β βΒ Β βββ run.py # Entry point for QAMNIST analysis
βΒ Β βββ rl
βΒ Β Β Β βββ train.py # Training code for RL environments
βΒ Β Β Β βββ utils.py # RL-specific helper functions
βΒ Β Β Β βββ plotting.py # Plotting utils specific to this task
βΒ Β Β Β βββ envs.py # Custom RL environment wrappers
βΒ Β Β Β βββ scripts/
βΒ Β Β Β βΒ Β βββ 4rooms/
βΒ Β Β Β βΒ Β βΒ Β βββ *.sh # Training scripts for MiniGrid-FourRooms-v0 environment
βΒ Β Β Β βΒ Β βββ acrobot/
βΒ Β Β Β βΒ Β βΒ Β βββ *.sh # Training scripts for Acrobot-v1 environment
βΒ Β Β Β βΒ Β βββ cartpole/
βΒ Β Β Β βΒ Β βββ *.sh # Training scripts for CartPole-v1 environment
βΒ Β Β Β βββ analysis/
βΒ Β Β Β βββ run.py # Entry point for RL analysis
βββ data # This is where data will be saved and downloaded to
βΒ Β βββ custom_datasets.py # Custom datasets (e.g., Mazes), sort
βββ models
βΒ Β βββ ctm.py # Main model code, used for: image classification, solving mazes, sort
βΒ Β βββ ctm_*.py # Other model code, standalone adjustments for other tasks
βΒ Β βββ ff.py # feed-forward (simple) baseline code (e.g., for image classification)
βΒ Β βββ lstm.py # LSTM baseline code (e.g., for image classification)
βΒ Β βββ lstm_*.py # Other baseline code, standalone adjustments for other tasks
βΒ Β βββ modules.py # Helper modules, including Neuron-level models and the Synapse UNET
βΒ Β βββ utils.py # Helper functions (e.g., synch decay)
βΒ Β βββ resnet.py # Wrapper for ResNet featuriser
βββ utils
βΒ Β βββ housekeeping.py # Helper functions for keeping things neat
βΒ Β βββ losses.py # Loss functions for various tasks (mostly with reshaping stuff)
βΒ Β βββ schedulers.py # Helper wrappers for learning rate schedulers
βββ checkpoints
Β Β βββ imagenet, mazes, ... # Checkpoint directories (see google drive link for files)
To set up the environment using conda:
conda create --name=ctm python=3.12
conda activate ctm
pip install -r requirements.txt
If there are issues with PyTorch versions, the following can be ran:
pip uninstall torch
pip install torch --index-url https://download.pytorch.org/whl/cu121
Each task has its own (set of) training code. See for instance tasks/image_classification/train.py. We have set it up like this to ensure ease-of-use as opposed to clinical efficiency. This code is for researchers and we hope to have it shared in a way that fosters collaboration and learning.
While we have provided reasonable defaults in the argparsers of each training setup, scripts to replicate the setups in the paper will typically be found in the accompanying script folders. If you simply want to dive in, run the following as a module (setup like this to make it easy to run many high-level training scripts from the top directory):
python -m tasks.image_classification.train
For debugging in VSCode, this configuration example might be helpful to you:
{
"name": "Debug: train image classifier",
"type": "debugpy",
"request": "launch",
"module": "tasks.image_classification.train",
"console": "integratedTerminal",
"justMyCode": false
}
We also provide analysis and plotting code to replicate many of the plots in our paper. See tasks/.../analysis/* for more details on that. We also provide some data (e.g., the mazes we generated for training) and checkpoints (see here). Note that ffmpeg is required for generating mp4 files from the analysis scripts. It can be installed with:
conda install -c conda-forge ffmpeg
You can download the data and checkpoints from here:
- checkpoints: https://drive.google.com/drive/folders/1vSg8T7FqP-guMDk1LU7_jZaQtXFP9sZg
- maze data: https://drive.google.com/file/d/1cBgqhaUUtsrll8-o2VY42hPpyBcfFv86/view?usp=drivesdk
Checkpoints go in the checkpoints folder. For instance, when properly populated, the checkpoints folder will have the maze checkpoint in checkpoints/mazes/...
(Update: 29 December 2025)
Please note that all pull requests will be assessed for their contributions. Those that involve either a bug-fix or an efficiency upgrade will likely pass review.
On the other hand, any pull requests that change the fundamental functionality of this repository will be immediately rejected and closed: the purpose of this repository is to accompany the Continuous Thought Machines paper and, therefore, it must remain an accurate reflection of that work.
That being said, we are happy to consider advances and deviations from the original repository as long as they:
- Are sufficiently self-contained (single feature additions per pull request).
- Well-justified with quantitative evidence.
- Not overclaiming to be fundamental.
- Do not change the default behavior of the original codebase. For example, while early stopping is useful, this should be set up in a way that is an optional extra, as opposed to the default behavior.