Future Lens: Anticipating Subsequent Tokens from a Single Hidden State (2023)

Abstract: We conjecture that hidden state vectors corresponding to individual input tokens encode information sufficient to accurately predict several tokens ahead. More concretely, in this paper we ask: Given a hidden (internal) representation of a single token at position t in an input, can we reliably anticipate the tokens that will appear at positions ≥ t + 2? To test this, we measure linear approximation and causal intervention methods in GPT-J-6B to evaluate the degree to which individual hidden states in the network contain signal rich enough to predict future hidden states and, ultimately, token outputs. We find that, at some layers, we can approximate a model’s output with more than 48% accuracy with respect to its prediction of subsequent tokens through a single hidden state. Finally we present a “Future Lens” visualization that uses these methods to create a new view of transformer states.

## introduction do hidden states in large language models (llms) encode tokens farther than a single token ahead? if so, how can we decode this sequence of tokens from a single state? in this work we empirically investigate these questions using gpt-j-6b @xcite . we train models to predict hidden states several tokens ahead of a given position t based only on a contextualized representation of the input at this position. auto-regressive transformer language models are typically trained to predict one token ahead, but recent work has hinted that individual hidden states may contain more information than just probabilities of the following token. for example, @xcite trace information flow from subject tokens to associated attribute predictions many steps ahead. elsewhere, @xcite suggest that neurons in early layers are dense with information, while middle layers have dedicated neurons that represent high-level contextual features. other related efforts have passed hidden intermediate states directly to the decoder head (skipping in-between layers) to "verbalize" such embeddings @xcite @xcite . studies of memorization @xcite @xcite have identified the presence of very long memorized sequences generated by language models, and @xcite shows that progressively dropping layers during computation can still achieve a similar prediction output of the model when compared against their fully computed model run. in this work we ask: to what extent can we extract information about future (beyond subsequent) tokens from a single hidden token representation? to answer this, we conduct three experiments. first, extending the ideas of tuned lens @xcite and the logit lens (nostalgebraist, 2020), we train linear models to approximate future model predictions several tokens in the future, in order to reveal the extent to which individual hidden states may directly encode subsequent tokens. second, we perform a causal intervention study in which we transplant individual hidden states from one context to a completely different context and measure the extent to which future tokens that were predicted in the original context can be predicted in the foreign context. finally, we fit a "soft prompt" to explicitly learn an optimal prompt that permits reading out information about subsequent tokens from a hidden state. ## methods to unveil the information about "future" tokens implicitly encoded in a single transformer state vector, we develop and compare several methods for predicting future tokens from a single hidden state. each of our methods has the same goal: extract accurate predictions of a model's probability distribution several tokens ahead, based on the information in only one hidden state at a single layer at one token of the transformer. for our evaluations we use an autoregressive transformer @xcite language model defined as a function g : g takes in a sequence of tokens @xmath0, ...., x t ] ∈ x, x @xmath1, 1] dv , which (greedily) predicts the next-token x t +1 = argmax y t . to generate additional tokens, the top predicted token x t +1 is added to the sequence of tokens [x 1 , ...., x t , x t +1 ] and the process is repeated until the next n tokens are produced. to calculate each predicted probability distribution from an input sequence x, the transformer performs a sequence of computations at l layers; this can be decomposed as: where the first step e :→ r d h embeds each input token into an initial hidden representation, e( when generating a sequence of tokens beyond the given starting prefix of length t , we write: our goal is to devise methods that can anticipate what g will predict for y t +1 through y t +n from only a single hidden state at h l t . ## related work knowledge prediction and manipulation recent works have delved into llm internals to better understand how such models predict the next token at each computation step. @xcite , for instance, find that the feed-forward layers in transformers operate as key-value memories, allowing one to intervene at those layers to modify the next token output @xcite . frameworks such as rome @xcite and memit @xcite scale such manipulations to edit knowledge in stored in llms. the consensus that has emerged in these papers is that some early-middle and late layer calculations contribute the most to the final predicted token. tools such as logit lens (nostalgebraist, 2020) and tuned lens @xcite allow us to look at the top-k values of the transformer at every layer and token to see early next-token predictions. @xcite used logit lens to visualize semantic information flow in gpt-2 models. in contrast to these approaches, we aim to characterize how the current hidden state would affect the prediction of not only the next token, but also tokens farther ahead. early exit decoding to optimize the running time and space requirements of training models, prior work has looked at "early exit" strategies, which usually involves stopping at earlier layers of computation and estimating the final predictions based on those computations @xcite @xcite @xcite . the takeaway from these methods is that it is possible to achieve prediction performance comparable to that observed when all layers are used even when dropping a couple of computational layers for each token. for instance, din and colleagues (2023) used linear transformations to predict a later layer's hidden representation from an earlier layer at the same token. this approach was able to preserve ∼95% of the full transformer model outputs on gpt-2 @xcite and bert @xcite . this result implies that initial model layers encode information that to a large degree determines the final output. in this work we test the limits of this phenomenon by evaluating the degree to which a single hidden representation for the token at position t can be used to predict tokens multiple steps ahead (i.e., at t + n ). ## discussion in this paper we explored the degree to which we are able to decode multi-token outputs subsequent to a particular token on the basis of its hidden representation alone. the results in to further explore the contexts in which these methods seem better (or worse) able to predict subsequent tokens, we categorize input token (the last original context token) into eight (non-mutually exclusive) categories, shown in while all categories of token types are predicted better by the learned prompt than by the linear model, the relative improvement is highest when the last context token is a lowercase token preceded by a space, or a longer token. this suggests that information about how to complete long words may not be immediately accessible by a linear model decoder, but that they can be made accessible by using the parameters of the pretrained model as done by the learned prompt intervention method. we have also observed that the accuracy of predicting subsequent tokens is correlates with the model's confidence in its next token prediction. in the case of @xmath4, for instance, the learned prompt intervention method's calibrated accuracy is 26%, 57%, 77%, and 95% for model confidence groups of 0-30%, 30-60%, 60-90%, and 90%-100%, respectively. these trends appear in @xmath5. this suggests that we might gainfully use this decoding method as a probing tool, trusting that predicted future tokens are generally accurate when the model is confident. does future information appear only in the presence of higher-level concepts? for example, one might hypothesize that in cases the language model predicts an entire named entity, that the probing method might decode future predictions more accurately. to investigate this, we performed sub-group analyses on test results to characterize how well the best probing method performed specifically for multi-token named entities. interestingly, we found little difference: when examining just the named entity cases, we observe similar or slightly lower accuracy: 44%, 42% and 37% for @xmath6, 2, 3, suggesting that future information is present broadly, not only for long entity names. in sum, we have found that a single hidden state encodes information about outputs more than one token ahead, and we have demonstrated three different methods that can decode them for gpt-j-6b. application: future lens we apply the learned prompt intervention method to create a novel probing tool we call the future lens. given a soft prompt, we perform the intervention using the states arising from the user's prompt to provide a view into what the hidden states encode about future tokens. in figure 6 , we show an example for the prompt: "marty mcfly from". the future lens reports the anticipated four tokens from every hidden state in the model (across layers). in the future lens visualization, every cell represents a hidden state from a particular layer ("l{digit}") at a specific token. the shade of each cell indicates the average confidence of the model with respect to the corresponding token predictions (darker shades indicate greater confidence). for example, at the cell representing the hidden state at layer 25 at the token "from", we can see that the confidence in the predicted tokens "back to the future" is strong. this particular state suggests that the llm already knows that marty mcfly is related to the back to the future movie. interestingly, the model also assumes "marty" to have the surname donough. returning to the predictions at token "from", we see that the early layers seem to first predict countries such as australia or cities such as boston. however, through future predictions, we can see the model begins to associate marty mcfly with a movie around layer 6. hence, through this tool, we can gain further insights about the model's chain of predictions at every hidden state. all code and data for demo and implementation is made available at: https: //github.com/koyenapal/future-lens