Do Language Models Plan Ahead for Future Tokens?

Authors: wilson  wu
john x morris
lionel  levine
Venue: colmweb_org_COLM_2024_Conference
Published: 2024-03-21 23:38:45

Abstract: Do transformers ``think ahead'' during inference at a given position? It is known transformers prepare information in the hidden states of the forward pass at time step $t$ that is then used in future forward passes $t+\tau$. We posit two explanations for this phenomenon: pre-caching, in which off-diagonal gradient terms present during training result in the model computing features at $t$ irrelevant to the present inference task but useful for the future, and breadcrumbs, in which features most relevant to time step $t$ are already the same as those that would most benefit inference at time $t+\tau$. We test these hypotheses by training language models without propagating gradients to past timesteps, a scheme we formalize as myopic training. In a constructed synthetic data setting, we find clear evidence for pre-caching. In the autoregressive language modeling setting, our experiments are more suggestive of the breadcrumbs hypothesis, though pre-caching increases with model scale.

Main paper: 
 introduction: 
 Humans are known to think ahead while speaking; decades of linguistics research (Huettig, 2015; Miller, 1951) have shown evidence that human language users internally predict upcoming language input, words and sometimes sentences ahead (Barthel et al., 2016) .Unlike humans, contemporary language models allocate a fixed amount of information processing for each token when "speaking" (Vaswani et al., 2017) . Do language models, like humans, think ahead? Recent work (Pal et al., 2023; Hernandez et al., 2024) has shown that tokens beyond the immediate next token can be predicted by probing the hidden state of the language model. Model outputs at future tokens can be predicted to some extent using linear probes on model hidden states, and interventions on hidden states can predictably alter future outputs. 1These findings indicate that model activations at a given timestep are at least somewhat predictive of future outputs. However, it remains unclear why this might be: is this just a happenstance property of the data, or because the model is deliberately preparing information for future timesteps, at the expense of degrading performance on the current position?We observe that gradients during training optimize weights for both the loss at the current token position as well as for tokens later in the sequence. We question to what extent current transformer weights dedicate resources to the current token vs. allocating it for future tokens.Figure 1 : At which position is the computation required to correctly answer this math problem taking place? Cognitive science tells us that humans think ahead while speaking; we investigate the extent to which language models do the same.We consider two possibilities: the pre-caching hypothesis, in which the transformer learns to compute features at time step t that are irrelevant to the inference task at that current time step but may be useful for future time steps t + τ, and the breadcrumbs hypothesis, in which the features most relevant to time step t are already identical to those that would most benefit inference at time t + τ. To evaluate which hypothesis might be correct, we propose a myopic training scheme that does not propagate gradients from the loss at the current position to hidden states from previous positions. We then evaluate the myopia gap in performance between myopically trained and vanilla transformers as a measure of pre-caching.To consider whether language models might directly implement pre-caching, we design a synthetic scenario where the task can only be completed via explicit pre-caching. We configure a task where the model must precompute information for the next token, because otherwise the correct answer could not be accurately computed in a single forward pass. In this synthetic scenario, we find clear evidence that the transformer learns to pre-cache. When transformer-based sequence models must precompute information to minimize loss, they do so.We then consider whether breadcrumbs or pre-caching is demonstrated in natural language models. Our experiments with myopic training suggest that, with small language models like GPT-2 (Radford et al., 2019) , much less pre-caching occurs in this setting, pointing towards the breadcrumbs hypothesis. That is, we claim language models on this scale do not intentionally prepare information for the future to a significant extent. Instead, they compute features that are useful to predicting the immediate next token, which turn out to then be helpful at future steps; there is not a significant tradeoff between greedily optimizing for next token loss and ensuring future predictive performance. However, we also find evidence that the importance of pre-caching increases with scale, becoming non-negligible with larger models, e.g. Pythia 2.8B (Biderman et al., 2023) . This suggests that these larger models are "planning for the future" in a way that small models cannot. 
related work: 
 Future token meta-prediction. Several recent works (nostalgebraist, 2020; Belrose et al., 2023; Pal et al., 2023; Cai et al., 2024) observe that transformer hidden states can be used to predict current and future tokens in a sequence, typically via linear probing. Notably, Hernandez et al. (2024) show that more complicated relationships are encoded linearly in hidden states, such as subject-object relations, implying that future tokens can also be predicted in specific cases. This future token predictivity has also been applied to speeding up inference by decoding future tokens in parallel (Stern et al., 2018; Cai et al., 2024) . Unlike these works, we focus on the question of how the model learns to prepare hidden states that are useful for future prediction, possibly at the expense of current-token predictivity.Probing. Our synthetic data experiments make use of probing, a technique where a simple auxiliary model is used to predict properties from target models' representations (Belinkov & Glass, 2019; Shi et al., 2016; Hewitt & Liang, 2019; Pimentel et al., 2020; Belinkov, 2021) . Probing-based approaches are known to overestimate latent information if the classifier learns to do a task on its own (Belinkov, 2021) , and probing analyses may only be informative when compared to probing a reasonable baseline (Hewitt & Liang, 2019) . In our probing experiments, we avoid these pitfalls by ensuring that the function to be learned cannot possibly be computed by the probe itself.Mechanistic interpretability. Our analysis of transformer models in a synthetic setting relates to the subfield of mechanistic interpretability, which seeks to understand models by isolating and explaining the behavior of their components (Olah et al., 2020; Bau et al., 2020; Meng et al., 2023; Nanda et al., 2023) . Some of these works (Nanda et al., 2023; Li et al., 2023; Zhong et al., 2023) practice mechanistic interpretability by studying models trained on synthetic data. We apply some mechanistic interpretability techniques in a synthetic setting to study the problem of whether language models "think ahead" for future tokens. However, our approach also differs from that of mechanistic interpretability by analyzing the effect of the training procedure on the learned model. 
theory: pre-caching or breadcrumbs?: 
 Consider a generic causal sequence-to-sequence prediction task (x 1 , . . . , x n , y 1 , . . . , y n ) ∼ D, where D is a data distribution supported on X n × Y n for some domains X, Y. The task is to estimate the conditional expectations E D (y i | x 1 , . . . , x i ) for 1 ≤ i ≤ n.2 Note that we recover the autoregressive setting by setting Y = X and y i = x i+1 .Transformer models trained on such tasks have been observed (Pal et al., 2023) to store information in hidden states during inference at position i that is then used in future inference at j > i. However, since the loss associated with each step i depends only on how well the model does at the immediate task of predicting y i , it may not be immediately obvious how this preparation for the future arises. We give names to two competing explanations:• Pre-caching: The model "deliberately" computes and stores features that are expected to be useful for the future, even if they are irrelevant to the present.• Breadcrumbs: The features that most benefit the present inference task are the same as those that are most useful to the future. When the model performs the present forward pass, it "unintentionally" leaves a trace ("breadcrumbs") that is then picked up by future passes.To disentangle these two explanations, we introduce a notion of myopic transformer models, which we show to be incapable of deliberate pre-caching-for these models, the extent to which past features are beneficial to the future is decided purely by the breadcrumbs explanation. Thus, the gap between vanilla and myopic transformer models is a quantitative measure of how much pre-caching is taking place.Suppose, for the sake of exposition, that the transformer model G uses independent parameters for each position. 3 Let p be the parameter count of each forward pass of G. Then, lettingθ i ∈ Θ = R p be all parameters used by G at position i, a transformer G is a parameterized function G : X n × Θ n → Y n , (x 1 , . . . , x n ; θ 1 , . . . , θ n ) → ( ŷ1 , . . . , ŷn ). For 1 ≤ i ≤ n, let G i (x 1 , . . . , x n ) ∈ Y bethe output of G's ith forward pass. Because of the causal masking within G, this depends only on x 1 , . . . , x i and θ 1 , . . . , θ i . That is, with slight abuse of notation, we may write ŷi = G i (x 1 , . . . , x n ; θ 1 , . . . , θ n ) = G i (x 1 , . . . , x i ; θ 1 , . . . , θ i ).Now, letting L : Y × Y → R + be some choice of loss function, the expected loss ℓ of a transformer model with parameters θ 1 , . . . , θ n isℓ(θ 1 , . . . , θ n ) := E (⃗ x,⃗ y)∼D n ∑ i=1 L(G i (x 1 , . . . , x i ; θ 1 , . . . , θ i ), y i ) =: n ∑ i=1 ℓ i (θ 1 , . . . , θ i ),the sum over 1 ≤ i ≤ n of the expected loss ℓ i at position i. (We suppress the dependence on G and D for concision.) In practice, we always tie the weights across position. That is, all θ i are set equal to the same θ ∈ Θ. Then, by the chain rule,∇ θ ℓ(θ, . . . , θ) = n ∑ i=1 ∇ θ i ℓ(θ 1 , . . . , θ n ) θ 1 =...=θ n =θ = n ∑ i=1 n ∑ j=i ∇ θ i ℓ j (θ 1 , . . . , θ j ) θ 1 =...=θ n =θ,a sum over an upper-triangular expected Jacobian "matrix". The off-diagonal terms i < j, corresponding to the expected gradient of the model's future loss at position j with respect to its weights at position i, are the training signals that encourage pre-caching.We say a model is myopic when each forward pass G i optimizes only ℓ i without regard for future ℓ j at j > i. In the untied weights case, the right definition is then apparent.Definition 1. The parameters ( θ1 , . . . , θn ) ∈ Θ n are untied-myopic if they satisfy θi ∈ arg minEQUATIONDefinition 2. Let M be the feasible set of the constraints in Equation 1. The untied myopia gap is the smallest possible gap between the expected loss attained by a myopic model and the optimal model: p * := minEQUATION)In the tied weights case, it is perhaps not immediately clear what the right definition of myopia should be. It does not suffice to simply constrain the minimizations in Equation 1to θ1 = . . . = θn , since min θ ℓ i (θ, . . . , θ) is optimizing for pre-caching (the dependence on arguments j < i) as well as the present inference (the dependence on argument i). Instead, the right notion is a choice of tied parameters such that the model is, aggregated over positions, optimal for the present task when conditioned on a fixed past. That is, forward passes do not compute features for the future if they can compute other features more beneficial to the present.Definition 3. The parameters θ ∈ Θ are (tied-)myopic if they satisfyθ ∈ arg min θ∈Θ n ∑ i=1 ℓ i ( θ, . . . , θ, θ).(3)The (tied) myopia gap is then defined analogously to Definition 2.The breadcrumbs hypothesis states that the myopia gap is small-near-optimal performance can be attained even when each forward pass is computing features relevant to only its own immediate inference task, with no regard to pre-caching for the future.If the breadcrumbs hypothesis does not hold, we say that the model is pre-caching. It is important to remember that the ℓ i depend on a choice of transformer model G and dataset D. That is, breadcrumbs and pre-caching are properties of the model architecture and the data considered as a whole.Although a small myopia gap reveals that one can do just as well without pre-caching, it does not say much about any specific model. To measure pre-caching within a given model, we examine the extent to which its parameters violate the myopia constraints. Definition 4. The (tied) local myopia bonus at θ * ∈ Θ isĉ(θ * ) := max θ∈Θ n ∑ i=1 (ℓ i (θ * , . . . , θ * ) -ℓ i (θ * , . . . , θ * , θ)) ≥ 0.For further interpretation of the myopia gap and myopia bonus, see Appendix A.Our heuristic remark in Section 3.2, that the off-diagonal gradient terms are responsible for pre-caching, is justified by Theorem 13 below. It states that, given certain regularity conditions on the loss terms ℓ i , performing gradient descent with the off-diagonal terms removed results in a myopic model in the sense of Definition 3. We call this myopic descent.For myopic descent to be stable in the tied-weights case, we need, roughly speaking, for the model to depend more on the parameters associated with the present forward pass than those from the past. This is a plausible condition-dependence on the past is mediated purely by the attention mechanism, while the present forward pass depends on both attention and feedforward parameters. The precise condition we use is forward bias (Definition 11); see Appendix C for details and the proof of the theorem.Theorem 13. Let f ( θ, θ) := ∑ n i=1 ℓ i ( θ, . . . , θ, θ).If f is forward-biased, σ-strongly convex, and L-smooth, then, for some step size η > 0, the iterates of myopic descent with tied weights θ (t+1) = θ (t) -η∇ θ f ( θ, θ) θ=θ=θ (t)(x n ) N n=1 , (y n ) N n=1 , (z n ) N n=1where, for each n,• x n ∼ N (0, 1) (standard Gaussian)• z n ∼ Ber(p) (Bernoulli with probability p)• y n = z n ∑ a i=1 sin(bx n-i ) + (1 -z n )x nand {x n } n∈N ∪ {z n } n∈N are mutually independent. In our experiments, we always set the parameters a = b = 10 and N = 64, so for convenience notate D p := D 64 p,10,10 .The intuition is that a transformer regression model G trained on D p would benefit from pre-caching sin(bx n ) during its forward pass at position n, even though this computation is irrelevant to its task of predicting y n . One simple strategy that makes use of this pre-caching is Algorithm 1. 4The motivation for the Bernoulli variables z i is that, as p decreases, the expected first time when sin(bx n ) becomes useful advances further into the future. In addition, when p is sufficiently small, the probability (1p) a that the value sin(bx n ) is never useful at all becomes non-negligible. We will show that, even in this case, the transformer model learns to pre-cache.Investigating myopia. Suppose that we train a myopic model (Section 3.4) on the same task. Since this model lacks off-diagonal gradient terms, we do not expect it to learn to pre-cache sin(bx n ) at position n. One possible strategy that does not use pre-caching is Algorithm 2. We expect this brute force algorithm to perform significantly worse given the same parameter count-it computes an a-dimensional nonlinear function within a single layer, while each layer of Algorithm 1 computes only scalar nonlinear functions. := z n ∑ a i=1 sin(bx n-i ) + (1 -z n )x n . return ŷn . 
evaluation: linear probing: 
 To determine if the transformer model is computing sin(bx n ) at position n, we fit linear probes on the hidden states. We additionally compute the correlations between sin(bx n ) and each individual dimension (i.e., each neuron) of each hidden state.6 See Section D.1.1 for details.For varying p, we train two-layer transformer models with embedding dimensions of 128 on D p using using both ordinary and myopic gradient descent. Full architecture and training details are provided in Section D.1.Examining the performance of each linear probe against sin(bx n-i ) for varying i, we find strong evidence that the transformer model with vanilla training is indeed pre-caching sin(bx n ), possibly in order to implement Algorithm 1. Indeed, in Figure 2,  • The zeroth hidden state (i.e., the sum of the input and position embeddings) at position n is correlated with only x n .• The first hidden state is correlated with sin(bx n ) but not correlated with any sin(bx n-i ) for i > 0. Further, looking at the per-neuron correlations in Figure 3 , we see that sin(bx n-i ) for 1 ≤ i ≤ a are all correlated with a single 1-d subspace of the second hidden state (they share the same striping pattern); this is the subspace storing ∑ a i=1 sin(bx n-i ). Mean- while, sin(bx n ), as well as many of the x n-i , are located in various other 1-d subspace of the second hidden state; these terms are all left over in the residual stream from previous layers, and are cleaned up only by the output unembedding.On the other hand, in Table 1 , the myopic models perform significantly worse. The perneuron correlations in Figure 3 suggest that the myopic model may be implementing a crude approximation of Algorithm 2. This suggests that the synthetic setting has an inherently high myopia gap-it is impossible for the transformer model to do well without pre-caching.In addition to the above D p synthetic data setting, we also measure the myopia gap on the task of natural number multiplication. In particular, we find evidence suggesting that pre-caching is responsible for model computation on filler tokens, in the sense of Pfau et al. (2024) . See Appendix D.3 for details. 
gpt-2's myopia gap: 
 In order to measure the extent to which transformer models learn to pre-cache on natural language data, we estimate both the myopia gap (Definition 3) in this setting as well as the local myopia bonus (Definition 4) of a transformer model with vanilla pre-training. Experiments in this subsection use the 124M-parameter GPT-2 architecture; see Table 4 in Appendix D for configuration details.We train all models (vanilla and myopic) from random initialization for one epoch on 4.6M sequences from the MS MARCO dataset (Nguyen et al., 2016) , truncated to length 64. To estimate the local myopia bonus of the vanilla model, we train another model from random initialization with the same architecture, but with past hidden states sourced from the frozen As baseline, we also train a "transformer bigram" model, a model with an identical architecture but all off-diagonal key/value states zeroed out.From Table 2 , the estimated myopia gap in this setting is 3.40 -3.28 = 0.12 cross entropy, while the local myopia bonus of the vanilla model is 3.28 -3.26 = 0.02.The nonzero myopia gap suggests that pre-caching may provide a small positive benefit. Indeed, in Figure 4 , we see that the myopic model outperforms the vanilla model at the beginning of the sequence, since it can allocate all compute to next-token prediction, but quickly falls behind as the length of the past increases, since it suffers from a lack of precached information from earlier forward passes. However, note that this gap is much smaller than that between the vanilla model and the transformer bigram model (Table 2 ). That is, the myopic model is still able to leverage past information (breadcrumbs) to a significant extent, even if they optimized only for the present inference task. That the local myopia gap is near zero further supports this direction-the model learned through vanilla training does not trade off significantly between features useful for the present and pre-caching for the future.Myopic model One might suppose that the relatively small myopia gap of GPT-2 is due to the relative simplicity of the small architecture we consider, and that larger language models might exhibit a more pronounced myopia gap.To test this, we train both vanilla and myopic transformers from the Pythia LLM suite (Biderman et al., 2023) , ranging in size from 14M to 2.8B parameters, on one epoch of 10M sequences of 64 tokens each subsampled from the Pile dataset (Gao et al., 2020) . (We use the same subsampled dataset for every training run.) We report validation cross-entropy loss (Figure 7 in Appendix D) as well as performance on a variety of natural language experiments (Figure 5 ). Note that, unlike in the GPT-2 experiments (Section 5.1), which start from random initialization, we start all training for Pythia models from the pre-trained checkpoints provided by Biderman et al. (2023) -for the larger architectures, the 10M sequence dataset we use is not sufficiently large to use for pre-training from random initialization. 
discussion and future work: 
 Using a synthetic dataset, we demonstrate that pre-caching can indeed be learned by a transformer model. On the other hand, our experiments with natural language suggest that the breadcrumbs hypothesis is more explanatory for that setting, especially with smaller models, but that the importance of pre-caching increases with scale.If the myopia gap is indeed not large in practice, there may be several applications of myopic training. We hypothesize that myopic transformers may have advantages in terms of safety and/or interpretability-it may be easier to understand what a model is doing if we know that everything that it is computing on each forward pass is directly towards the goal of predicting the immediate next token. For example, as seen in Appendix D.3, myopic models may not be able to make use of computation on the forward passes of filler tokens in the sense of Pfau et al. (2024) .Another possibility is that of automatically swapping in a locally myopic model (Section 5.1) on forward passes where we detect it is beneficial to sacrifice future performance in favor of immediate next-token accuracy (for example, on especially important tokens, or near the end of a text). We leave these possible applications to future work.Notice that the myopia gap consists of two pieces: a myopia bonus, the improvement that can be obtained at the current forward pass by ignoring the future forward passes; and a myopia malus, the cost to the future forward passes that is incurred by not pre-caching for them. To be precise, in the untied case,9 given a choice of myopic θ1 , . . . , θn satisfying constraints (1) of Definition 1, write ℓ( θ1 , . . . , θn )minθ 1 ,...,θ n ℓ(θ 1 , . . . , θ n ) = n ∑ i=1 min θ i+1 ,...,θ n ℓ( θ1 , . . . , θi-1 , θi , θ i+1 , . . . , θ n ) -min θ i ,...,θ n ℓ( θ1 , . . . , θi-1 , θ i , θ i+1 , . . . , θ n ) = n ∑ i=1 ℓ i ( θ1 , . . . , θn-1 , θi ) -ℓ i ( θ1 , . . . , θn-1 , θ * i i ) + n ∑ i=1 n ∑ j=i+1 ℓ j ( θ1 , . . . , θi-1 , θi , θ * i+1 i+1 . . . , θ * i+1 n ) -ℓ j ( θ1 , . . . , θi-1 , θ * i i , θ * i i+1 , . . . , θ * i n ) .where we defineθ * i i , . . . , θ * i n ∈ arg min θ i ,...,θ n ℓ( θ1 , . . . , θi-1 , θ i , θ i+1 . . . , θ n ).That is, the myopia gap is the sum c + d = ∑ i c i + ∑ i d i of the myopia bonuses c i ( θ1 , . . . , θn ) := ℓ i ( θ1 , . . . , θn-1 , θi ) -ℓ i ( θ1 , . . . , θn-1 , θ * i i ) ≤ 0, Published as a conference paper at COLM 2024 (the inequality following from the myopia constraints (1)), and the myopia malusesd i ( θ1 , . . . , θn ) := n ∑ j=i+1 ℓ j ( θ1 , . . . , θi-1 , θi , θ * i+1 i+1 . . . , θ * i+1 n ) -ℓ j ( θ1 , . . . , θi-1 , θ * i i , θ * i i+1 , . . . , θ * i n ) ≥ 0, with d i + c i ≥ 0 for each 1 ≤ i ≤ n by the definition of the θ * i j .A priori, a small myopia gap does not necessarily imply a small (in magnitude) myopia bonus c and malus d. Indeed, in the case when the myopia gap is small, a large value for c (and thus a corresponding large value for d) means precisely that the transformer model is committing significant resources to pre-caching that could otherwise have been used to improve inference on the current position. On the other hand, it is possible that both c and d are small; that is, there is not much cost associated with pre-caching for the future, as the present forward pass already results in information (breadcrumbs) useful for that purpose.In practice, the myopia bonus may be difficult to estimate, as it depends on the result of O(n) separate optimization problems (each of which, in practice, is a full transformer model training run). Thus, we instead compute the local myopia bonus of Definition 4.The dependence of forward pass G i on previous forward passes G j for j < i is mediated through hidden states h 1 , . . . , h i-1 :G i (x 1 , . . . , x i ; θ 1 , . . . , θ i ) = Gi (h 1 , . . . , h i-1 , x i ; θ i )where the h i are themselves recursively defined parameterized functions h i = h i (h 1 , . . . , h i-1 , x i ; θ i ). (Note that we are making a choice here to consider transformers as functions of hidden states, and not their key/value states. This has implications for how myopic descent is defined: when hidden state h j is attended to by forward pass i > j, we consider the key and value weights W K and W V , respectively, to belong to forward pass i. Thus, they are updated by the gradient wrt. ℓ i .)With the hidden states explicitly written out, the gradient wrt the loss is a sum over all possible paths to the present: for j < i,∂G i ∂θ j (x 1 , . . . , x i ; θ 1 , . . . , θ i ) = ∑ j=i 1 <...<i m <i ∂h j ∂θ j ∂ Ĝi h i m m-1 ∏ k=1 ∂h i k ∂h i k-1 .where the sum is over all partitions j = i 1 < . . . < i m < i.An important primitive that we use when implementing the myopic gradient descent of Section 3.4, the local myopic bonus of Definition 4, and the transformer bigram in Section 5.1 is an attention mechanism that uses key/value states for past forward passes differing from those it uses for the current pass, while still computing all forward passes in parallel. We call our construction the myopic attention mechanism. We use it to implement several distinct transformer training methodologies:• When training with myopic descent, the past key/value states are the result of key/value weights W K and W V , respectively, multiplying a cloned and detached copy of the previous hidden states. Let X = (x 1 , . . . , x n ) ⊤ ∈ R n×d be the sequence of residual stream hidden states per position, with each row representing one position's hidden state in R d . Let W Q , W K , W V ∈ R d×h be the query, key, value weight matrices for one attention head, of dimensionality h, in the transformer G. DenoteQ := XW Q , K := XW K , V := XW Vand let Q, K, Ṽ be the alternate states we wish to use for off-diagonal attention terms. We adopt the convention that lowercase letters with subscripts represent rows of matrices; e.g. q i is the ith row of Q. For simplicity of presentation, we omit causal masking; the modifications that should be made in the presence of a mask are straightforward.Recall that the vanilla attention mechanism for G isY = σ(QK ⊤ )V,where σ is row-wise softmax. Writing this out token-wise,y i = Z -1 i n ∑ j=1 exp(q ⊤ i k j )v j ,where Z i is the partition functionZ i := n ∑ j=1 exp(q ⊤ i k j ).The myopic attention mechanism, on the other hand, is written tokenwise asỹi = Z-1 i exp(q ⊤ i k i )v i + ∑ j̸ =i exp(q ⊤ i kj ) ṽj = Z-1 i n ∑ j=1 exp(q ⊤ i kj + δ ij q ⊤ i (k j -kj ))( ṽj + δ ij (v j -ṽj )) = n ∑ j=1 a ij ṽj - n ∑ j=1 δ ij a ij (v j -ṽj )where Zi := exp(q⊤ i k i ) + ∑ j̸ =i exp(q ⊤ i kj ) = n ∑ j=1 exp(q ⊤ i kj + δ ij q ⊤ i (k j -kj )), a ij := Z-1 i exp(q ⊤ i kj + δ ij q ⊤ i (k j -kj )), and δ ij is the Kronecker delta. Now, notate(diag A) ij := δ ij A ij .That is, diag A is the diagonal matrix that has the same entries as A along the diagonal and is zero elsewhere. We are now able to write the myopic attention mechanism in matrix form:Ỹ = A Ṽ + (diag A)(V -Ṽ ), where A := σ(Q K⊤ + diag(Q(K ⊤ -K⊤ ))).We use two assumptions on the loss function to be minimized. These are standard in the first-order methods literature. Definition 6. A function f : R n → R is called L-smooth if it is continuously differentiable with L-Lipschitz gradient. That is, for all x, y in the domain,∥∇ f (x) -∇ f (y)∥ 2 ≤ L∥x -y∥ 2 .Definition 7. A function f : R n → R is called σ-strongly convex if, for all x, y in the domain,f (y) ≥ f (x) + ∇ f (x) ⊤ (y -x) + σ 2 ∥y -x∥ 2 2 .In particular, recall that strong convexity implies the existence of a unique minimum. Hence, it makes sense to write, for example, x * = arg max x f (x) without ambiguity.Theorem 8. Assume ℓ : Θ n → R is σ-strongly convex and L-smooth for some σ, L > 0. Consider ordinary gradient descent with untied weightsθ (t) i = θ (t-1) i -η∇ θ i ℓ(θ (t-1) 1 , . . . , θ (t-1) i ) ∀i ∈ {1, . . . , n}.Then, for θ * 1 , . . . , θ * n = arg min θ 1 ,...,θ n ℓ(θ 1 , . . . , θ n ), for small enough η > 0,∥θ (t) i -θ * i ∥ 2 2 ≤ 1 - 2ησL σ + L t ∥θ (0) i -θ * i ∥ 2 2 ∀i ∈ {1, . . . , n}.Proof. This is a standard convergence result for gradient descent on strongly convex functions. For example, see Nesterov (2018) .Theorem 9. Assume ℓ is σ-strongly convex and L-smooth, and consider ordinary gradient descent with tied weightsθ (0) 1 = . . . = θ (0) n θ (t+1) i = θ (t) i + η n ∑ i=1 ∇ θ i ℓ(θ 1 , . . . , θ n ) ∀i ∈ {1, . . . , n}.There exists η > 0 such that ∥θ(t) i -θ * ∥ 2 2 ≤ 1 - 2 √ nησL σ + L t ∥θ (0) i -θ * ∥ 2 2 ∀i ∈ {1, . . . , n}.where θ * = arg min θ ℓ(θ, . . . , θ).Proof. This is again the standard convergence result, now applied to descent with step size √ nη on the σ-strongly convex L-smooth function θ → ℓ(θ/ √ n, . . . , θ/ √ n). Alternatively, one may think of this as projected gradient descent constrained to the subspace θ 1 = . . . = θ n with a step size of √ nη. Projected gradient descent inherits the same convergence properties as unconstrained gradient descent (Beck, 2017) .Theorem 10. Assume each of the ℓ 1 , . . . , ℓ n are σ-strongly convex and L-smooth. Consider myopic gradient descent with untied weightsθ (t) i = θ (t-1) i -η∇ θ i ℓ i (θ (t-1) 1 , . . . , θ (t-1) i ) ∀i ∈ {1, . . . , n}.There exists η > 0 such that θ i t→∞ --→ θi for all i, where θ1 = arg minθ 1 ℓ 1 (θ 1 ) θ2 = arg min θ 2 ℓ 2 ( θ1 , θ 2 ) . . . θn = arg min θ n ℓ n ( θ1 , θ2 , . . . , θn-1 , θ n ).Proof. Let θ1 , . . . , θn be as in the theorem statement. We proceed by induction. For the base case, note that the myopic descent iterates for θ (t)1 are independent of θ (t) j for j > i. Thus, the standard convergence theorem gives that θ(t) 1 → θ1 as t → ∞.(t) i t→∞ --→ θi for all i < k. Thus, for any ε > 0, for sufficiently large t, ∥(θ(t) i ) k-1 i=1 -( θi ) k-1 i=1 ∥ 2 < ε. Hence, since ℓ k is L-smooth, for any θ k , ∥∇ θ k ℓ k (θ (t) 1 , . . . , θ (t) k-1 , θ k ) -∇ θ k ℓ k ( θ1 , . . . , θk-1 , θ k )∥ 2 < Lε. Expanding and rearranging, ⟨∇ θ k ℓ k (θ (t) 1 , . . . , θ (t) k-1 , θ k ), ∇ θ k ℓ k ( θ1 , . . . , θk-1 , θ k )⟩ ≥ 1 2 ∥∇ θ k ℓ k ( θ1 , . . . , θk-1 , θ k )∥ 2 2 - 1 2 L 2 ε 2is bounded away from zero as long as, say,EQUATIONk-1 , θ k ) is a descent direction for ℓ k ( θ1 , . . . , θk-1 , θ k ). This is a sufficient condition for θ k to converge to a (L + 1)ε σ -neighborhood of θk given a small enough step size (Beck, 2017) . Since ε > 0 is arbitrary, this completes the inductive step.Definition 11. A function f (x, y) : R a × R a → R with continuous second derivatives is ρforward-biased if, for all y ∈ R a and x = y, H y,y f (x, y) + H x,y f (x, y) ≻ 0 and κ(H y,y f (x, y) + H x,y f (x, y)) < ρ, where κ is the condition number, and we write the Hessian of f as a block matrix: Lemma 12. If A ∈ R a×a is positive definite with κ(A) < ρ, then there exists some η > 0 such that Iη A is a (1ρ -1 )-contraction. Explicitly, for any y ∈ R a , ∥(Iη A)y∥ 2 ≤ (1ρ -1 )∥y∥ 2 .H f (x, y) = H x,x f (x, y) H x,y f (x, y) H y,x f (x, y) H y,y f (x, y) =         ∂ f (xProof. By symmetry, E[sin(bx)] = 0. We calculate Cov(x, sin(bx)) = (2π) -1/2 ∞ -∞x sin(bx)e -x 2 /2 dx = be -b 2 /2 .We already computed the variance in Lemma 14, so ρ(x, sin(bx)) = Cov(x, sin(bx))Var(x)Var(sin(bx)) = 2be 3b 2 /2 e 2b 2 -1 .In our experiments we set b = 10, so ρ(x, sin(bx)) < 10 -20 .We use a smaller version of the GPT-2 architecture, adapted to regression tasks. That is, the token embedding and unembedding layers are replaced with a trained linear map from the input space to the embedding space and from the embedding space to the output space, respectively. For each p ∈ {0.01, 0.1, 0.3, 1} models are trained using ordinary and myopic descent on one epoch of 30M sequences of length 64 sampled from D 64 p,10,10 . See Table 3 Given a transformer model trained on D p , we sample the hidden states at each layer when the model is given as input 50 000 evaluation sequences from the same distribution D p .For each layer and targets sin(bx n-i ) for varying i > 0, we fit a linear regression model on the hidden state of that layer to the target. The in-sample R 2 of each linear model is then reported in Figure 2 . Figure 6 is a visualization of the linear probe's performance on the vanilla transformer. For both vanilla and myopic training, we train the GPT2-small architecture from random initialization for one epoch on 4.6M sequences from the MS MARCO dataset (Nguyen et al., 2016) , truncated to length 64. To estimate the local myopia bonus of the vanilla model, we train another model from random initialization with the same architecture, but with past hidden states provided by the vanilla model. Note that this "local myopic" model attains slightly better performance than the vanilla model; each forward pass can focus purely on next-token prediction, since past hidden states are supplied by a separate model. As a baseline, we also train a "transformer bigram" model, which is a transformer model whose key/value states are zeroed out during training and evaluation. See Table 4 for configuration details.Code to reproduce our results can be found at https://github.com/wiwu2390/FutureGPT2-publicFor classification tasks, we are typically interested in the conditional probabilities Pr D (y n = c | x 1 , . . . , x n ) for each class c. However, this can be subsumed into the generic case by letting Y be the probability simplex over all classes.For example, this is true of absolute position embedding weights.We think of each transformer layer as a read operation, performed by the attention mechanism, followed by a compute operation, performed by the feedforward block.For example,Shen et al. (2022) provide upper bounds on MLP error that degrade exponentially in dimensionality given a fixed parameter and layer count.Note that, in order for linear probing to be meaningful, we must first ensure that there is no pre-existing linear relationship between the inputs and the quantities we are probing for. Since the (x n , z n ) n are mutually independent, this follows from Lemma 15, stating that x n and sin(bx n ) have near-zero correlation for large enough b. In our experiments, we set b = 10, in which case ρ(x n , sin(bx n )) < 10 -20 . In other words, there is low predictive V-information from the inputs to the target sin(bx n ), where V is the class of linear models(Xu et al., 2020).Note that this "local myopic" model attains slightly better performance than the vanilla model; each forward pass can focus purely on next-token prediction, since past hidden states are supplied by a separate model.We use a sliding window overPG-19 (Rae et al., 2019) samples, which comprise longer sequences, in order to reduce noise in our per-position loss estimates. The distribution shift between MS MARCO and PG-19 does not affect the relative comparison of myopic and vanilla GPT-2.The tied case is analogous, so we do not write it explicitly. 
 
 Appendix:  


Decision: This paper analyses if/how ("small", 124M) LLMs plan ahead, i.e. their hidden state contains information useful for only the next token (breadcrumbing), or for future time steps (pre-caching).
Experiments with modified/myopically restricted model unable to pre-cache and a standard transformer on a synthetic task are used to shed light on this, but do leave many open questions.

All reviewers (URPJ, 4NSr, ZSJf) were positive in general, and mentioned the "valid analysis" (URPJ), say that the myopia gap is of interest (4NSr), and that it is "interesting", and that the experimental setup is good (ZSJf). The critiques are the that there could be more NLP discussion (URPJ), that the scale is too small for broad conclusions (4NSr, ZSJf), which was countered by the authors providing extra results in the rebuttal phase as well.

All in all the paper makes for a good contribution despite the (initial) small scale, so in my view the pros outweigh the cons, especially since some weaknesses could be mitigated for camera ready.

[comments from PCs] Please refine and incorporate the results mentioned by the AC.

Reviews: review: paper_summary: paper_summary: The paper investigates the idea of a transformer-based LM taking decisions by "planning ahead" (pre-caching) or merely making use of prior states that were deemed optimal for the prior decisions (breadcrumbs). They do this by implementing a reduced "myopic" transformer which does not update prior states during updates with later decisions during training.

The authors perform a synthetic experiment which shows that a myopically restricted model is unable to pre-cache prior results whereas a standard transformer does for this specifically designed synthetic task.
The authors then train both kinds of models as well as a model which strictly limits attention to bigramming on natural language data and find that the myopic model performs much closer to the standard than the bigram model. The authors take this as an indication that a transformer is largely breadcrumbing and that the performance increase from using a model that is able to pre-cache is only moderate, which leads to the conclusion that little pre-caching is going on.
strength_weakness: strength_weakness: The implications of this research, in particular whether a model might be pre-caching or breadcrumbing remain open. 
The paper would be much more approachable and engaging if it could elaborate on what's hinted at in Figure 1 (which is not referenced from the text!), i.e., that humans start pondering about computations (and the like) in a more structured form. Is there also evidence that this is helpful? Maybe could one come up with ways of training "focally" instead of "myopically", i.e. lay caches in certain ways that then yield better results?

The paper is nice theoretically but I would like to see a closer discussion of NLP topics. For example, when you discuss future token meta-prediction, it would help to highlight that of course next tokens can be predicted from state as that's what's the whole point of language modelling. I presume that the cited works (nostalgebraist is an incomplete citation!) take precautions (synthetic data or in the evaluation?) that future-token prediction doesn't simply stem from the fact that a next-to-next token will clearly depend on the next token and can hence be predicted from the next token's hidden state with some accuracy.
questions: questions: See above. Some remarks:
- Figure 4 is not referenced from the text. (It an also probably be omitted.) I'd put Table 1 earlier, into Section 4.
- When you refer to Table 4, it would help to remark that it can be found in Appendix F.
- There's an Appendix G which is not referenced from the main text. What does it do?
- I share your belief that humans "pre-cache" in some way. However, is there reason to believe that boosting this behaviour (beyond what happens in your standard transformer) would lead to better language modelling results? If so, in what ways? Would it be possible to design corresponding follow-up work? Can you measure this in ways that are different from cross-entropy? (E.g., would it help an LM to pre-compute when it sees the second addend rather than think about what type of second object (apple) it is?
paper_summary: paper_summary: This paper investigates whether the hidden states of large autoregressive language models encode information solely useful to predict future tokens (pre-caching hypothesis), or whether any such information is a side effect of being able to predict the following word accurately (breadcrumb hypothesis).

It proposes a quantity called the myopia gap: the difference in loss between a globally optimal LM and one where the parameters at each time step are greedily optimized for the loss at that position alone. It proposes to estimate the myopia gap by comparing the loss of a typically trained model, and one where the token-level loss is backpropagated only with respect to the weights at the current position. While weights at each time step are tied, the paper engages in a thought experiment, first supposing parameters to be untied across time steps, and develops "myopic gradient descent" to optimize LM parameters only for predicting the subsequent position.

The paper constructs a synthetic data set, and demonstrates that the myopia gap of a 2-layer transformer trained on this data is high (by construction). By applying a linear probe, it also demonstrates that the representation of the first hidden layer is strongly correlated with a value consistent with the pre-caching hypothesis, which is not true if the model is trained myopically. Training an $\approx$100M parameter LM on MS MARCO shows that the myopia gap exists for larger LMs, although the difference is slight, and is a function of the token position as well.
strength_weakness: strength_weakness: It is unclear from the paper what constitutes a significant amount of myopia gap. How does one test the hypothesis that a model is explicitly using pre-caching? The paper does not answer this explicitly, just alludes to the fact that myopically trained model incurs a penalty of $\approx0.12$ in terms of cross-entropy loss. In addition, the actual model ($\approx$100M parameters) and training data ($\approx$250M tokens) sizes are modest. It is unclear whether the natural language findings regarding myopia gap would generalize to more realistic LLMs, trained to convergence on larger corpora.
questions: questions: - One may consider discussing work on blockwise parallel [1]/non-autoregressive decoding, which explicitly train models to predict a block of subsequent tokens in parallel.
- Why use PG-19 for evaluation in Figure 5 when MS MARCO was used as the training set? Wouldn't this be conflating domain shift with general LM quality?
- What are the practical implications of whether the model is employing pre-caching vs. breadcrumbs, conditioned on similar model performance? Is one a priori more desirable?
- Why does the myopically-trained model incur lower loss on the first few token positions than the standard LM (Figure 5)?
- Appendix C: tilde is used inconsistently (e.g., when defining Q, K, V, and also over X).
- Appendix G: This appears to be orthogonal to the experiments in the body text and is never referred to from the body. It is unclear how this relates to the myopia gap, or if this was a separate effort at understanding the tendency of LLMS to pre-cache.

[1] Stern, Mitchell, Noam Shazeer, and Jakob Uszkoreit. "Blockwise parallel decoding for deep autoregressive models." Advances in Neural Information Processing Systems 31 (2018).
paper_summary: paper_summary: This paper studies whether Transformer language models encode information in middle layers that is useful for future token predictions, or are they mostly "greedy" in a sense that they learn features relevant for the current token prediction and these feature end up being useful for future predictions. The authors formalize this research question and call the two hypotheses "pre-caching" and "breadcrumbs" respectively. They go ahead and design experiments for this question by avoiding the gradient updates to previous timesteps and measuring the gap in perfomance between optimal models trained with past awareness and without. On a synthetic dataset that requires continuously memorizing the past for predicting the next token, the authors find evidence for pre-caching. On small scale LM training, it appears that the breadcrumbs case is more likely.
strength_weakness: strength_weakness: Overall I think the study is interesting and worth sharing with the community. Yet, I would note that the LM experiments are on only 124M GPT-2 models 4.6M texts of length 64 from MS MARCO. This is much smaller in model size, context length and token scale compared to today's LLMs. So the conclusion of natural language learning following the breadcrumbs hypothesis might only hold for this limited studied setting and not hold for realistic settings. I'd be curious to know this, but given the cost of such an experiment I understand that it might be infeasible to include.
questions: questions: Other than studying the behavior of models, do you think there are any practical conclusions? for example, if the breadcrumbs hypothesis is true, maybe some more efficient pretraining schema is possible?

Comments: title: To all reviewers: details of new experiments
comment: thank you for your response
Thanks to all reviewers. Here are some additional details on the new results we intend to include in the final revision of the paper. Note that all experiments have already been run, but we only include a subset of them in this comment due to space considerations.

**Scaling of myopia gap on Pythia family**

We investigate the myopia gap on a range of model sizes from the Pythia family, ranging from 70M to 2.8B parameters. In particular, we are interested in how the myopia gap scales with model size — multiple reviewers suggested that the insignificant gap observed in GPT-2 may be explained by it being a relatively small model. Our experiments reveal a trend towards myopia gap increasing with scale; this suggests that pre-caching is indeed more explanative of the behavior of larger models. Further, we measure the gap not only on validation cross-entropy loss but also a suite of standard NLP benchmarks (LAMBADA, PIQA, WinoGrande, SciQ, ARC), which may be a more realistic measure of the role of pre-caching in real-world language tasks. Comparisons of the myopia gap between these benchmarks may offer insight into what kinds of tasks benefit more from pre-caching.

As a sample, here are the results on the LAMBADA benchmark [1]:

size | 14M | 31M | 70M | 160M | 410M | 1B | 1.4B | 2.8B
---  | --- | --- | --- | ---  | ---  | -- | ---  | ---
vanilla accuracy | 0.06 | 0.10 | 0.19 | 0.28 | 0.43 | 0.47 | 0.54 | 0.59
myopic  accuracy | 0.02 | 0.07 | 0.13 | 0.21 | 0.29 | 0.35 | 0.38 | 0.40
myopia gap       | 0.04 | 0.03 | 0.06 | 0.07 | 0.14 | 0.12 | 0.16 | 0.19

Note that, unlike in the case of GPT-2, we are unable to train the Pythia models from random initialization (due to limitations in compute). Thus, our methodology is to start from the provided pre-trained checkpoints, and to fine-tune using either vanilla or myopic gradient descent on a further 640M tokens subsampled from the Pile dataset.

**Multiplication and filler tokens**

We also train and evaluate the GPT-2 architecture, using either vanilla or myopic gradient descent, on the task of integer multiplication. We again observe a subtle but nonzero myopia gap in this setting.

Furthermore, we investigate the role of pre-caching in computations on the forward pass of filler tokens, in the sense of Pfau et al [2]. We train vanilla and myopic transformers on multiplication of up to five-digit numbers, but with each multiplicand zero-padded either to length 5 or length 10. Our experiments show that vanilla transformers perform better on this task given more zero-padding tokens, while myopic transformers actually do worse. We interpret this as showing that pre-caching is necessary for making effective use of these filler tokens — without pre-caching, these filler tokens are mere noise, and actually degrade performance. The below table contains the accuracies of both models on multiplication of two five-digit multiplicands, each either zero-padded to length 5 or length 10.

model | pad5 acc. | pad10 acc.
-- | -- | --
vanilla | 0.40 | 0.99
myopic  | 0.39  | 0.05

[1] Paperno et al. "The LAMBADA dataset: Word prediction requiring a broad discourse context."

[2] Jacob Pfau, William Merrill, Samuel R. Bowman. “Let's Think Dot by Dot: Hidden Computation in Transformer Language Models.”