PaperlayerView exact revision on arXiv

Source-linked AI summary

Memory Retrieval and Consolidation in Large Language Models through Function Tokens

Shaohua Zhang, Yuan Lin, Hang Li

arXiv:2510.08203v1cs.CLcs.AI

TL;DR

LLM memory retrieval and consolidation remain poorly understood despite progress in understanding neuron activations. The paper advances the Function Token Hypothesis and reports experiments whose token-feature analyses and pre-training loss patterns support its account.

  • Problem

    How LLMs retrieve memory during inference and consolidate it during pre-training remains poorly understood.

    Source: §1 · p. 2

  • Method

    The paper proposes the Function Token Hypothesis and tests it with token-feature analysis and pre-training loss comparisons across function/content token pairs.

    Source: §1 · p. 2

  • Results

    Function tokens reactivate predictive context features during inference, while predicting content tokens following function tokens dominates pre-training and feature growth.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The experiments provide strong evidence for the Function Token Hypothesis as an account of LLM memory mechanisms.

    Source: §7 · p. 11

    How function tokens acquire dynamic predictive-feature activation remains unresolved; its emergence likely involves model architecture, data, training loss, and learning algorithm.

    Source: §7 · p. 12

Abstract

from arXiv · show

The remarkable success of large language models (LLMs) stems from their ability to consolidate vast amounts of knowledge into the memory during pre-training and to retrieve it from the memory during inference, enabling advanced capabilities such as knowledge memorization, instruction-following and reasoning. However, the mechanisms of memory retrieval and consolidation in LLMs remain poorly understood. In this paper, we propose the function token hypothesis to explain the workings of LLMs: During inference, function tokens activate the most predictive features from context and govern next token prediction (memory retrieval). During pre-training, predicting the next tokens (usually content tokens) that follow function tokens increases the number of learned features of LLMs and updates the model parameters (memory consolidation). Function tokens here roughly correspond to function words in linguistics, including punctuation marks, articles, prepositions, and conjunctions, in contrast to content tokens. We provide extensive experimental evidence supporting this hypothesis. Using bipartite graph analysis, we show that a small number of function tokens activate the majority of features. Case studies further reveal how function tokens activate the most predictive features from context to direct next token prediction. We also find that during pre-training, the training loss is dominated by predicting the next content tokens following function tokens, which forces the function tokens to select the most predictive features from context.

1 Introduction

The section frames two unresolved questions—how LLMs retrieve memory at inference and consolidate it during pre-training—and proposes analyzing function and content tokens to address them.

1 Introduction

  • Sparse Autoencoders decompose neuron activations into interpretable features that offer insight into the composition of Transformer circuits.The paper motivates this approach as a way to study LLM neuron activations, while noting that memory mechanisms remain poorly understood.Source: §1 · p. 2
  • For experiments, tokens are classified by corpus frequency as approximations of linguistic function and content words; function words include punctuation, articles, prepositions, and conjunctions.The frequency-based token classification is presented as an approximation of the linguistic categories.Source: §1 · p. 2
  • Bipartite analysis and case studies support a two-stage account: function tokens reactivate predictive context features at inference, while function→content prediction dominates pre-training and feature growth.The authors report that function tokens activate a large proportion of features, although they are few in number.Source: §1 · p. 2
  • Memory retrieval denotes activation of features and circuits, whereas consolidation denotes parameter learning that forms and expands features and circuits.The paper studies GPT-type Transformer decoder models and uses Sparse Autoencoders to derive features from activations.Source: §1 · p. 2
  • The Function Token Hypothesis assigns function tokens predictive-feature selection at inference and parameter-updating feature expansion during pre-training from content-token prediction.The authors propose this account to explain LLM memory mechanisms.Source: §1 · pp. 2–3

2 Preliminary

The preliminary section frames FFN layers as key-value memories and uses feature decomposition to describe their internal representations. It defines function and content tokens by corpus frequency, then characterizes their frequency and document coverage.

2.1 Model Memory and Superposition Phenomenon

  • FFN layers are described as key-value memories: input similarity to keys weights value vectors whose sum forms the output.The authors interpret Transformer knowledge as represented in FFN parameters; attention layers also form key-value memories using softmax weighting.Source: §2.1 · p. 3
  • Sparse Autoencoders decompose neuron activations into interpretable features, which can outnumber neurons when neurons represent multiple meanings.The decomposition is learned through a reconstruction task and represents activations as linear combinations of features.Source: §2.1 · pp. 3–4

2.2 Function Tokens and Content Tokens

  • The authors label 122 tokens covering 40% of sampled corpus occurrences as function tokens, assigning the remaining tokens to the content-token group.They tokenize SlimPajama-627B with the LLaMA-3.1 tokenizer and use frequency as a proxy for the linguistic distinction, with punctuation among the exceptions.Source: §2.2 · p. 4
  • 22.58% of the corpus is accounted for by the 15 most frequent tokens, consistent with a Zipf distribution in which a few tokens are frequent and most are rare.Figure 2 presents token-frequency statistics.Source: §2.2 · p. 4, Figure 2 · p. 4
  • High-frequency tokens typically occur across most documents, whereas low-frequency tokens are concentrated in a smaller fraction of documents.Figure 3 partitions documents into equal bins; the function token “of” appears similarly across documents, while “Tokyo” occurs in only a few.Source: §2.2 · pp. 4–5, Figure 3 · p. 5

3 Memory Retrieval through Function Tokens

The section examines how function tokens relate to model features during inference. A small set of function tokens activates most features, and case studies show that these tokens can reactivate context-dependent predictive features that influence model outputs.

3.1 A Few Function Tokens Activate Most Features

  • A small set of function tokens activates most features; the top 10 frequent tokens activate more than 70% of features in the middle layer.The analysis constructs a token-feature bipartite graph from SAE-decomposed activations.Source: §3.1 · p. 6

3.2 Feature Reactivation via Function Tokens

  • Function tokens such as ‘:’, ‘the’ and ‘\n’ propagate and recreate features first activated by context tokens, with different feature combinations for prompts differing only in the named country.In the examples, ‘Chinese’ and ‘Russia’ initially activate the ‘Speak Chinese’ and ‘Russia’ features; replacing ‘Russia’ with ‘UK’ changes the combination and output.Source: §3.2 · p. 7
  • Editing activations at the final ‘\n’ can change the answer: activating ‘Speak Chinese’ switches ‘Japan’ to ‘日本’, while activating ‘Russia’ changes the response to ‘Russia’.Jointly activating ‘Speak Chinese’ and ‘UK’ yields ‘英国’; feature steering can also produce contextually appropriate responses such as ‘Moscow State University’.Source: §3.2 · p. 8, Figure 7 · p. 8

4 Memory Consolidation through Function Tokens

The section tracks feature growth and token-group prediction losses during pre-training. Function→content prediction has the highest loss and dominates optimization, while learned features grow and function-token prediction becomes easy earlier.

4.1 Pre-Training Setup

  • The study trains 1.5B- and 8B-parameter models from scratch on one epoch of SlimPajama-627B using LLaMA-3.1-8B architecture and training hyperparameters.The 1.5B model has two layers; the 8B model uses the original 32 layers.Source: §4.1 · p. 8

4.2 Memory Consolidation as Feature Expansion

  • The number of learned features grows across training checkpoints, while function tokens consistently activate most features and the disparity from content tokens widens.The analysis uses 500,000 sampled activations at each of three checkpoints: 3,000, 50,000 and 130,000 steps.Source: §4.2 · p. 9, Figure 8 · p. 9

4.3 Loss on Function and Content Tokens

  • Function→content has the highest loss in both models, dominates optimization, and primarily drives feature growth during pre-training.The authors connect this difficult prediction task with function tokens developing the capability to reactivate predictive context features.Source: §4 · p. 8, §4.2 · p. 9, §4.3 · p. 10
  • For both model sizes, loss falls faster and converges lower when predicting function tokens than when predicting content tokens.Function-token prediction reaches very low loss early; representative tokens ‘the’, ‘of’ and ‘,’ converge rapidly within the first 3,000 steps.Source: §4.3 · pp. 9–10, Figure 9 · p. 9
  • Scaling from 1.5B to 8B parameters reduces function→content loss from 4.88 to 4.27 and content→content loss from 3.69 to 3.08.The reported reductions are larger than for content→function (1.90 to 1.64) and function→function (2.12 to 1.87).Source: §4.3 · p. 10
  • Generation progresses from random output at step 100 to function tokens at 3,000, locally coherent phrases at 50,000, and later long-range predictions in the 8B model.Predicting a token after ‘...struggle with’ requires recalling the earlier context ‘learning to read’.Source: §4.3 · p. 10, Figure 10 · p. 10

5 Function Token Hypothesis

The authors propose the Function Token Hypothesis to explain memory retrieval and consolidation in LLMs. It describes function tokens as activating predictive context features during inference and content-token prediction after function tokens as driving parameter updates and feature learning during training.

5 Function Token Hypothesis

  • During inference, function tokens activate the most predictive context features to direct next-token prediction; during pre-training, predicting content tokens after function tokens drives parameter updates and feature learning.The authors state that their investigations validate the hypothesis while acknowledging the complexity of LLM memory mechanisms.Source: §5 · pp. 10–11
  • The authors explain the proposed mechanism through nested text chunks: predicting after a function token starts a harder next-chunk task requiring context-wide meaning.They propose that this challenge compels predictive-feature activation during training and reactivation during inference.Source: §5 · p. 11

6 Related Work

The related work reviews neural memory and interpretable feature analysis, alongside prior findings about function-token-like patterns in LLMs. The authors position their hypothesis as an account of how such tokens relate to memory retrieval and consolidation.

6 Related Work

  • Prior work on neural memory spans Hopfield networks, restricted Boltzmann machines, key-value memories and Transformer features uncovered through dictionary learning.Existing studies show that decomposed features can steer model behavior by controlling access to memories and directing generation.Source: §6 · p. 11
  • Prior LLM findings concern separator-token activations and attention, formatting tokens, high-entropy reasoning tokens and pivot tokens; the authors group these patterns as function-token behavior.The passage cites separator tokens [45], attention [6], formatting [23] [26] [56] [57], high-entropy tokens [53] and pivot tokens [1].Source: §6 · p. 11

7 Conclusion and Open Questions

The paper concludes that function tokens support memory retrieval during inference and feature learning during pre-training. It identifies unresolved questions about how these mechanisms arise and change across training and model layers.

7 Conclusion and Open Questions

  • The function token hypothesis links inference-time predictive-feature activation to next-token prediction and pre-training feature learning.The authors report strong experimental evidence for the hypothesis.Source: §7 · p. 11
  • How function tokens acquire their ability to activate predictive features, and how architecture, data, training loss, and learning algorithms interact, remain open questions.The authors say investigating this interaction is essential to understanding the phenomena.Source: §7 · p. 12
  • How post-training changes function-token activation patterns remains unresolved, despite reported capability gains from few training steps and function-token-only reinforcement learning.The passage connects those gains to the possibility that post-training activates latent pre-training capabilities.Source: §7 · p. 12
  • Scaling data, computation, and model size reduced pre-training loss while increasing learned features; function tokens activated most features with a scale-free degree distribution.The dynamics of feature formation, the reason for the scale-free property, and whether specific principles govern them remain unclear.Source: §7 · p. 12
  • Case studies found stronger interpretability and steerability in middle layers, but why steerability concentrates there rather than in shallow or deep layers remains unknown.Source: §7 · p. 12

A Steering Method for Large Language Models

The steering method identifies a trait-specific feature by comparing trait-enforcing and non-trait prompts, selecting an informative layer, then applying a feature-specific activation to the prompt’s last function token.

A Steering Method for Large Language Models

  • Two prompt sets—a trait-enforcing prompt and 20 prompts without the trait—identify the feature and test steering effectiveness.The trait-enforcing prompt identifies the relevant layer and feature; non-trait prompts form the test set.Source: A Steering Method for Large Language… · p. 18
  • The method selects the layer where adding the final function-token activation from the trait prompt yields the highest test-prompt success rate.For “Speak Chinese,” “Russia,” and “UK,” the selected layer was 26.Source: A Steering Method for Large Language… · p. 18
  • At a selected layer, features are ranked by activation strength in descending order, and binary search identifies the smallest k whose top-k features enable the trait while top-(k−1) does not.For Speak Chinese, the identified feature is ID 15261 at layer 26; other examples are Russia (feature ID 9591, layer 26) and UK (feature ID 13751, layer 26).Source: A Steering Method for Large Language… · p. 18
  • The identified feature-specific steering vector is added to the prompt’s last function token to induce traits such as speaking Chinese, Russia, or UK.Source: A Steering Method for Large Language… · p. 18

B Additional Case Study

The appendix presents additional examples of steering by editing the final newline function-token activation.

B Additional Case Study

  • The authors present more interesting examples of steering activations on function tokens, with Figure 11 cited as showing those examples.Source: B Additional Case Study · p. 18

C SAE Training Details

This appendix describes the JumpReLU-SAE, its reconstruction-and-sparsity objective, and the training settings used to obtain feature representations at different checkpoints.

C SAE Training Details

  • The JumpReLU-SAE encodes an activation vector into features and decodes those features to reconstruct the input.The encoder and decoder are specified by equations 7 and 8.Source: C SAE Training Details · p. 18
  • SAE training balances reconstruction quality against sparsity using a reconstruction loss and an L0 penalty; cross-entropy reconstruction scores range from 0 to 1, with higher scores indicating more faithful reconstructions.Source: C SAE Training Details · p. 19
  • Training uses one billion pre-training activations, a dictionary 16 times the activation dimension, and a learning rate of 1 × 10^-5.The specified batch size is 4096; JumpReLU bandwidth is 0.02 and its initialization threshold is 0.01.Source: C SAE Training Details · p. 19
  • To make feature counts comparable across early, intermediate, and late checkpoints, the authors tuned λ to similar reconstruction scores; feature representations became more complex and harder to decompose during pre-training.They used λ = 10, 4, and 2.5 for the 3,000-, 50,000-, and 130,000-step checkpoints, respectively.Source: C SAE Training Details · p. 19, Figure 12 · p. 20

D Function Token List

Table 2 lists the function tokens identified in the experiments, ranked by frequency in SlimPajama-627B; tokens absent from the table are classified as content tokens.

D Function Token List

Loading 2510.08203v1…