Continuous Thought Machines

Figure˜8 show the synapse model which is the recurrent structure that shares information across neurons in the CTM. It is implemented by choosing a depth of $k$ (always even), where each subsequent layer width is chosen to linearly reduce the dimensionality until a width of 16 is reached, and then increase thereafter, using skip connections to retain information. The synapse model also takes $\mathbf{o}^{t}$ (the output of attention) as input.

The CTM operates using recurrence on a latent representation, $\mathbf{z}$, that is $D$-dimensional. $\mathbf{z}$ unfolds over time and the synchronization between some chosen neurons form the new kind of representation that the CTM enables.

There are $\frac{D\times(D+1)}{2}$ unique pairs of neurons, making for a substantially larger set of neuron synchronization pairs than there are neurons themselves. This motivates the need for a selection process. Over the development of the CTM we came up with three approaches to selecting neurons:

1. 1.

   Dense pairing: in this setup we select ${J}$ neurons and compute synchronization for every possible $(i,j)$ pair of the ${J}$ neurons. For $\textbf{S}^{t}_{\text{out}}$ we choose $J_{\text{out}}$ neurons and for $\textbf{S}^{t}_{\text{action}}$ we choose non-overlapping $J_{\text{action}}$ neurons. Selecting $J_{\text{out}}$ neurons for dense pairing results in an output synchronization representation of $D_{\text{out}}=\frac{J_{\text{out}}\times(J_{\text{out}}+1)}{2}$, and similarly for the action representation.

   This approach essentially creates a strong bottleneck where all gradients must flow through the selected neurons, which can be advantageous for some tasks.

2. 2.

   Semi-dense pairing: in this setup we open up the aforementioned bottleneck twofold by selecting two different subsets, $J_{1}$ and $J_{2}$, such that the left neurons of the synchronization dot product, $i$, are taken from $J_{1}$ and the right neurons, $j$, are taken from $J_{2}$. The same dense computation as before is then applied.

   The bottleneck width in this case is $2\times$ as wide as before. Output and action selections are, once more, not overlapping.

3. 3.

   Random pairing: in this setup we randomly select $D_{\text{out}}$ or $D_{\text{action}}$ pairs of neurons and compute the synchronization between each pair as opposed to doing so densely between all selected neurons. We also intentionally compute the $(i,i)$ dot products between $n_{\text{self}}$ neurons in each case in order to ensure that the CTM could recover a snapshot representation if it wanted to.

   This opens up the bottleneck much more than before. We allow overlapping selections in this case.

We used the maze-dataset repository to generate mazes for this work. We generated mazes of size $19\times 19$, $39\times 39$, and $99\times 99$. In each case we generated 50000 mazes and split them into train sets of size 45000 and test sets of size 5000. We used the $39\times 39$ for training in this technical report and tested generalization on the $99\times 99$. We provide all three maze datasets³³3We found the $19\times 19$ beneficial for debugging, hence we provide it too. in the CTM code repository, made available upon publication.

We used the following hyperparameters:

• •

  $39\times 39$ mazes and a ResNet-34 backbone, where the keys and values were taken as features after the second hyper-block, resulting in a down-sample to $10\times 10$

• •

  $D=2048$ (the width of $\mathbf{z}^{t}$ and $\mathbf{a}^{t}$)

• •

  $k=16$ (synapse depth, $8$ layers down and $8$ layers up)

• •

  $d_{\text{input}}=512$ (the width of attention output, $\mathbf{o}^{t}$)

• •

  ${n}_{\text{heads}}=16$

• •

  Dense pairing for neuron selection (see Section˜C.2)

• •

  $J_{\text{out}}=32$ (width of $\textbf{S}^{t}_{\text{out}}$ synchronization representation)

• •

  $J_{\text{action}}=32$ (width of $\textbf{S}^{t}_{\text{action}}$ synchronization representation)

• •

  $T=75$ (internal ticks)

• •

  $M=25$ (FIFO rolling memory input to NLMs)

• •

  $d_{\text{hidden}}=32$ (width of MLPs inside NLMs)

• •

  $p_{\text{dropout}}=0.1$ (dropout probability for synapse model)

• •

  No positional embedding

We used the following settings for optimization:

• •

  Trained using a batch size of 64 on 1 H100 Nvidia GPU

• •

  1000000 iterations for training using AdamW [47]

• •

  A learning rate of 1e-4 with a linear warmup of 10000 iterations and decaying to zero using a cosine annealing learning rate scheduler

• •

  No weight decay

This resulted in a model with 31,998,330 parameters.

We adapted the loss function for solving mazes to include a curriculum element. Before computing $t_{1}$ and $t_{2}$ for the loss in Equation˜11 we first altered the loss at each internal tick to only account for those steps in the maze that were correctly predicted with plus 5 additional steps along the path. This effectively lets the model (CTM or LSTM baselines) slowly learn to solve the maze from start to finish. Listing LABEL:listing:maze-loss shows how this is implemented when computing the loss and is used for both CTM and LSTM training.

We tested a number of LSTM baselines for solving this task, but struggled with stability during training (see Figure˜4), particularly beyond a single LSTM layer or with a higher number of internal ticks. Hence we tested three LSTM configurations of depths 1,2, and 3. For each model we tested with 75 internal ticks to match the CTM, and 50 internal ticks for stability. We also tested a feed-forward model, projecting the feature space (before average pooling) into a hidden layer of the same width as the CTM, yielding a slightly higher parameter count. We kept all hyperparameters constant, yielding the following setups:

• •

  LSTM, 1 layer, $T=50$ and $T=75$: 42,298,688

• •

  LSTM, 2 layers, $T=50$ and $T=75$: 75,869,504 parameters

• •

  LSTM, 3 layers, $T=50$ and $T=75$: 109,440,320 parameters

• •

  Feed-forward, with a hidden layer width of 2048 (and GLU activation thereon): 54,797,632 parameters

Figure˜9 gives the loss curves for the maze solving models in Section˜4.1, showing how the CTM is more stable and performant when trained on this task.

Internal models of the world and cognitive maps represent crucial aspects of intelligent systems [43, 48, 49]. In this case, we consider a world model to be an internal representation of the external environment, encapsulating an agent’s knowledge about the world’s structure, its dynamics, and its actionable place therein. A good world model should enable an agent to reason about the world, plan, and predict the consequence of its actions. Cognitive maps [48] specifically focus on spatial relationships and navigation. The ability to construct and utilize these internal representations is a strong indicator, and arguably a prerequisite, for sophisticated intelligence. The notion of ‘episodic future thinking’ [44] is even considered a hallmark feature of human intelligence. An agent devoid of a world model would be limited to reactive behaviors. Similarly, lacking a cognitive map would severely restrict an agent’s ability to navigate and interact effectively within complex spatial environments. Therefore, the presence and sophistication of world models and cognitive maps can serve as a benchmark for evaluating intelligence.

To this end, we designed the maze task such that it would require a good internal world model to solve. This was achieved by (1) requiring the model to output a route directly, as opposed to solving the maze with a local algorithm [37], and (2) forgoing any positional embedding in the image representation, meaning that the model must build its own spatial cognitive map in order to solve the task [48]. Indeed, we saw that the NLMs and synchronization components of the CTM enables it to solve our 2D maze task, far surpassing the best baselines we trained. Our results suggest that the CTM is more capable of building and utilizing an internal model of its environment.

We used a constrained version of the classic ResNet architecture [41] for this task, which we adapted from https://github.com/huyvnphan/PyTorch_CIFAR10. It differs from the standard implementation for ImageNet in that the first convolution is constrained to use a kernel size of $3\times 3$ as opposed to $7\times 7$. We used a ResNet-152 structure and took the output prior to the final average pooling and projection to class logits. We used input images of size $224\times 224$ which yielded $14\times 14$ features for the keys and values used in cross attention.

We used the following hyperparameters:

• •

  $D=4096$ (the width of $\mathbf{z}^{t}$ and $\mathbf{a}^{t}$)

• •

  $k=16$ (synapse depth, $8$ layers down and $8$ layers up)

• •

  $d_{\text{input}}=1024$ (the width of attention output, $\mathbf{o}^{t}$)

• •

  ${n}_{\text{heads}}=16$

• •

  Random pairing for neuron selection (see Section˜C.2)

• •

  $D_{\text{out}}=8196$ (width of $\textbf{S}^{t}_{\text{out}}$ synchronization representation)

• •

  $D_{\text{action}}=2048$ (width of $\textbf{S}^{t}_{\text{action}}$ synchronization representation)

• •

  $n_{\text{self}}=32$ (for recovering a snapshot representation)

• •

  $T=50$ (internal ticks)

• •

  $M=25$ (FIFO rolling memory input to NLMs)

• •

  $d_{\text{hidden}}=64$ (width of MLPs inside NLMs)

• •

  $p_{\text{dropout}}=0.2$ (dropout probability for synapse model)

• •

  No positional embedding

We used the following settings for optimization:

• •

  Trained using a batch size of 64 across 8 H100 Nvidia GPUs

• •

  500000 iterations for training, using a custom sampling such that each minibatch was sampled with possible replacement

• •

  AdamW [47]

• •

  A learning rate of 5e-4 with a linear warmup of 10000 iterations and decaying to zero using a cosine annealing learning rate scheduler

• •

  Gradient norm clipping set to a norm of 20

• •

  No weight decay

Listing LABEL:listing:image-loss shows the python code for the image classification loss function used to train the CTM on ImageNet-1K. This accompanies the loss defined in Section˜3.5.

We showcase a series of additional demonstrations of CTMs classifying images from ImageNet in Figures˜12, 13, 14 and 15. We encourage the reader to view similar demonstrations as videos (we will provide links upon publication; see supplementary video ‘imagenet.mp4’)

To visualize the dynamics of the neuron activations we used UMAP [50] to project the neurons to a 2D feature space. We used the history of post-activations for 200 different ImageNet images as the high dimensional inputs to UMAP. Specifically, this input was $200\times T=200\times 5=1000$ dimensional. The resulting 2D projection assigns each neuron to a position in space determined by its activation ‘profile’ – its response pattern both over time and over multiple stimuli. Visualizing this mapping over internal ticks reveals low frequency structures propagating across the feature space. We show snapshots of this visualization over the internal ticks in  Figure˜11, noting that these visualizations are best viewed in video form (see supplementary video ‘umap.mp4’).

Importantly, the CTM generates this structure in an emergent fashion, without any explicit driving signal. Analogous phenomena occur in networks of Kuramoto oscillators [51]; in our case, waves propagate across a learned feature map in an all-to-all network. Concurrent work also explores explicitly encoding traveling waves for long-range communication [52]. We do not assign functional meaning to these observed waves but highlight their distinct presence during the CTM’s thought process.

The parity of a binary sequence is given by the sign of the product of its elements. When processing a sequence element by element, an RNN could conceivably compute parity by maintaining an internal state, flipping an internal ‘switch’ whenever a negative number is encountered. However, if the entire sequence is provided simultaneously, the task increases in difficulty due to the increasing number of distinct patterns in the input. Previous work [18] has addressed this challenge using recurrent models, which can learn sequential algorithms for statically presented data. Computing parity, as posed in this manner, is well-suited for testing the capabilities of the CTM.

We apply the CTM to the task of computing the parity of a 64-length sequence containing the values 1 and -1 at random positions. Unlike [18], we set up the task such that the model computes the cumulative parity at every index of the sequence, not just the final parity. An example is shown in Figure˜6(a). The values -1 and 1 are embedded as learnable vectors combined with positional embeddings, using attention to ingest input data. We train the CTM with the loss function described in Section˜3.5. As a baseline we also trained an LSTM, but set $t_{2}$ to be the final iteration since this gave the best results and stability for LSTM training.

In this section, we describe some additional results for the parity experiments.

The input data for the parity task is a vector of length 64, where at each position is either a $-1$ or $1$. The target for each sample is a vector of the same size, where at each position is the parity of the sequence up to that position, which we refer to as the cumulative parity. This data is generated on the fly, each time a new batch is fetched.

The following architecture is used for the experiments in the parity task. Inputs to the model are of shape ($B$, 64), where the size of the minibatch and sequence length are $B$ and 64, respectively. Each of the values in the 64-length sequence are either -1 or 1, at random. First, the values of -1 and 1 are converted into embeddings in $d_{\text{embed}}=d_{\text{input}}$ and positional embeddings are added. The resulting embeddings are passed through a linear layer with layer normalization to form (identical) attention keys and values. As described in section 3, the synchronization between $J_{\text{action}}$ neurons is computed and from this representation an attention query is formed. This query is then used to compute the attention values, which are concatenated to the activated state to be processed by the synapses and the neuron-level models. For the synapses we use a shallow feedforward network. This process repeats for $T$ internal ticks, where at each internal tick $t$, the synchronization can be computed between $J_{\text{out}}$ neurons and projected to the logit space.

In the parity task, we experimented with how the model performs with a varying number of internal ticks and memory length. As a baseline, we use single-layer LSTMs which are both parameter matched and use the same number of internal ticks as the CTM.

All CTM models share a common set of architectural hyperparameters, listed below. The Table˜3 shows the subset of hyperparameters that vary across experimental configurations.

• •

  $d_{\text{model}}=1024$

• •

  $d_{\text{input}}=512$

• •

  $d_{\text{hidden}}=4$

• •

  $k=1$

• •

  $p_{\text{dropout}}=0$

• •

  $n_{\text{heads}}=8$

• •

  $J_{\text{action}}=32$

• •

  $J_{\text{out}}=32$

• •

  Semi-dense pairing was used for selecting neurons for synchronization

• •

  Absolute positional encoding was added to the input features

The CTM was trained using the certainty-based loss function described in section 3.5, whereas the LSTM baselines utilized the cross-entropy loss computed at the final internal tick. This choice was made due to initial difficulties in training the LSTM effectively with the certainty-based loss. In Figure˜18, we compare training accuracy curves for the LSTM baselines with 10 and 25 iterations, trained with either the final or certainty-based loss. Generally, both loss functions lead to unstable training for LSTMs with multiple internal ticks.

We used the following settings for optimization:

• •

  Trained using a batch size of 64 on 1 H100 Nvidia GPU.

• •

  200000 iterations for training using AdamW [47].

• •

  A learning rate of 1e-4 with a linear warmup of 500 iterations and decaying to zero using a cosine annealing learning rate scheduler.

In this section we test the CTM using CIFAR-10, comparing it to human performance, a feed-forward (FF) baseline, and an LSTM baseline. For the model-based baselines, we used a constrained featurization backbone in order to emphasize the differences owing to the model structure post-featurization (i.e., CTM versus LSTM versus FF). We also used 50 internal ticks to give the CTM and LSTM ‘time to think’. The human and model baselines were set up as follows:

• •

  Human baseline. We used two datasets of human labels for CIFAR-10; we call these CIFAR-10D [53] owing to its calibration of difficulty levels, and CIFAR-10H [54] originally used to quantify human uncertainty.⁴⁴4CIFAR-10D can be found at https://sites.google.com/site/hophuoctien/projects/virec/cifar10-classification; CIFAR-10H can be found at https://github.com/jcpeterson/cifar-10h We used CIFAR-10D to determine easy versus difficult samples, and CIFAR-10H as a direct human baseline.

• •

  FF baseline. A feed-forward only baseline (denoted FF). An MLP was applied to ResNet features after average pooling, where the width of the hidden layer was set to match the parameter count of the CTM for this experiment.

• •

  LSTM baseline. An LSTM set up to unroll with an internal thought dimension, with a hidden width set to match the parameter count of the CTM. The LSTM could attend to the image at each step and used the same loss as the CTM for valid comparison.

Figure˜21 shows the training curves of the CTM, FF, and LSTM models, and calibration plots for each, including an estimation of human calibration using CIFAR-10H. The FF baseline reaches a high training accuracy early on, but also demonstrates a poor generalization gap. The LSTM is less stable during training (we had to set the learning rate to $0.0001$ for all experiments because of this) and yields a marginally improved test accuracy. The CTM is more stable and performant.

For the human calibration we used the probabilities provided in CIFAR-10H, which were computed using guesses from multiple humans. We computed calibration here as we did for ImageNet-1K (see Figure˜5b in the main paper): we compute the predictive probability as the average probability for the chosen class over all internal ticks. None of the models are perfectly calibrated, but the CTM demonstrates the best calibration, even when compared to humans. Strikingly, the CTM has even better calibration than humans, while the LSTM follows the human under-confidence.

Figure˜22(a) compares models and CIFAR-10H against the difficulty determined using the CIFAR-10D dataset. Each model and humans have similar trends in this case, although the CTM follows most closely to CIFAR-10H. Figures˜22(b) and 22(c) compare the uncertainties of the CTM and LSTM to the uncertainties of humans (using reaction times from CIFAR-10H as a proxy for uncertainty). We compute the CTM and LSTM uncertainties using the normalized entropies (see Section˜3.5 in the main paper) averaged over internal ticks as this approximates the total uncertainty each model has regarding the observed data. Both the CTM and LSTM exhibit trends similar to human reaction times.

Figure˜23 shows the neural activities for the CTM and the LSTM baseline. The CTM yields rich, diverse, and complex dynamics with multiple interesting features, including periodic behavior (there is no periodic driving function). The distinct difference between the CTM and LSTM neural activities is evidence that the two novel elements of the CTM (NLMs and synchronization as a representation) enable neural dynamics as a fundamental computational tool.

In this section we explore two aspects of the CTM: (1) width (i.e., number of neurons), and (2) number of internal ticks. We used CIFAR-100 in the experiments discussed below as it is a more challenging dataset than CIFAR-10, while remaining relatively low-demand regarding compute.

Core to the design of the CTM are the NLMs and the use of synchronization as a representation. To understand how both of these components contribute to the performance of the CTM, we carried out ablations using a variant of the maze task of Section˜4, with $15\times 15$ sized mazes. Our ablations include four model configurations:

1. 1.

   The standard CTM as detailed in Section˜3.

2. 2.

   A CTM without NLMs, where additional synapse model layers are introduced to maintain an equivalent number of parameters.

3. 3.

   A CTM without synchronization, in which attention queries and output projections are formed directly from the post-activations $z^{t}$.

4. 4.

   An LSTM with synchronization, where the synchronization is computed using the time series of each dimension of the LSTM’s hidden state.

Each configuration was parameter matched to approximately 9M parameters and trained for $100000$ iterations. A full list of hyperparameters is shown in Table˜4.

The training curves for each of the four variants are shown in Figure˜26, with the final results summarized in Table˜5. We find that the two augmented variants of the CTM and the LSTM with synchronization were unable to achieve a solve rate above $50\%$. The standard CTM, on the other hand, achieves a solve rate of $66\%$ and a per-step accuracy of $95\%$. These findings suggest that the combination of neuron-level models and synchronization as a representation is key to the success of the CTM.

In this section, we apply the CTM to the task of sorting 30 numbers drawn from the normal distribution. Sorting real numbers was a task explored by [18] when designing RNNs for adaptive compute, and it provides a test bed for understanding the role of compute for an adaptive-compute system, such as the CTM. In this case the CTM does not use attention, but rather ingests the randomly shuffled input data (30 real numbers) directly. This is implemented by replacing the attention mechanism with a straightforward concatenation, replacing in Figure˜3.

To assess the CTM’s capabilities for memory, retrieval, and arithmetic computation, we devise a Question and Answering (Q&A) MNIST task, reminiscent of [55] or  [56]. In this task, the model sequentially observes a series of MNIST digits [57], followed by an interwoven series of index and operator embeddings that determine which of the observed digits to select and which modular operation to perform over them. This allows us to probe whether the CTM can simultaneously recognize hand-drawn digits, recall previous observations, and perform logical computation on them without any prior knowledge of the digits depicted in the images or the relationships between them. Furthermore, by applying more operations at inference time than observed during training time, we can test the generalizability of the CTM.

Specifically, the model first observes $N_{d}$ MNIST digits sequentially for $t_{d}$ internal ticks each. Next, the model receives an interwoven sequence of $N_{\text{idx}}$ index embeddings (indicating which digit to select) and $N_{\text{op}}$ operator embeddings (specifying either modular addition or subtraction, where each intermediate result is taken modulo 10 to keep answers within the range 0–9), each presented for $t_{\text{idx}}$ and $t_{\text{op}}$ internal ticks, respectively. Finally, the model observes a zero tensor for $t_{\text{ans}}$ internal ticks, signaling the model to produce its answer. The target, between 0 and 9, results from the composition of all specified modular arithmetic operations. An example is shown in Figure˜29.

We trained CTMs and parameter-matched LSTMs with two different configurations, varying how many internal ticks were used to process each input. Digits and embeddings were observed for either 1 or 10 internal ticks, with corresponding answering times of 1 or 10 internal ticks. The number of digits and the number of operations were sampled uniformly between $1$ and $4$. Memory lengths for the $1$ and $10$ internal ticks per input CTMs were set to $3$ and $30$ steps, respectively. We highlight that with these observation and memory length configurations that the digit observations will always lie outside of the memory length-sized window during the answering stage. In this way, the CTM must organize its activations such that it can recall the digits at later time steps. The CTM is trained with the loss defined in Section˜3.5 (main paper), computed only over the final $t_{\text{ans}}$ steps. Once more, we set $t_{2}$ to be the final iteration for the LSTM for stable training.

We have previously shown that the CTM can process sequentially on non-sequential tasks via its decoupled internal recurrence. Here, we extend the CTM to sequential decision-making tasks involving interactions with external environments. Specifically, we train CTMs using reinforcement learning (RL), where the model learns action-selection policies based on environmental observations and trial-and-error interactions. In this setting, the CTM processes one or more internal ticks before producing an action that transitions the environment to the next state. To achieve this, we continuously maintain the neuron dynamics across these internal ticks over successive environment steps, allowing previous environmental observations to influence current internal states via the NLMs. A central goal of this section is to provide evidence that the CTM can be set up to learn in a continuous environment.