SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 4: A. The predictions of the representation straightening hypothesis for the ground-truth (natural, human-produced) sentences (red line) vs. model-generated sentences (blue line). The hypothesis predicts that the trajectories of the model-generated sentences should be straighter given that linear extrapolation in the internal state space is hypothesized to serve as the critical prediction mechanism. In other words, if the model is internally producing a low-curvature trajectory, then self-generated sequences should have lower curvature than sequences generated by humans (given that next-word prediction is not the only objective that guides human language production). B. Curvature changes across the layers of the network for 5,815 pairs of 10-token sentences: the ground-truth sentences (red line; each dot is a layer) come from the UDsubset8408 set and the model-generated sentences (blue) are generated from the prompt consisting of the first three tokens of the ground-truth sentences (the model generates the subsequent 7 tokens using a greedy approach; Methods). Model-generated sentences show a greater drop in curvature reduction relative to the ground-truth sentences. An example of ground-truth vs. model generated sentence is shown on top of the panel.
3.4 Experiment 4. The curvature of sentence representations is correlated with sentence surprisal.
In Experiments 1-3, we have focused on language models and established that models reduce the curvature of sentence trajectories in their internal state spaces. Using a set of 8408 sentences, we found consistent curvature reduction in the deep model layers across this corpus. However, we also observed some variability among the individual sentences in their curvature patterns across the layers (Figure 1C). In Experiment 4, we attempted to connect this variability in the geometry of sentence representations in the model to some linguistic features of the sentences that we know affect human language processing. In particular, we focused on surprisal—a measure of how expected a word is given context. Surprisal has been shown to affect behavioral (Levy, 2008; Smith and Levy, 2013) and neural (Willems et al., 2016; Shain et al., 2020) responses to language. So we asked whether average sentence surprisal (we used 3-gram surprisal, averaging across the words to derive a single measure for each sentence; see Methods) relates to the curvature of the sentence's trajectory in the model space. If higher curvature in the model space corresponds to greater difficulty in predicting the next word, we should observe a positive relationship between curvature and surprisal: sentences that are overall less predictable (more surprising) should have higher curvature. Moreover, this correlation should only be observed in deeper layers (Figure 5A).
The results are shown in Figures 5B-C. As expected, for the untrained model, we found no relationship between surprisal and curvature (0 correlation across layers, except for layer 0). However, for the trained model, the correlation starts to increase from the early layers toward the deeper layers, peaking around layer 20 (Figure 5C). This result suggests that the degree to which the model straightens a sentence trajectory internally is associated with how surprising the sentence is behaviorally. These results also generalize to smaller models (not shown) and to other surprisal metrics (other n-gram metrics and PCFG-parser-based surprisal; see Supplementary Information).
In this work, building on Henaff and colleagues' proposal for primate vision (Henaff et al., 2019; Henaff, 2018; Henaff et al., 2021), we established neural trajectory straightening as a representational hypothesis about how neural network language models perform next-word prediction. Models consistently reduced sentence curvature from early to middle layers, and this effect was only observed
Figure 5: A. A hypothesized relationship between curvature and surprisal in the early layers (top) and the middle layers (bottom) of the network. In early layers the representations are close to input, so the curvature is not correlated with sentence surprisal. In deep layers however the curvature is predictive model states, so there is a consistent correlation between the sentence curvature and sentence surprisal. B. The relationship between average sentence surprisal (3-gram surprisal; see Methods) and average sentence curvature for sentences in the UDsubset8408 set in layer 1 (top) vs. layer 20 (bottom). The lines show a linear regression fit to the data (for illustrative purposes). C. Pearson correlation between average sentence surprisal (3-gram) and average sentence curvature across the layers of the network for an untrained model (gray-contour dots) and a trained model (black-contour dots). The data points for the layers where the correlation reached significance are color-filled (grey for untrained; red for trained). For the untrained model, no relationship is observed in any of the layers, except for layer 0. For the trained model, there is a consistent increase in correlation between curvature and surprisal from the early to middle layers.
for trained models (Figure 2). Model size and training dataset size affected the model's ability to reduce curvature (Figure 3), and model-generated sentences exhibited lower curvature compared to natural human-generated sentences (Figure 4). Finally, average sentence curvature correlated with average sentence surprisal in the middle layers of the model (Figure 5).
Our results sit squarely within the efficient coding framework and establish temporal prediction as a form of efficiency that is born out of training autoregressive transformer models on next-word prediction. These findings may also explain why representations from the middle layers of transformer language models align best with human neural responses during language processing (e.g., Schrimpf et al., 2021; Goldstein et al., 2022; Caucheteux and King, 2022; Caucheteux et al., 2023; Toneva and Wehbe, 2019; Jain and Huth, 2018).
Using representational geometry to understand the relationship between the internal workings of artificial and biological systems and their behavior has gained momentum in recent years (Wang et al., 2018; Remington et al., 2018; Mante et al., 2013; Chung et al., 2018 for a review, see Chung and Abbott, 2021). For example, in a deep neural network of vision, Cohen et al., 2020 showed how the geometry of object manifolds changes across model layers and how it relates to model performance in object categorization. In the domain of lang, Hewitt and Manning, 2019 found that the representations in the middle layers of BERT best capture the hierarchical structure of sentences. In another line of work, Mamou et al., 2020 used manifold analysis to uncover how different features of words (e.g., part of speech) and sentences become separable in the deep layers of language models like BERT and GPT. More recently, Valeriani et al., 2023 investigated how geometric properties, such as intrinsic dimensionality, change across the layers of the transformer models, and found that dimensionality increases across layers before sharply decreasing in deeper layers of the bidirectional transformer models. These prior studies provide insights into the geometric properties of vision and language models, but, to our knowledge, no prior study has evaluated a representational-level hypothesis that connects language model behavior (next-word prediction) to neural trajectories of individual sentences.
In the representation straightening hypothesis, the problem of predicting the next token/word is reformulated as predicting the next state in the model's internal representation. Upon predicting the next state, the model can connect this new state to the features of the input/output (i.e., to words). Explicitly defining model's objective to build a predictive representation over its internal
representation, and not output, and would be an interesting future direction (Olshausen and Field, 1996). Another direction would be to evaluate representation straightening in human behavioral or neural responses to language. Finally, if representation straightening is a general mechanism for temporal prediction, it should be evident in other systems, biological and artificial, that have a core prediction objective; it would be important to evaluate the generality of this mechanism.
It is important to note that in the representation straightening hypothesis, the problem of predicting the next token/word is reformulated as predicting the next state in the model's internal representation. Upon predicting the next state, the model can then connect this new state to the features of the input/output (i.e., to words). Explicitly defining such transformation stages, would allow the model to build a predictive representation without any apparent behavior and would be an interesting future direction. Another exciting direction would be to evaluate representation straightening in human neural language data direction. Finally, if representation straightening is a general mechanism for temporal prediction, it should be evident in other systems, biological and artificial, that have a core prediction objective; it would be important to evaluate the generality of this mechanism.
5 Broader Impact and Limitations
Our work puts forward and provides support for a general hypothesis at the representational level about a mechanism that allows large language models to achieve good performance on next word prediction and potentially downstream tasks. This work adds to the growing body of research on the interpretability of AI models. A better, more mechanistic understanding of these models, and potentially other models with prediction objectives, can both i) suggest ways to improve model efficiency and robustness, and ii) provide insights into the relationship between neural networks and the human language system.
We acknowledge that our work could be improved in several respects. The results as they stand are compatible with at least two possibilities: (i) that predicting future inputs intrinsically gives rise to implicit next-state prediction, thus directly favoring linear state dynamics, or (ii) that predicting future inputs in domains like language benefits from slowly-changing contextual information, thus indirectly favoring slower (and more approximately linear) state dynamics. These possibilities can be distinguished in the future by training models on artificially created datasets that vary in the length of context that affects the predictability of an incoming element. If the effects we report here obtain across these different training datasets, that would support the first possibility; if instead the effects only hold for models trained on data where relatively long predictive contexts, that would support the second possibility.
Furthermore, we have not evaluated the effects on sentence curvature of other training objectives or fine-tuning for downstream tasks. Doing so can help understand the selectivity of the observed effects (i.e., do sentence representations get straightened in the middle layers only under the pressure of the next-word prediction objective?) and their robustness to adding other objectives to a pre-trained model. We have also not causally tested the straightening hypothesis, which would require ablating the model in such a way that only curvature is affected, and testing how next-word prediction behavior changes.
It is also important to note that we are not claiming that representation straightening is the only mechanism that models rely on to gain linguistic competence. However, to the extent that prediction is a core part of language learning and processing (in artificial as well as biological systems), we are showing that targeted inspection of geometric properties of sentence representation gives rise to a hypothesis about how prediction may be implemented in language models.
The authors would like to thank Eero Simoncelli, Olivier Henaff, Yoon Bao, and Cory Shain for helpful insights and discussions, Anya Ivanova, Chengxu Zhuang for feedback on the manuscript, as well as members of EvLab at MIT. EF was supported by NIH grants R01-DC016950 and U01-NS121471, and by funds from the McGovern Institute for Brain Research, the Brain and Cognitive Science Department, the Simons Center for the Social Brain, and MIT Quest for Intelligence.
Shadmehr, R., Smith, M. A., and Krakauer, J. W. (2010). Error correction, sensory prediction, and adaptation in motor control. Annu. Rev. Neurosci., 33:89-108.
Shain, C., Blank, I. A., van Schijndel, M., Schuler, W., and Fedorenko, E. (2020). fMRI reveals language-specific predictive coding during naturalistic sentence comprehension. Neuropsychologia, 138:107307.
Shannon, C. E. (1949). Communication theory of secrecy systems. Bell Syst. Tech. J., 28(4):656-715.
Smith, N. J. and Levy, R. (2013). The effect of word predictability on reading time is logarithmic. Cognition, 128(3):302-319.
Tishby, N., Pereira, F. C., and Bialek, W. (2000). The information bottleneck method.
Toneva, M. and Wehbe, L. (2019). Interpreting and improving natural-language processing (in machines) with natural language-processing (in the brain). In Wallach, H., Larochelle, H., Beygelzimer, A., dAlché-Buc, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32, pages 14954-14964. Curran Associates, Inc.
Valeriani, L., Doimo, D., Cuturello, F., Laio, A., Ansuini, A., and Cazzaniga, A. (2023). The geometry of hidden representations of large transformer models.
van den Oord, A., Li, Y., and Vinyals, O. (2018). Representation learning with contrastive predictive coding.
Wang, J., Narain, D., Hosseini, E. A., and Jazayeri, M. (2018). Flexible timing by temporal scaling of cortical responses. Nat. Neurosci., 21(1):102-110.
Willems, R. M., Frank, S. L., Nijhof, A. D., Hagoort, P., and van den Bosch, A. (2016). Prediction during natural language comprehension. Cereb. Cortex, 26(6):2506-2516.
Wiskott, L. and Sejnowski, T. J. (2002). Slow feature analysis: unsupervised learning of invariances. Neural Comput., 14(4):715-770.
Wolf, T., Debut, L., Sanh, V., Chaumont, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, J., Shleifer, S., von Platen, P., Ma, C., Jernite, Y., Plu, J., Xu, C., Le Scao, T., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. (2019). HuggingFace's transformers: State-of-the-art natural language processing.
Zhu, Y., Kiros, R., Zemel, R., Salakhutdinov, R., Urtasun, R., Torralba, A., and Fidler, S. (2015). Aligning books and movies: Towards story-like visual explanations by watching movies and reading books.