Continuous thought: the horse and its blinders
Recently, I noticed that every time my agent starts thinking, it repeats what I asked before producing any meaningful original thought. I got really curious as to why that happens. Humans tend to do this, repeating what they heard to align their thoughts. Why do humans do that though?
Humans have a very small working memory. We can only keep very few things in our head at a time. There are very contested numbers on this. Most popular estimates put it around $7 \pm 2$ things at a time. If we have to make space for something, something else needs to go. Also the things that do stay in our memory tend to behave sort of like an LRU cache but with some weight to repetitions. So we repeat to keep the thing we want to think about in our working memory.
It would be anthropomorphizing to say that LLMs do that for the same reason. Its entire working memory is its context and nothing needs to go. The amount of information it can keep is extraordinarily high compared to humans. Yet still it repeats, why? There could be two reasons. One, the training dataset that is used to do Chain of Thought (CoT) training is predominantly human thought traces or modeled on them. So it could be a training artifact, it's trained to think that way, so it does, case closed. But there could be a second reason, it could be circumventing an architectural constraint just like humans do. What is that?
Problems with causal mask
Long time ago, when prompt engineering was in vogue, a 2023 paper titled Re-Reading Improves Reasoning in Large Language Models was published. They tried repeating the question to see if repetition has any effect on the reasoning quality of the models. And they indeed found that the models seem to do better if some of the complex prompts were repeated.
This should be extremely surprising result. Why would repetition in prompt help at all? Just as with humans, this seems to circumvent a crucial architectural constraint we have in our modern foundation models called causal mask (the paper's claim). The causal mask is how a token can see what was there before it but not after it during the self attention side. By repeating the sentence again, now we have let the model see per token what the future holds for them.
Let's take the following sentence cat sat on a mat.
Here due to the causal mask the cat doesn't know that it's sitting on a mat. Though the mat knows there is a cat sitting on it. The graph of
relations between the tokens seem kind of incomplete. But the moment we repeat the sentence again, the second cat can attend to previous mat and
knows not only it is doing the sitting and also that the sitting is on a mat. But repeating made the entire graph of relations readily available.
In my agent's case the repetition is in its thinking not in my prompt, but that makes no difference. It might have learnt to do so on a short
term due to it being represented in training dataset and eventually was getting better results due this hack and started using it more consistently?
What the agent has done is fake bidirectional attention by repeating the question back. Looks like some of form of reasoning is effective if the attention is bidirectional. The erstwhile encoder-decoder models or encoder only models like T5 and BERT are not restricted by the causal mask and have been shown to understand prose better than the decoder only models. Even the system one decision models that are getting popular these days predominantly use bidirectional attention. Why aren't we doing the bidirectional attention in generative language models if it clearly seems to help?
We have made tremendous progress on decoder only models, in both training and inference. The generation is simpler and more efficient because of the KV cache stays valid. The KV cache for a token once created is not going to get affected by newer generated tokens and can be effectively used. But if we have bidirectional attention, the KV cache is no longer stable. Although the original encoder-decoder uses bidirectional attention only on the prompt, still uses causal mask on the newly generated token, so cache is stable for this turn. But the cache is invalid for next conversation turn, and the cache has to be regenerated. We might have caught the tiger by its tail. We can't let go of the decoder models as we might have to give up at least some of the progress we made here. Looks like we have to live with more tokens than necessary.
Repetition is one kind of overhead, but the other is the sampling of tokens in CoT reasoning steps. That raises the question of, why the thought has to be concrete tokens at all? To put some substance on the bandwidth lost by sampling tokens, modern foundation models have vocabularies of about 200K tokens, so a token carries at most $\log_2(200000) \approx 18$ bits of information. But the embedding vector has close to a 7k+ (Kimi K2.6, DeepSeek-V3) floating point values. Even at 1 bit quantization that's 400 times the information capacity there. Though that's raw capacity, the usable capacity is smaller, but still far more than what we have with concrete tokens. Can we have better utilization of that extra capacity by not thinking in concrete tokens?
The Tao that can be told
Bandwidth loss is not the only problem with sampling thinking tokens. By sampling a token we are also committing to a particular chain of thought. The abstract thought before sampling can be seen as holding several possible continuations at once. But by forcing the model to produce a token, we are committing to one path very greedily before knowing how that thought would end up. This raises the question again, why can't an abstract thought stay abstract? By making it concrete we are losing so much.
Lao Tzu put it in the Tao Te Ching: The Tao that can be told is not the eternal Tao. The moment you put it into words, you have already left something out.
Although, sampling of thinking tokens has some real uses too, in terms of safety and interpretability of the models. It is sort of like an aircraft's black box. We are able at least investigate what happened in OpenAI/Hugging Face incident because we have legible reasoning tokens. But for a moment let's assume that we don't have any of those issues.
To understand how to keep the thought abstract, we need to have a mental model of how transformers work. The simplest way to picture a transformer is that of residual streams. In simple terms, the whole thing roughly looks like this
def transformer(xs):
for block in blocks:
# causal: position k sees only xs[:k+1]
xs = xs + block.attention(xs)
xs = xs + block.fully_connected(xs)
return xs
xs = [embed(t) for t in prompt]
while True:
last_state = transformer(xs)[-1] # take only last token stream
token = sample(softmax(unembed(last_state)))
if token == end:
break
# This is the part we are interested in
# Our proposition is instead of this,
xs.append(embed(token))
# How about xs.append(last_state)
I generally don't like overwriting the same variable over and over again. But here I meant to write it this way to get across the residual stream idea. This is however a very simple and naive model as we need to take care of positional encoding at every layer, multiple heads, MoE routing, layer normalization, the looped models doing variable number of transformer blocks some of them repeating and so on. But this is simple enough to understand and characterize a lot of the behaviors of the transformer models.
Look at the append line, it's pretty clear we are unembedding and embedding back again. In case of output tokens, the unembedding clearly makes sense, but
when we are reasoning it does lose some information. We lost some of the richness of it by sampling one token and re-embedding again. Why not just use the last hidden state
as the input to the next reasoning step? The sampling seems to have no value, but it turns out it's the foundation on which everything else works.
Why the horse needs blinders
There are two foundational reasons why this isn't standard yet. The first one is the semantic collapse problem and the other is the training parallelism problem. Let's first look at the semantic collapse problem. In order to understand that, we take an even simpler language model, one that just understands a single word. There used to be this great model called word2vec which was trained to produce embeddings for words. I was mind blown when it first came out. The demos were incredible. One of the examples was
$\overrightarrow{\text{king}} - \overrightarrow{\text{man}} + \overrightarrow{\text{woman}} \approx \overrightarrow{\text{queen}}$
The underlying premise was extremely simple, the word's meaning lies in its context (You shall know a word by the company it keeps - J.R.Firth). We take a large enough corpus of text, and train a model to say if one word is in the context of another (There were two variants, CBOW and skip-gram. This is CBOW). Our model is trained to find embeddings for words such that the dot product of the embeddings of the word and the average embedding of every instance of its context words is maximized. We apply a sigmoid to the dot product, read as the probability of one word appearing in the company of selected others. Something like this,
def classify(word, context):
x = embed_pivot(word) # Just look up
ys = [embed_context(y) for y in context] # Again but usually different map
return sigmoid(x.T @ average(ys))
There is a trivial solution to this problem, map every word to the same vector and the dot product would be maximal. Since the dot product is unbounded, it keeps increasing in magnitude too as training progresses. Loss is basically zero and the embedding is meaningless. This is a very simple example of a semantic collapse. There is no incentive for each word to be taken to a different embedding by just showing it positive samples. I suppose non-examples are just as important for understanding as examples are, even for neural networks. The word2vec solves this by doing negative sampling, ie., for each pivot word, they sample words at random from the vocabulary and call it a negative sample. This course corrects it, the identical vector solution is no longer maximal. The vectors can't keep growing forever or it would score badly on non-examples. Now the embedding has to be meaningful. Moral of the story is that neural networks need blinders like horses do. You need to hide the distracting shortcuts from it.
But LLMs aren't trained with negative sampling? Where are the blinders?
Here the blinder takes the form of softmax. The softmax cannot give all the words maximal probability. Probability is traded. If one token becomes more probable, another has to become less probable. Softmax being bounded makes each token become negative of others and vice versa. Cross entropy loss on softmax is negative sampling in disguise, with the whole vocabulary as negative samples. The continuous thought is removing exactly that blinder. Unless we substitute something else the horse is going to go astray.
Taking off the blinders bit by bit
COCONUT (Chain of Continuous Thought, Hao et al., FAIR at Meta) tried to loosen the blinders a little bit at a time. They split the CoT training dataset into $n$ stages. Stage 0 is regular CoT training with sampling. At stage $k$, they replaced each of the first $k$ reasoning steps in each thinking trace with $c$ (hyperparameter, fixed per training run) continuous thoughts. A continuous thought is where we take the last hidden state as is, as the next input embedding, without sampling and re-embedding. Rest of the reasoning steps still stayed as concrete language tokens. The loss computations are also disabled for the first $k$ steps, however the gradient flow from the later steps still flow back through them. So the blinders are loosened for first $k$ steps for every reasoning chain. But they do get indirect supervision from the later language tokens.
They were able to see a breadth-first like search going on in the continuous thought. The first thought has high probability spread across multiple tokens, which they interpreted by unembedding the continuous thought and calculating the softmax probability for like say top few tokens. It sort of had frontier tokens for each possible CoT. The subsequent thoughts started narrowing down on what thought seemed more viable. So individual continuous thought looked like a superposition of frontier thoughts, and as it progressed it settled on most viable ones. The model seemed to be hedging on certain thoughts and stayed non-committal, sort of what we expected to happen. On ProsQA dataset, the team claims that this reduces hallucinations, because it doesn't have to live with bad decisions taken earlier. In essence, on ProsQA at least, the continuous thought made the models better at planning too. However it was significantly worse on math, benchmarked with GSM8k dataset.
When they raised the continuous thoughts per step to $c = 3$, the loss spiked at stage transitions. My guess is that the only gradient flow is through the tokens further down, and it sort of starts behaving like recurrent neural networks at that point. The gradient flow is not going to be consistent (vanishing or exploding gradients). So the moral of the story is the blinders can't be loosened too fast, otherwise the horse bolts. But we do have concrete tokens after $k$ steps. Can't we do RL by assigning rewards for these continuous thoughts once we have concrete thoughts? Wouldn't that fix the gradient issue?
Why we can't just train it with RL
Why not reinforcement learning here? Can we stop being thought police of LLMs? Can we just instruct the models if their thoughts were good or bad purely based on outcomes? Let the horse graze the pasture without blinders like cows do and come back to the shed at the end of the day? This ties back to the training parallelism problem that we promised earlier.
The big advantage of transformers is that, though the network is autoregressive on its tokens, the training need not be. We can do teacher forcing and train an entire block of text at a time. It's a much cheaper way to train such a large network. This is why we don't need a KV cache during training. This is also why we are able to scale the networks to trillions of tokens and hence billions of parameters (scaling law). When we have continuous thought we are stuck with sequential training one step at a time.
Most reasoning models are already trained with RL on verifiable rewards and the update is still teacher-forced to scale the training. I mean, we just feed what token/embedding the rollout sampled , on each position in parallel, without waiting for that position to be fully realized. However in case of continuous thoughts, there is nothing to teacher force. So RL or no RL the training has to be sequential at least for the continuous thought. The inability to teacher force is the primary problem with continuous thought. The other problem is evaluating policy gradient, the sampled token would give us log probability of the actions taken at the end of the reasoning steps (through the probability of the tokens sampled along the way), even if we are rewarding only the eventual actions. Every sampling gives you both the probability and a token, continuous thought doesn't have either. We have no way to ascertain the probability distribution of the eventual action token or the thought's contribution to it. COCONUT trained GPT-2 with small datasets with clean step boundaries and scaling continuous thoughts to large datasets is still an open problem. Thus sampling and indirectly softmax kept the blinders on. It allowed us to guide which thoughts are valuable, by clean supervised gradients or RL rewards through policy gradients. It gave hand rails in training in terms of token probability.
What to ask next time
So whenever someone says that their model reasons in latent space, the first question you need to ask is what the blinders are. JEPA, Yann LeCun's family of models, is the most committed attempt at an out-and-out latent space model. The hardest problem is exactly what we discussed here, semantic collapse. LeCun is trying to build world models based on images and videos, but instead of trying to reproduce pixel by pixel the occluded part of an image or videos, they try to produce a latent embedding of what should be there. Like if there is a boy throwing a ball on the visible part, it predicts the latent embedding of the occluded part, which includes a ball. We might not be able to get the color of it right, but we might be able to guess the size of the ball. There are multiple versions of JEPA. The latest, LeJEPA, constrains the embedding to spread out like a Gaussian distribution, doing what the softmax did in our language models. Invariably there are always blinders. We don't need the words, but we haven't yet figured out how to ride the horse without blinders.