Do Language Models Plan Ahead for Future Tokens?
Authors: Wilson Wu
John Xavier Morris
Lionel Levine
Venue: COLM (2024)
Keywords: probing, mechanistic interpretability, future tokens, transformers
TL;DR: We propose a theoretical notion of transformer models "pre-caching" (preparing information useful for future token inference); we exhibit pre-caching in a synthetic data setting and measure its occurence in language models.

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.

Published as a conference paper at COLM 2024
Do Language Models Plan Ahead for Future Tokens?
Wilson Wu
John X. Morris Department of Computer Science Cornell University
Department of Mathematics
University of Colorado Boulder
wiwu2390@coloraade edu
jxm3@cornell.en
Lionel Levine
Department of Mathematics
Cornell University
levine@math.
cornell. Levine@math.coreell.ccoo
Abstract
Do transformers "think ahead" during inference given position? It is known transformers prepare information in the hidden states of the forward pass time step that is then used in future forward passes f+ T. We posit two explanations for this phenomenon: pre-caching, in which off-diagonal gradient terms present during training result the model computing features at irrelevant to the present inference task but useful for the future, and breadcrumbs, in which features most relevant to time step are already the same those that would most benefit inference at time + T. We test these hypotheses by training language models without propagating gradients to past timesteps, scheme we formalize myopic fraining. In constructed synthetic data setting, we find clear evidence for pre-caching. In the autoregressive language modeling setting, our experiments more suggestive of the breadcrumbs hypothesis, though pre-caching increases with model scale.
1 Introduction
Humans are known think ahead while speaking; decades 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 be predicted to some extent using linear probes on model hidden states, and interventions on hidden states can predictably alter future outputs.'
These findings indicate that model activations a given timestep are least somewhat predictive of future outputs. However, it remains unclear why this might be: is this just 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.
Code to reproduce results be found at http:://g.thh.ocooo/////////tttt

---

Published as a conference paper at COLM 2024
& ("twelve"l""Alic has..")
Human
5 + == 12
Vo
Unembedding
Layer m +
Alice has five apples. Bob has seven apples. total,
Laver m
apples
Layer m
GPT
Sequence position t
Figure 1: At which position is the computation required 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 f that are irrelevant to the inference task at that current time step but may be useful for future time steps t T, and the breadcrumbs hypothesis, in which the features most relevant time step are already identical those that would most benefit inference at time / + To evaluate which hypothesis might be correct, we propose myopic training scheme that does not propagate gradients from the loss 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 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 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 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" way that small models cannot.
2

---

Published as a conference paper at COLM 2024
2
Related work
Future token meta-prediction. Several recent works (nostalgebraist, 2020; Belrose et al., 2023; Pal 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 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 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; P'mentel et al., 2020; Belinkov, 2021). Probing-basee 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 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 undersland 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 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.
3 Theory: Pre-caching or breadcrumbs?
Consider a generic causal sequence-to-seqeene prediction task
(xi, Xuy .,...., .,n) ~ D,
where is a data distribution supported on X" x Y for some domains X, Y. The task is to estimate the conditional expectations Ep (y; xI,..., for 1 S n.² Note that we recover the autoregressive setting by setting Y X and y; X;+1
Transformer models trained on such tasks have been observed (Pal et al., 2023) to store information in hidden states during inference at position that is then used in future inference at j > i. However, since the loss associated with each step depends only on how well the model does at the immediate task of predicting yi, 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 ex- pected 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, l "unintentionally" leaves trace ("breadcrumbs") that is then picked up by future passes.
2or classification tasks, we are typically interested in the conditional probabilities Prp c Xi,-.., xn) for each class c. However, this can be subsumed into the generic case by letting Y be the probability simplex over all classes.
3

---

Published as a conference paper at COLM 2024
To disentangle these two explanations, introduce 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 a quantitative measure of how much pre-caching is taking place.
3.1 Causal sequence modeling
Suppose, for the sake of exposition, that the transformer model G uses independent parame- ters for each position. Let be the parameter count of each forward pass of G. Then, letting 8 E R be all parameters used by G at position i, a transformer G is a parameterized function
G: X" x O"
Y",
(xi...., Xa; ,,....,0,1 + (f1....,'w')
For 1 sis H, let (X.,..... Xn) E be the output of G's ith forward pass. Because of the causal masking within G, this depends only on x1, X; and oj.... 0,. That is, with slight abuse of notation, we may write
9;
Xa. 0I............... ......, 0n) (x . x;; 0 ... 0:)
3.2 Off-diagonal gradient terms
Now, letting L: Y x Y + R+ be some choice of loss function, the expected loss & of transformer model with parameters oI,..., O, is
E( 0n) = (Z Σ [(G x 0 yi). =: Σ L (0 (0)
the sum over 1 < i S of the expected loss at position (We suppress the dependence on G and DD for concision.) In practice, we always tie the weights across position. That is, all 0j are set equal to the same O. Then, by the chain rule,
Vel(o, 0) = Σ
0n)
( 0;)lo Or
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.
3.3 Measuring pre-caching: The myopia gap
We say a model is myopic when each forward pass G, optimizes only L, without regard for future Ej jt > i. the untied weights case, the right definition then apparent.
Definition 1. The parameters (ØI...,, 0m) O are untied-myoptt if they satisfy
0, E arg min 00 0;)
Vi (1....,
(1)
6,
Definition 2. Let IM be the feasible set of the constraints in Equation The untied myopia gap the smallest possible gap between the expected loss attained by a myopic model and the optimal model:
min e( p" (M)
min 8.).....)))6
e(o 0,)
(2)
In the tied weights case, it is perhaps not immediately clear what the right definition of myopia should be. It does not suffice simply constrain the minimizations in Equation to θ,, since ming (,,...., 0) is optimizing for pre-caching (the dependence on arguments j < i) well as the present inference (the dependence on argument i). Instead,
For example, this true of absolute position embedding weights.

---

Published as a conference paper at COLM 2024
the right notion is 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 i E + are (tied-)myopic if they satisfy
11
i min 0.0
Σ Ei
(3)
The (tied) myopia gap is then defined analogously to Definition
The breadcrwmbs hypothesis states that the myopia gap is nmall-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 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 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 The (tied) local myopia bonus at 0 E e
E(O') max
(e
0*) ( 0,,,0"...,0000))
0)) 2 0.
For further interpretation of the myopia gap and myopia bonus, see Appendix A.
3.4 Myopic gradient descent
Our heuristic remark in Section 3.2, that the off-diagonal gradient terms responsible for pre-caching, is justified by Theorem 13 below. states that, given certain regularity conditions the loss terms l,, performing gradient descent with the off-diagonal terms removed results in myopic model in the sense of Definition 3. We call this myopic descent.
For myopic descent to be stable in the tied-weights case, 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-dppendenn mediated purely by the attention mechanism, while the present forward pass depends on both attention and foedforward parameters. The precise condition we use forward bias (Definition 11); see Appendix C for details and the proof of the theorem.
on the past is
Theorem 13. Let (0,0) Σ (,, ..., 0). If f is arward-biasedd o-strongly convex, and L-smooth, then, for some step size rj > 0, the iterates of myopic descent with tied weights
0,
@it+1) ( ef, 0=0=0())
converge to θ 2 satisfying the myopia constraints of Equation 3.
Synthetic data experiments
4.1 The Dp task
To demonstrate simple example where the transformer model learns pre-cache (and thus the myopia gap is large), we construct the following synthetic dataset.
Definition The data distribution DN p.a is defined as the joint distribution of real- nall-ollud random variables (xn (x N ( where, for each n,

---

Published as a conference paper at COLM 2024
Xy (0, (standard Gaussian)
Zn Ber(p) (Bernoulli with probability p)
Yn In Dizı sin(bx,-i) + (1- Zn)xn
and (Xx) neN U (Z IEEN are mutually independent. In our experiments, we always set the parame- ters a = b = 10 and N = 64, so for convenience notate Dp := -64 p.10,10
The intuition is that a transformer regression model trained on Dp would benefit from pre-caching sin (bx,) during its forward pass at position н, even though this computation is irrelevant to its task of predicting Yn One simple strategy that makes use of this pre-caching Algorithm
The motivation for the Bernoulli variables z is that, as decreases, the expected first time when sin(bx;) becomes useful advances further into the future. In addition, when p is sufficiently small, the probability (1 p) that the value sin (bxx) is never useful at all becomes non-negligible. We will show that, even this case, the transformer model learns
to pre-cache.
Investigating myopia. Suppose that we train myopic model (Section 3.4) on the same task. Since this model lacks off-diagonal gradient terms, we not expect to learn to pre-cache sin(bx,) at position n. One possible strategy that does not 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 single layer, while each layer of Algorithm 1 computes only scalar nonlinear functions."
Algorithm Pre-caching algorithm
Algorithm Brute force algorithm
At position
At position n, input Xn,7n layer 1 compute
input Xn. 7m layer compute :- sin(bxn))
layer read F for i 1,.... a.
layer 2 read Xリ-i for 1,, 1,...,a.
layer 2 compute
layer 2 compute
yn == Zu D Fu-i + (1 kn)
8n := Zn Ei-1 sin (bxnb) (1 Zn Xn
return 8n.
return 8..
4.1.1 Evaluation: linear probing
To determine if the transformer model is computing sin(bx,) at position we fit linear probes on the hidden states. We additionally compute the correlations between sin and each individual dimension (i.e., each neuron) of each hidden state. See Section D.1.1 for details.
*We think of each transformer layer read operation, performed by the attention mechanism, followed by compute operation, performed by the feedforward block.
Foor 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 (Xn, Zu : mutually independent, this follows from Lemma 15, stating that x, and (bn)),,) have near-zero correlation for large enough b. In our experiments, we set b 10, which case p(x,, sin (bxx))))) 10 10-20, In other words, there is low predictive y-information from the inputs to the target sin (bx,)), where is the class of linear models (Xu al., 2020).

---

Published as a conference paper at COLM 2024
4.1.2 Results for Dp
For varying p. we train two-layer transformer models with embedding dimensions of 128 on Dp 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, i) for varying i, we find strong evidence that the transformer model with vanilla training is indeed pre-caching sin(bx, ((nb),)) ) 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 is correlated with only Xx-
The first hidden state is correlated with sin(bx,) but not correlated with any sin(bx,-i) for i > 0.
The second hidden state (immediately before the output unembedding) is correlated with -i) for each
Further, looking at the per-neuron correla- tions in Figure 3, we see that sin sin(bx,-i) for 1 S I a are all correlated with sin- gle 1-d subspace of the second hidden state (they share the same striping pattern); this is the subspace storing Ei= sin sin(bx (bx Mean- while, sin (bx,), as well as many of the X 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.
p
Vanilla
Myopic
0.01
0.096
1.10
0.1
0.016
0.97
0.3
0.0030
1.03
1.0
0.0074
1.26
Normalized Huber loss C/p for vanilla and myopic models trained and eval- uated on Dp for each p in our synthetic set- ting. For reference, the trivial model that always outputs zero attains Huber loss of 1.26.
Table 1:
On the other hand, in Table 1, the myopic models perform significantly worse. The per- neuron correlations in Figure suggest that the myopic model may be implementing a crude approximation of Algorithm 2. This suggests that the synthetic setting has an in- herently high myopia gap-it it impossible for the transformer model to do well without pre-caching.
4.2 Multiplication
In addition to the above Dp synthetic data setting, we also measure the myopia gap 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.
5
Natural language experiments
5.1 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 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
7

---

Published as a conference paper at COLM 2024
between hidden state sin(10X;-)
p=0.01
Ryer lyyer
0.8
0.4
14
pw 0.3
p
0.8
0.8
0.6
0.2
0.2
12
14
10
14
Figure Empirical R² between linear probes fit each layer of vanilla transformer models trained on Dp for 0.1,0,0.1033,1111 to targets sin(bx, i). Computed over 50 000 samples from D1.
vanilla model during both training and evaluation.? See Appendix for implementation details.
As baseline, we also train "transformer bigram" model, model with identical architec- ture but all off-diagonal key / /value states zeroed out.
5.1.1 GPT-2 results
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 3.28 3.26 0.02.
The nonzero myopia gap suggests that pre-caching may provide small positive benefit. Indeed, in Figure 4, we see that the myopic model outperforms the vanilla model the beginning of the sequence, since it can allocate all compute to next-token prediction, but quickly falls behind the length of the past increases, since suffers from lack of pre- cached information from earlier forward passes. 8
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 in- formation (breadcrumbs) to significant extent, even they optimized only for the present in- ference 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.
Model
Cross-entropy
Vanilla
3.28
Myopic
3.40
Local myopic 3.26
Transformer bigram
5.33
Table 2: Validation cross-entropy loss ob-
tained by GPT-2 with vanilla and myopic training
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 separate model.
WWe use a sliding window over PG-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 affect the relative comparison of myopic and vanilla GPT-2.
8

---

Published as a conference paper at COLM 2024
pllayer 0,Xo Xa-i)
plaayer O. sin(10x,-i)
plaayer O, X--i)
pllayer G. sini sin1100
0-
0-
lookback Index
lookback Index
5
10
10
10-
10-
15
15
13
15
50
neuron index
neuron index
neuren index
neuron index
pllayer 1, Ko
pllayer 1. sin(10x .1]
p(layer 1. X, .
pllayer 1, sin[10x,
lookback Indox
lookback Index
5-
5-
10
10
10
10
15
15
13
30
neuron Index
neuron index
neuron Index
neuron Index
pllayer 2, Ko-.)
ptlayer 2, sin/10x, |1
p(layer Xx
pilayer 2, sin(10x, )1
looksace Index
lookback Indext
20
10
10
20
25
25
125
25
50
25
50
75
neuron index
neuron index
neuron index
Vanilla model
Myopic model
Figure 3: Empirical correlations between each hidden state neuron and X;-i sin (bxxn.i). Models are vanilla (left two columns) and myopic (right two columns) transformers trained Do.3.
Loss by position
Myopla gap by position
myoppi
dhandart
6.00
3.75
mrpppia gep
5,50
5,25
0.05
-0.10
position
posibio
Figure 4: Cross-entropy loss of vanilla and myopic GPT-2 models by token position, and their difference. Evaluated on sliding window over 100K-token sample text from the PG-19 dataset (Rae al., 2019). Aggregate cross-entropy losses on this sample are 4.67 (vanilla) and 4.77 (myopic).
5.2 Myopia gap scaling
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 more pronounced myopia gap.
To test this, we train both vanilla and myopic transformers from the Pythia LLM suite (Bi- derman et al., 2023), ranging in size from 14M to 2.8B parameters, one epoch of 10M

---

Published as a conference paper at COLM 2024
0.-
077
0.70
044
035
10P
10*
301
parameeeers
LAMBADA (Paperno et al., 2016)
PIQA (Bisk et al., 2020)
030
0.55
myppi
0.50
0.55
0.5
0.0
0.0
0.40
0.75
033
0.50
030
10a
20*
paramates
SciQ (Welbl et al., 2017)
ARC-Easy (Zhang et al., 2018)
Figure 5: Benchmarks of Pythia models fine-tuned on the Pile dateset using vanilla and myopic descent.
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 ments (Figure 5). Note that, unlike in the GPT-2 experiments (Section 5.1), which start from random initialization, points provided by Biderman et al. (2023)- for the larger architectures, the 10M sequence dataset use is not sufficiently large to for pre-training from random initialization.
a variety of natural language experi-
start all training for Pythia models from the pre-trained check-
Discussion and future work
Using a synthetic dataset, we demonstrate that pre-caching can indeed be learned by 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 terms of safety and/ /nd/or interpretability-ii may be easier to understand what a model is doing if we know that everything that is computing on each forward pass 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 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 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.
10

---

Published as a conference paper at COLM 2024
Acknowledgments
LL thanks Lukas Berglund and David Schneider-Joseph for inspiring conversations. This research was partly supported by Open Philanthropy and the Berkeley Existential Risk Initiative.
References
Maik Barthel, Sebastian Sauppe, Stephen C Levinson, and Antje Meyer. The timing of utterance planning in task-oriented dialogue: Evidence from novel list-completion paradigm. Frontiers in psychology, 7:1858, December 2016. doi: 10.3389 ftts//psy.001 2er..2166.. 2016.01856
David Bau, Jun-Yan Zhu, Hendrik Strobelt, Agata Lapedriza, Bolei Zhou, and Antonio Torralba. Understanding the role of individual units in deep neural network. Proceedings of the National Academy of Sciences, 117749):30771-3007 September 2020. ISSN 1091-6490.
doi:
URL http://dx..oi..rr// 1073/ pnas//pnsss....... 1987375117.
Amir Beck. First-Order Methods in Optimization. Society for Industrial and Applied
Mathematics, Philadelphia, PA, 2017. doi: 01.1137/1978161117499999
URL https:
//epubs. siam. org/ /tt.o///ab//.... 11771 977711119111199977777
Yonatan Belinkov. Probing classifiers: Promises, shortcomings, and advances, 2021.
Yonatan Belinkov and James Glass. Analysis methods in neural language processing: A survey, 2019.
Nora Belrose, Zach Furman, Logan Smith, Danny Halawi, Igor Ostrovsky, Lev McKinney, Stella Biderman, and Jacob Steinhardt. Eliciting latent predictions from transformers with the tuned lens, 2023.
Stella Biderman, Hailey Schoelkopf, Quentin Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar Van Der Wal. Pythia: suite for analyzing large language models across training and scaling. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. Proceedings of the AAAI Conference on Artificial Intelligence, 34(05):7432-7439 Apr. 2020. doi: 10.1609/ v34i.53400. URL https://ojs. aaai wtt.aa/.rrr....e...//////a/aaaaa
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple Ilm inference acceleration framework with multiple decoding heads, 2024.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling, ar arXiv preprint arXiv:2101.0002) 2020.
Evan Hernandez, Arnb Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau. Linearity of relation decoding in transformer language models, 2024.
John Hewitt and Percy Liang. Designing and interpreting probes with control tasks, 2019.
Falk Huettig. Four central questions about prediction in language processing. Brain Research, 1626:118-135, 2015. ISSN 0006-8993. doi: https:/ http:///doooor/////// 20l.://0...000010/.aaaare......0 URL https:/ //ts:///ww... sciencedirect. wttpeedettttttccccceeeeeeeeeeeee http://errrrreepp///////00033310 146. Pre- dictive and Attentive Processing in Perception and Action.
com/science/ http:/aaticccppeeeeee
11

---

Published as a conference paper at COLM 2024
Kenneth Li, Aspen K. Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Emergent world representations: Exploring a sequence model trained on synthetic task, 2023.
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations gpt, 2023.
George A. Miller. Language and communication. McGraw-Hill, New York, NY, US, 1951. doi: 110.003711133-000.
Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023.
Yurii Nesterov. Lectures on Corvex Optimization. Springer Publishing Company, Incorporated, 2nd edition, 2018. ISBN 3319915770..
Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. MS MARCO: A human generated machine reading comprehension dataset. In Tarek Richard Besold, Antoine Bordes, Artur S, d' Avila Garcez, and Greg Wayne (eds.), Proceedings of the Workshop Cognitive Computation: Integrating neural and symbolic approaches 2016 co-located with the 30th Annual Conference on Neural Information Processing Systems (NIPS 2016), Barcelona, Spain, December 9, 2016, volume 1773 of CEUR Workshop Proceedings. CEUR-WS.org, 2016. URL https: http://///euurwwoooo////// wttg//oll1773/COoo paper pdf.
2016
nostalgebraist. Interpreting gpt: The logit lens, 2020.
Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter. Zoom in: An introduction to circuits. Distill, 2020. doi: 10.23915/ 0002..0.00024.001... https:/ http:://dssiill.p//22/00 2tt//////////////////////a strcuis/zzoomi/
Koyena Pal, Jiuding Sun, Andrew Yuan, Byron Wallace, and David Bau. Future lens: Antici- pating subsequent tokens from a single hidden state. In Proceedings of the 27th Conference on Computational Natural Language Learning (CoNLL). Association for Computational Lin- guistics, 2023. doi: w..11653/w1/2223..2 conlll............... 2023. c023.con11-1.37
URL http://xx doi org/ httgg/....1133///1 18653/v1/
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The LAMBADA dataset: Word prediction requiring broad discourse context. In Katrin Erk and Noah A. Smith (eds.), Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp- 1525-1534, Berlin, Germany, August 2016. Association for Computational Linguistics. doi: 10.18653/v1/P16111 URL BuRLhttps://acaanhholloo org/P16-1144.
Jacob Pfau, William Merrill, and Samuel R. Bowman. Let's think dot by dot: Hidden computation in transformer language models, 2024.
Tiago Pimentel, Josef Valvoda, Rowan Hall Maudslay, Ran Zmigrod, Adina Williams, and Ryan Cotterell. Information-theerrtiic probing for linguistic structure, 2020.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and llya Sutskever. Language models are unsupervised multitask learners. 2019. URL h
semaanticccholl....rrr//oooooopo 160025533.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, Chloe Hillier, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arviv preprini, 2019. URL
Ruoqi Shen, Sébastien Bubeck, Ronen Eldan, Yin Tat Lee, Yuanzhi Li, and Yi Zhang. Posi- tional description matters for transformers arithmetic, 2023.
12

---

Published as a conference paper at COLM 2024
Zuowei Shen, Haizhao Yang, and Shijun Zhang. Optimal approximation rate of relu networks in terms of width and depth. Journal de Mathématiques Pures et Appliquées, 157: 101-135, 2022. ISSN 0021-7824. doi: https://doi..rr http:///t.ooog/000./////aaaaaaap
2021.tp000..22...00000.0...... URL http:///wwe.....ec.c.c.cceaaaaaa
Xing Shi, Inkit Padhi, and Kevin Knight. Does string-based neural mt learn source syntax? Pp. 1526-1534, 01 2016. doi: 10.18653/v1/D16--155/
Mitchell Stern, Noam M. Shazeer, and Jakob Uszkoreit. Blockwise parallel decoding for deep autoregressive models. In Neural Information Processing Systems, 2018. URL https:
seampniisannccccholl wwwp/oorpgssD.......ooo 53:30333.... 53208380.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, L ukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https://proceedins..eerr httpr//ggg..rrisscc///ppp/eee//e ff5ee2255444eeffffo0sscccccaaaa3
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice sci- ence questions. ArXia, abs/ babs/1707.062099 2017. URL / semanticscholaaa. .rrg/
CorpusID: 1553193.
Yilun Xu, Shengjia Zhao, Jiaming Song, Russell Stewart, and Stefano Ermon. A theory of usable information under computational constraints, 2020.
Yuyu Zhang, Hanjun Dai, Kamil Toraman, and Le Song. kg²: Learning reason science exam questions with contextual knowledge graph embeddings. ArXiv, abs 18055 2018. URL https: / /api.. semanticscholar wttpricccoooooooorrroooooooeo [ 44098100.
Ziqian Zhong, Ziming Liu, Max Tegmark, and Jacob Andreas. The clock and the pizza: Two stories in mechanistic explanation of neural networks, 2023.
A Myopia bonus and malus
Notice that the myopia gap consists of two pieces: myopia bonus, the improvement that can be obtained at the current forward pass by ignoring the future forward passes; and 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," given choice of myopic o,,..., 0 satisfying constraints (1) of Definition 1, write
E(O 0x)
min ((0)...., e
Σ
min {(ō 0 0, 0,000 0,,....,,, 0n) min 00.1,,,,,,,.,,,, 0 0.....0r
.0)
0++1...,0a On
0 (4,(0 0 _ 0,
Σ (e,(0,....
...,0, ( ..., 0 - 10,0,..,,,,,,.0 .7,....,
-1j
where we define
07 ...,, 0 E arg aig ((0.....,0011,,,,,00 ......
That is, the myopia gap is the sum d E d; of the myopia bonuses
c: (..,,,,, 0m) = ( 0 0,) , 0,¹) < 0, 9hhe tied is analogous, do not write it explicitly.
13

---

Published as a conference paper at COLM 2024
(the inequality following from the myopia constraints (1)), and the myopia maluses
d;(0),...,,,,
Σ (,(0,....,,,,,,0,00 ...,0.0. 00 0 0 0 1++1
2
with d; + 2 0 for each 1 i n by the definition of the 0 A priori, a small myopia gap does not necessarily imply small (in magnitude) myopia bonus and malus d. Indeed, in the case when the myopia gap is small, large value for (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, is possible that both c and d are small; that is, there 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 (()) separate optímization problems (each of which, practice, a full transformer model training run). Thus, we instead compute the local myopia bonus of Definition 4.
A.1 Explicit gradient paths
The dependence of forward pass G on previous forward passes Gj for j < i is mediated through hidden states hy...., h;
(X x, 0 0;) = (h,,,..., 1, 1x
where the h; are themselves recursively defined parameterized functions h; hich.....,l x;; 0,), (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ú is attended to by forward pass >> j, we consider the key and value weights Wx and Wy, respectively, to belong to forward pass i. Thus, they are updated by the gradient wrt. fz.))
With the hidden states explicitly written out, the gradient wrt the loss is a sum over all possible paths to the present: for j <
aG; 30;
ahj ac, m-1 ah;
[x1
X (Q.......) = Σ
00, him
dhi-1
jui
k=
where the sum is over all partitions j j= = <... < im <.
B The myopic attention mechanism
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 / vallue 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 key/value states are the result of key /value weights Wx and Wy, respectively, multiplying cloned and detached copy of the previous hidden states.
During both training and inference of a local myopic model, the past key value states come from a separate frozen pre-trained transformer model.
During both training and inference of a transformer bigram model, the past key value states are simply zeroed out.
14

---

Published as a conference paper at COLM 2024
Let X Xx) E Roxd be the sequence of residual stream hidden states per position, with each row representing one position's hidden state in R Let Wo, Wg. Wy Rdxh be the query, key, value weight matrices for one attention head, of dimensionality h, in the transformer G. Denote
Q = XWQ : XWW, V = XWV
and let Q K, V 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. 4i is the ith row of Q. For simplicity of presentation, we omit causal masking; the modifications that should made in the presence of a mask are straightforward.
Recall that the vanilla attention mechanism for G is
== cokk((QK)) )),
where o is row-wise softmax. Writing this out token-wise,
yi Z exp(q) (q: kj)v,,
where Z is the partition function
2, = Σ exp kj)
The myopic attention mechanism, on the other hand, written tokenwise as
9: 2,- exp(q) k;)v; + Exxp(1 k,)v,)
2,
exp(q) kj + pوfr (kj k;)) + + 7,))
11
Σ a: Vj
Σ -=1
ōijaj(y; - -v,)
where
7, = exp(q) k;) Eep((q.1. k))
iti
Σ exp(1) T kj + 8ij97 (kj kj)),
ajj Z exp(q, kj + &jjq. (kj k,)),
and dij is the Kronecker delta. Now, notate
(diag A)i; == Aij
That is, diag A is the diagonal matrix that has the same entries A along the diagonal and zero elsewhere. We are now able to write the myopic attention mechanism in matrix form:
AV + (diag A) - VV,,
where
== (QK + diag( Q((K"QQ KT KT))).
C Proofs
We use two assumptions on the loss function to be minimized. These are standard the first-order methods literature.
15

---

Published as a conference paper at COLM 2024
Definition 6. A function f: R" + R is called L-smooth if it is continuously differentiable with L-Lipschitz gradient. That is, for all x, y in the domain,
Vf(y)ll2 S L||x - yll2. [
Definition 7. A function f: R² + IR is called o-strongly convex if, for all x, y in the domain,
f(y) 2 f(x) + Vf(x) (y - x) 2 lly - x112.
In particular, recall that strong convexity implies the existence of unique minimum. Hence, makes sense to write, for example, x² arg max, fx) without ambiguity.
C.1 Gradient descent with untied weights
Theorem Assume E: @" - IR is r-strongly convex and L-smooth for some ar, L >> O. Consider ordinary gradient descent with unlied weights
00 0f 1) 1 e(0 1) V,..., ") Vi nI.
Then, for 0
0% = arg min m Br..,, e fn), small enough > 0,
110|) 11 v (1 201l
|10(0) -10!0/-0111 vi nI.
o +
Proof. This is standard convergence result for gradient descent on strongly convex func- tions. For example, see Nesterov (2018).
C.2 Gradient descent with tied weights
Theorem 9. Assume e is o-strongly convex and L-smooth, and consider ordinary gradient descent with fied weights
0
0
0 (f11) 1)
+ Σ V ((00 E(0,,...,,0) 0n)
Vi nI.
There exists 1j > such that
1000 1011 < (1
2 2/nnoL.)
11000)
10
Vi (I.....n).
++
where 0* m E(),...,,)) 0).
Proof. This is again the standard convergence result, now applied to descent with step size /nn on the o-strongly convex L-smooth function 0 /n./vn,.... /n). Alternatively, one may think of this as projected gradient descent constrained to the subspace 01 0, with step size of /nn. Projected gradient descent inherits the same convergence properties as unconstrained gradient descent (Beck, 2017).
C.3 Myopic descent with untied weights
Theorem 10. Assume each of the e,...., en o-strongly convex and L-smooth. Consider myopic gradient descent with untied weights
0/0 0((-1) -n V. (0 (--1
((- vi n..
16

---

Published as a conference paper at COLM 2024
There exists rj > such that 0 + 0, for all i, where
0 = arg min (01)
8
0 = arg min (0, 02)
0, arg min en 02. ,, On 1. 0,),
0n
Proof. Let .j...., 0, be as the theorem statement. We proceed by induction. For the base case, note that the myopic descent iterates for 8 are independent of f for Thus, the standard convergence theorem gives that 0 0 as
Now, assume o( 1 Ú for all k. Thus, for any 0, for sufficiently large f,
(1 | (( (0.) <E (I(e))
Hence, since l is L-smooth, for any Oz.
V (0 ,..., (()) 0x) 02) Ve, (0),..... 0g 1, 1,0x)1l < Le.
Expanding and rearranging,
0x), (Ve (0 0b), V 4 (O,...., 0g 1. 0,)) 2 2 [ 0, (0,,... ..... 02 1. 0x)111
1
1 L222
is bounded away from zero as long as, say,
V ( 02-1, 02) 2 0 02112 > (L+ 1)e,
(I. + 1)e it is guar-
using the o-strong convexity of 4g That is, as long as ||0 - 0.l12 >
anteed that Ve, (( .,.., 1- 0x) is descent direction for (0),
0, 1 0x). This is
a sufficient condition for 0g to converge to a (L+1)e neighborhood of 0g given small arbitrary, this completes the inductive step.
o
enough step size (Beck, 2017). Since
C.4 Myopic descent with tied weights
Definition 11. A function f(x y): Rd x R IR with continuous second derioatives is p- forward-biased if, for all y E R" and x = y,
Hynf (x, Hxyf (x, (x, y) x
and
x x(Hyyf f (x, y) + Hxyy y)) < p, where x is the condition number, and we write the Hessian of f as a block matrix:
df(x,y) dx,āx;
df(x, y) ax,āy;
Hxxf(x,y) Hxyf (x,
Isisa 15i5a
1sisa
Hf(x,y)
13i5h
H (x,y)
(x,
af(x, y) dy,Ox;
a dy,ōy;
Isisb 15j5a
1sish
Lemma 12, A E Raaxa is positive definite with x(A) < p, then there exists some fj > 0 such that 1- 1- is (1 contrracittt... Explicitly, for any y IRa,
((T- 1Aylla (1 ----->>)lll llvll2
17

---

Published as a conference paper at COLM 2024
Proof. Set 11 1/ 1/Amaxx(44 (A). Then I 7/A 2 0 with
A Amn (A)
Amax( 1/A) 1/A (A) =
< <1-p-1
Amax
It immediately follows that I A/ is (1 p (1--p()11-cotapaaaaaaa --ontrracci.
Theorem 13. Let y) Rd x IR" + R be p-forward biased, o-strongly convex, and L-smooth with continuous second derivatives for some 07, L > Then, there exists y E RR" such that
y arg min f (g.y).
yeR*
Further, for some step size 1j > O, the iterates of myopic gradient descent with tied weights
((++1) = y() V yf (x, y) x=y=y/()
converge to y E R satisfying (4). We call such y a myopic solution.
We then recover the original sequential modeling setting by defining f(θ,0) == Σ (,, 0,0).
First, note that if myopic descent does converge, must be point y such that arg min, fg. y). Indeed, at convergence we must have Vyf(y, y- Ó, so strong convexity tells us that ÿ is optimal. Thus, if we establish that myopic descent with small enough step size > converges to some y, we automatically get the existence of a myopic solution. For this, it suffices to show that the gradient descent step
Proof.
87 (y) y 7/ V (x, y)xx=y strictly contractive, so that iterates of g converge to a fixed point.
Consider arbitrary y and y'. Then, by chain rule and the mean value theorem applied to the map y' + (xf(x,y)xx y-y''
Vyf(y', v') (y.(y = (Hxyfly", ,""""")) + (f))((""v"")) - y)
(5)
for some y" y', Using the definition of gy and substituting in (5),
(18)(y') 8n (y) (v) 12 l1(y' y) 7 (V vf uy. ,y) Vyf v))
= =II(1 (Hxyf (y"
(1"""""" .y") + "y" (3) y)
(1-p ()' -y), y),
where the last line is by p-forward bias and Lemma 12. That is, 8y is (1 p- p-))-contracttve, completing the proof.
C.5
Properties of sine
Lemma 14. Let x (0, Then
Var(sin(bx)) 1 e262
2
Proof.
2262
Var(sin(bx)) (2Tt) 1/2 I sinZ(bx)e x22 dx
1
2
Lemma 15. Let ~ (O, Then
2be362 /2
p(x, sin(())
E Ole O(e-b12),
6262
where p is the Pearson correlation coefficient.
18

---

Published as a conference paper at COLM 2024
Proof. By symmetry, [[sin[bx)] 0. We calculate
Cov(x, sin) (2/2)-1/2 [ x sin(bb)) x2/2 dx be -b212
1/2
We already computed the variance in Lemma 14, so
s(nx)))( = (in(x)sin(bx)) 2be
Vaa(xx)aa((sin(b)) e26²
In our experiments we set b = 10, p(x, sin(b)) < 10 20
D Details and additional experiments
D.1 Synthetic setting: Dp task
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 trained linear map from the input space to the embedding space and from the embedding space to the output space, respectively. For each p E {0.01, 0.3, 1} models are trained using ordinary and myopic descent on one epoch of 30M sequences of length 64 sampled from D64 architecture details.
p,10,10 See Table for
Configuration Key
Value
num_layers
2
num heads
2
embd dim
128
n innner inputddmm
512
2
output. dim
1
activation
relu
attn -dropp
0
embd embd.pdrop
0
resid pdrop Ir
0
1e-3
optimizer
AdamW
weight.decay
0.01
betas
scheduler
cosine
warmup
0.01
batch size
512
seq length
64
loss_fn
HuberLoss
Table 3: Transformer configuration used when training on synthetic data distribution Dp
D.1.1 Probe details
Given transformer model trained on Dp, we sample the hidden states at each layer when the model is given input 50000 evaluation sequences from the same distribution Dp. For each layer and targets sin(bx,-:) for varying i > 0, we fit linear regression model on the hidden state of that layer to the target. The in-sample R² of each linear model is then reported in Figure/ 2. Figure a visualization of the linear probe's performance on the vanilla transformer.
19

---

Published as a conference paper at COLM 2024
Layer 1 probe estimatee sin[10X,)
1.0
0.-
peobu
-1.5
Figure 6: Estimate of sin(bxn) by linear probe fit on layer of transformer with vanilla training Da.3. Computed over 50 000 samples from D1.
D.2 Natural language setting
D.2.1
GPT-2
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, 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 baseline, we also train "transformer bigram" model, which is transformer model whose key / value states are zeroed out during training and evaluation. See Table 4 for configuration details.
Configuration Key
Value
num layers
12
num heads
12
embd dim
768
n_inner vocab_size
3072
50257
activation
gelu.new
attn -drop
0.1
embd pdrop
0.1
resid _drop Ir
0.1
6.0 x 10-4
optimizer
AdamW
weight decay
0.01
betas
(0.9,0.999) 0.999)
scheduler
cosine
warmup
0.01
batch_size
512
seq length
64
loss fn
CrossEntropy
Table 4: GPT-2 small configuration used when training on natural language data.
20

---

Published as a conference paper at COLM 2024
D.2.2 Pythia
For experiments on the Pythia suite, we finetune with either vanilla or myopic descent on 10M sequences of 64 tokens each subsampled from The Pile (Gao et al., 2020). Learning rates and batch sizes for each model are presented Table 5; they are the same between vanilla and myopic descent. All other training and architectural hyperparameters are the same as those used by Biderman et al. (2023).
Model
Learning rate
Batch size
Pythia 14M
4.0 x 10
512
Pythia 31M
4.0 x 10-1
512
Pythia 70M
1.2 10
512
Pythia 160M
1.2 10
512
Pythia 410M
1.2 10 4
256
Pythia 1B
1.2 10
128
Pythia 1.4B
1,2 x 10
128
Pythia 2.8B
8.0 x 10 5
64
Table 5: Pythia suite hyperparameters for finetuning, Batch size is measured in sequences of 64 tokens each.
4,25
vanilla myopic
4,00
3.75
cross-entropy loss
3.25
3.00
2.50
10a
109
parameters
Figure 7: Cross-entropy loss of Pythia models fine-tuned on The Pile dataset using vanilla and myopic gradient descent. Starting from the 70M model, we see that the gap increases with parameter count.
D.3 Multiplication
In addition to the natural language experiments, also measure the myopia gap on the task of natural number multiplication. We use the same GPT2-small architecture as the natural language experiments starting from random initialization; Table 4. See Figure for an example input sequence.
7
0
0
5 5
Figure 8: Example multiplication input sequence.
21

---

Published as a conference paper at COLM 2024
2
4
5
6
7
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
0.99
1.00
1.00
1.00
1.00
0.98
0.92
0.21
0.10
1.00
1.00
0.99
0.99
0.59
0.16
0.02
0.02
1.00
1.00
1.00
0.96
0.16
0.03
0.00
0.00
8
1.00
1.00
0.99
0.51
0.06
0.01
0.00
0.00
Vanilla multiplication accuracy
2
5
7
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
0.70
0.50
0,34
0.26
1.00
1.00
1.00
0.99
0.45
0.16
0.07
0.03
1.00
1.00
1.00
0.96
0.27
0.05
0.01
0.00
1.00
1.00
1.00
0.60
0.11
0.02
0.00
0.00
8
1.00
1.00
1.00
0.41
0.05
0.00
0.00
0.00
Myopic multiplication accuracy
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
2
0.00
0.00
0.00
0.00
0.00
0.00
0.00
0.00
3
0.00
0.00
0.00
0.00
0.00
0.00
0.00
-0.00
4
0.00
0.00
0.00
0.00
0.30
0.50
0.65
0.73
0.00
0.00
0.00
0.01
0.53
0.77
0.13
0.07
0.00
0.00
0.01
0.03
0.32
0.11
0.01
0.02
0.00
0.00
0.00
0.36
0.05
0.01
-0.00
0.00
8 0.00
0.00
0.01 0.10
0.01
0.01
0.00
0.00
Vanilla-myopic accuracy gap
Figure 9: Accuracy of vanilla and myopic transformers trained on multiplication of up to 8-digit inputs. Row and columns correspond to the number of digits the first and second multiplicands, respectively.
We use several formatting tricks found by Shen et al. (2023) to improve performance on the multiplication task:
Characters are delimited by spaces, such that each digit is tokenized into separate token.
All numbers are written in the reverse of the standard order, i.e. such that the least significant digits come first.
All inputs are zero-padded to the same length.
Note for each multiplicand we first sample the number of digits d ~ Unif(n) uniformly in some range n, then uniformly sample natural number x ~ Unif(10 1) with no more than d digits. This distribution allocates increased weight to smaller numbers, and was found to result in superior performance.
We train both vanilla and myopic transformers on one epoch of 10M examples with no more than digits, then measure 0/1 accuracy (that is, the model is provided with the an input sample up to the '=' token, and scored if it completes the rest of the sequence exactly correctly and 0 otherwise) on 1024 independent random validation examples for each of the 8 x 8 possible pairs of digit counts for the two multiplicands. See Figure 9.
22

---

Published as a conference paper at COLM 2024
D.3.1
Filler tokens
We further hypothesize that, as in Pfau et al. (2024), the vanilla transformer may learn to perform computation even on forward passes corresponding to filler tokens, thus attaining better performance when trained on examples zero-padded to longer lengths. We expect that myopic transformers, on the other hand, are not incentivized to do this, since this extra computation holds relevance towards the immediate task of predicting the filler zero token. To test this hypothesis, we train vanilla and myopic transformers on each of two different multiplication datasets:
1. Both multiplicands have at most 5 digits, and are zero-padded to exactly 5 digits.
2. Both multiplicands have at most 5 digits, and are zero-padded to exactly 10 digits.
Again, all training runs consist of one epoch of 10M examples.
See Figure 10 for results. Note that the vanilla transformer indeed performs better when trained and evaluated on input sequences zero-padded to a longer length. However, the myopic transformer performs substantially worse with increased zero-padding, Our explanation is that, not only does the myopic transformer not learn to perform intermediate tokens during zero-token forward passes, the increased input length makes more difficult for the attention mechanism to correctly attend to the relevant tokens.
D.4 Gradient angles
Using the publicly available training checkpoints for Pythia-410M (Biderman et al., 2023), we measure the sizes of both the myopic component and the future component of the gradient of the loss w/r/t the parameters over the course of training. (Note that the future component is the difference between the total vanilla gradient and the myopic gradient.) We also measure the cosine similarity between the myopic and future components. See Figure 11. One observation is that the norm of the myopic gradient is consistently larger than that of the future gradient- thus, training is dominated by the effect of each forward pass's 's parameters on the immediate next-token prediction.
1
2
3
4
5
1
2
3
4
5
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
3
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
0.40
1.00
1.00
1.00
1.00
0.99
Vanilla accuracy, padded to length
Vanilla accuracy, padded to length 10
1
2
3
4
5
1
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
2
1.00
1.00
1.00
1.00
1.00
2
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1.00
1,00
0,99
1.00
1.00
1.00
0.98
0.79
1.00
1.00
0.99
0.43
0.24
1.00
1.00
1.00
0.82
0.39
1.00
1.00
0.97
0.22
0.05
Myopic accuracy, padded to length 5
Myopic accuracy, padded to length 10
Figure 10: Multiplication accuracy of GPT-2 with either vanilla or myopic training and with input multiplicands zero-padded to either length or 10. Row and columns correspond to the number of digits in the first and second multiplicands, respectively. Observe that padding improves performance the vanilla model, but decreases performance of the myopic model.
23

---

Published as a conference paper at COLM 2024
myopic grad norm
Future
300
mmp
250
150
100
pimpppic
-04
-0.5
101
10'
10a
20
training steps
Norms of myopic, future, and total gradients
Cosine similarity between myopic and future gradient
Figure 11: Myopic and future gradients of PytMia-410M during training.
24


Reviews:
[{"review_id": "FNsIGTo1Xo", "reviewer": "Reviewer_URPJ", "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.\n\nThe 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.\nThe 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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. \n\n> Maybe could one come up with ways of training \"focally\" instead of \"myopically\"\n\nIt is in principle possible to train with up-weighted off-diagonal gradient terms, thus resulting in a “hyperopic” model. We leave this to future work.\n\n> I would like to see a closer discussion of NLP topics.\n\n> Can you measure this in ways that are different from cross-entropy?\n\nThe final revision will contain more experiments on a range of models up to 2.8B params in size, evaluated on standard NLP benchmarks\n\n> it would help to highlight that of course next tokens can be predicted from state\n\nWe cite Pal et al. \"Future Lens\", which investigates this claim empirically.\n\n> I presume that the cited works take precautions that future-token prediction doesn't simply stem from the fact that a next-to-next token will clearly depend on the next token\n\nThe cited Pal et al. paper accounts for this by measuring the effect of causal interventions on the hidden state.\n\n> Figure 4 is not referenced from the text.\n\nWe will fix this. Fig 4 demonstrates the linear probe’s ability to detect $\\sin(bx_n)$ in the hidden state.\n\n> When you refer to Table 4, it would help to remark that it can be found in Appendix F.\n\nWe will do this.\n\n> There's an Appendix G which is not referenced from the main text.\n\nWe agree that Appendix G would make more sense as part of a separate work. It will be removed in the final revision.\n\n> However, is there reason to believe that boosting this behaviour (beyond what happens in your standard transformer) would lead to better language modelling results?\n\nWe do not claim that boosting pre-caching would benefit performance. However, this is an interesting question to investigate in future work.\n\n> The implications of this research, in particular whether a model might be pre-caching or breadcrumbing remain open.\n\nWe agree that much future work can be done investigating the practical implications of pre-caching vs breadcrumbing. One possibility is that of automatically swapping in a myopic model 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).\n\nAlso, if indeed the myopia gap remains acceptably small for production-sized models, there may be some efficiency gains to be had by training myopically, and thus not computing the off-diagonal gradient terms. \nWe leave these directions to future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "valid analysis of what a transformer does internally.\ndefinition and thorough discussion of myopic training that could be useful in future analyses", "reasons_to_reject": "The implications of this research, in particular whether a model might be pre-caching or breadcrumbing remain open. \nThe 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?\n\nThe 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_to_authors": "See above. Some remarks:\n- Figure 4 is not referenced from the text. (It an also probably be omitted.) I'd put Table 1 earlier, into Section 4.\n- When you refer to Table 4, it would help to remark that it can be found in Appendix F.\n- There's an Appendix G which is not referenced from the main text. What does it do?\n- 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?", "extra_scores": {}}, {"review_id": "zMn1D5AeCO", "reviewer": "Reviewer_4NSr", "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).\n\nIt 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.\n\nThe 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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> the actual model and training data are modest.\n\nThe final revision will contain further myopia gap experiments on models ranging up to 2.8B parameters in size.\n\n> It is unclear what constitutes a significant amount of myopia gap\n\nThe myopia gap of $\\sim 0.12$ is best understood in the context of relative comparisons. Eg this gap is much smaller than that between the standard transformer and the “transformer bigram” referenced in the text. The aforementioned new experiments will also give more points of comparison between models of different scale. We will also include results on NLP benchmarks, which may more meaningfully translate to performance on real-world tasks\n\n> One may consider discussing work on blockwise parallel /non-autoregressive decoding\n\nWe will add the suggested citation. Note that we also cite Cai et al. \"Medusa\". We will expand our discussion of these related works in the final revision\n\n> Why use PG-19 for evaluation in Figure 5\n\nThe PG-19 dataset includes longer texts, and thus allows us to evaluate using a rolling window over a single text. This reduces noise in the per-position loss plot, as every token is used in every position. The domain shift between MS MARCO and PG-19 does not affect our findings---what we care about is the relative comparison between the vanilla and myopic models, both evaluated on PG-19\n\n> What are the practical implications of whether the model is employing pre-caching vs. breadcrumbs?\n\nOne possibility is that of automatically swapping in a myopic model 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).\nAlso, if indeed the myopia gap remains acceptably small for production-sized models, there may be some efficiency gains to be had be training myopically, and thus not computing the off-diagonal gradient terms.\nWe leave these directions to future work\n\n> Why does the myopically-trained model incur lower loss on the first few token positions than the standard LM (Figure 5)?\n\nIntuitively, the myopic model can dedicate all of its compute on the immediate next token. This advantage outweighs the cost of not pre-caching for the first few tokens\n\n> Appendix G: This appears to be orthogonal to the experiments in the body text\n\nWe agree that Appendix G would make more sense as part of a separate work. It will be removed in the final revision\n\n> Appendix C: tilde\n\nWill fix", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The concept of \"myopia gap\" may be of interest to LLM researchers and this paper makes a valiant effort in justifying its use as a tool for understanding what information LLMs acquire during training. The paper shows that myopic training of LMs indeed results in a high myopia gap in a synthetic data scenario. While the experiments on larger LMs are light, the concept of myopia gap could be a useful tool for understanding how more realistically-sized/trained LLMs operate.", "reasons_to_reject": "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_to_authors": "- 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.\n- 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?\n- 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?\n- Why does the myopically-trained model incur lower loss on the first few token positions than the standard LM (Figure 5)?\n- Appendix C: tilde is used inconsistently (e.g., when defining Q, K, V, and also over X).\n- 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.\n\n[1] Stern, Mitchell, Noam Shazeer, and Jakob Uszkoreit. \"Blockwise parallel decoding for deep autoregressive models.\" Advances in Neural Information Processing Systems 31 (2018).", "extra_scores": {}}, {"review_id": "EbCo5jBLWb", "reviewer": "Reviewer_ZSJf", "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.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review.\n\n> I would note that the LM experiments are on only 124M GPT-2 models 4.6M texts of length 64 from MS MARCO.\n\nThe final revision will contain further myopia gap experiments on models from the Pythia suite ranging up to 2.8B parameters in size, evaluated on both cross-entropy and various NLP benchmarks.\n\n> Other than studying the behavior of models, do you think there are any practical conclusions?\n\nWe 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.\n\nOne manifestation of this is our finding (which will be included in the final version) that pre-caching is necessary to make use of computation on the forward passes of filler tokens, in the sense of Pfau et al. \"Let's think dot by dot\". We train and evaluate vanilla and myopic transformers on an integer multiplication task. 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.\n\nAnother possibility, which we have not yet explored, is that of automatically swapping in a myopic model 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).\n\nAlso, if indeed the myopia gap remains acceptably small for production-sized models, there may be some efficiency gains to be had be training myopically, and thus not computing the off-diagonal gradient terms. Again, we have not yet investigated this possibility, and leave it to future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "An interesting study on the dynamics of language modeling training with nice formulations and experimental setup to create feasible but interesting explorations.", "reasons_to_reject": "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_to_authors": "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?", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept