--- title: Continuous Thought Machine - Energy based Halting emoji: 🕰ī¸ colorFrom: blue colorTo: indigo sdk: docker sdk_version: "20.10.21" app_file: app.py pinned: false --- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference # 🕰ī¸ The Continuous Thought Machine 📚 [PAPER: Technical Report](https://arxiv.org/abs/2505.05522) | 📝 [Blog](https://sakana.ai/ctm/) | 🕹ī¸ [Interactive Website](https://pub.sakana.ai/ctm) | ✏ī¸ [Tutorial](examples/01_mnist.ipynb) ## Overview The **Continuous Thought Machine (CTM)** is a novel neural architecture designed to unfold and leverage neural activity as the underlying mechanism for observation and action. By introducing an internal temporal axis decoupled from input data, CTM enables neurons to process information over time with fine-grained temporal dynamics. ### Key Contributions 1. **Internal Temporal Axis**: Decoupled from input data, allowing neuron activity to unfold independently 2. **Neuron-Level Temporal Processing**: Each neuron uses unique weight parameters to process a history of incoming signals 3. **Neural Synchronisation**: Direct latent representation for modulating data and producing outputs, encoding information in the timing of neural activity The CTM demonstrates strong performance across diverse tasks including ImageNet classification, 2D maze solving, sorting, parity computation, question-answering, and reinforcement learning. --- ## đŸ”Ŧ Energy-Based Halting Experiment This repository includes an implementation of **Energy-Based Halting**, a mechanism that frames "thinking" as an optimization process where the model dynamically adjusts its internal thought process duration based on sample difficulty. ### Concept Instead of using heuristic certainty thresholds, we train a learned energy scalar that: - **Minimizes energy** for correct predictions (pushing the system to low-energy equilibrium) - **Maximizes energy** for incorrect predictions (pushing away from stable states) - **Enables adaptive halting** based on energy thresholds or convergence ### Implementation **Modified Components:** - `models/ctm.py`: Added energy projection head that maps synchronization states to scalar energy values - `utils/losses.py`: Implemented `EnergyContrastiveLoss` for training the energy function - `tasks/image_classification/train_energy.py`: Training script with energy halting - `inference_energy.py`: Adaptive inference that halts when energy drops below threshold or stabilizes - `configs/energy_experiment.yaml`: Configuration for energy experiments **Training:** ```bash # Local training pixi run accelerate launch tasks/image_classification/train_energy.py \ --energy_head_enabled \ --loss_type energy_contrastive \ --dataset cifar10 # Or with traditional python pixi run python tasks/image_classification/train_energy.py \ --energy_head_enabled \ --loss_type energy_contrastive ``` **Deployment to Hugging Face:** See [GUIDE_HF.md](GUIDE_HF.md) for instructions on deploying the training job to Hugging Face Spaces with GPU support. --- ## 🚀 Quick Start ### Setup with Pixi (Recommended) We use [Pixi](https://pixi.sh) for dependency management, which handles both Python packages and system dependencies like `ffmpeg`. ```bash # Install dependencies pixi install # Run training pixi run python tasks/image_classification/train.py ``` ### Alternative: Conda Setup ```bash conda create --name=ctm python=3.12 conda activate ctm pip install -r requirements.txt conda install -c conda-forge ffmpeg ``` If there are PyTorch version issues: ```bash pip uninstall torch pip install torch --index-url https://download.pytorch.org/whl/cu121 ``` --- ## 📁 Repository Structure ``` ├── tasks/ │ ├── image_classification/ │ │ ├── train.py # Standard training │ │ ├── train_energy.py # Energy halting training │ │ ├── analysis/run_imagenet_analysis.py │ │ └── plotting.py │ ├── mazes/ │ │ ├── train.py │ │ └── analysis/ │ ├── sort/ │ ├── parity/ │ ├── qamnist/ │ └── rl/ ├── models/ │ ├── ctm.py # Main CTM model (with energy head support) │ ├── modules.py # Neuron-level models, Synapse UNET │ ├── ff.py # Feed-forward baseline │ └── lstm.py # LSTM baseline ├── utils/ │ ├── losses.py # Loss functions (includes EnergyContrastiveLoss) │ ├── schedulers.py │ └── housekeeping.py ├── data/ │ └── custom_datasets.py ├── configs/ │ └── energy_experiment.yaml # Energy halting hyperparameters ├── inference_energy.py # Adaptive energy-based inference ├── Dockerfile # For HF Spaces deployment ├── GUIDE_HF.md # Hugging Face deployment guide └── checkpoints/ # Model checkpoints ``` --- ## đŸŽ¯ Model Training Each task has dedicated training code designed for ease-of-use and collaboration. Training scripts include reasonable defaults, with paper-replicating configurations in accompanying script folders. ### Image Classification Example ```bash # Standard CTM training python -m tasks.image_classification.train # Energy halting training python -m tasks.image_classification.train_energy \ --energy_head_enabled \ --loss_type energy_contrastive ``` ### VSCode Debug Configuration ```json { "name": "Debug: train image classifier", "type": "debugpy", "request": "launch", "module": "tasks.image_classification.train", "console": "integratedTerminal", "justMyCode": false } ``` --- ## 🔍 Analysis & Visualization Analysis and plotting code to replicate paper figures is provided in `tasks/.../analysis/*`. **Note:** `ffmpeg` is required for generating videos: ```bash conda install -c conda-forge ffmpeg # or with pixi (already included) pixi install ``` --- ## đŸ“Ļ Checkpoints and Data Download pre-trained checkpoints and datasets: - **Checkpoints**: [Google Drive](https://drive.google.com/drive/folders/1vSg8T7FqP-guMDk1LU7_jZaQtXFP9sZg) - **Maze Data**: [Google Drive](https://drive.google.com/file/d/1cBgqhaUUtsrll8-o2VY42hPpyBcfFv86/view?usp=drivesdk) Place checkpoints in the `checkpoints/` folder following the structure `checkpoints/{task}/...` --- ## 🤗 Hugging Face Integration This repository includes full support for training on Hugging Face infrastructure: - **Accelerate**: Multi-GPU and mixed precision training - **Hub Integration**: Automatic checkpoint uploading - **Spaces Deployment**: Run training jobs on GPU Spaces See [GUIDE_HF.md](GUIDE_HF.md) for detailed instructions. --- ## 📖 Interactive Resources - **[Interactive Website](https://pub.sakana.ai/ctm)**: Maze-solving demo, videos, and visualizations - **[Paper](https://arxiv.org/abs/2505.05522)**: Technical details and experiments - **[Blog](https://sakana.ai/ctm/)**: High-level overview and insights - **[Tutorial Notebook](examples/01_mnist.ipynb)**: Hands-on introduction --- ## 🙏 Citation ### The Continuous Thought Machine : Energy-Based Halting Extension This repository contains experimental extensions for Energy-Based Halting developed by **Uday Phalak**. ```bibtex @misc{ctmenergy2025, title={Energy-Based Halting for Continuous Thought Machines}, author={Phalak, Uday}, year={2025}, note={Experimental Extension of Continuous Thought Machines} } ``` Based on The Continuous Thought Machine ```bibtex @article{ctm2025, title={The Continuous Thought Machine}, author={Darlow, Luke and Regan, Ciaran and Risi, Sebastian and Seely, Jeffrey and Jones, Llion}, journal={arXiv preprint arXiv:2505.05522}, year={2025} } ``` --- ## 📝 License This project is released under the MIT License. See LICENSE file for details.