Memory as a Markov Matrix: Sample Efficient Knowledge Expansion via Token-to-Dictionary Mapping

Abstract: Continual incorporation of new knowledge is essential for the long-term evolution of LLMs. Existing approaches typically rely on parameter-update algorithms to mitigate catastrophic forgetting, yet they suffer from fundamental limitations: 1) forgetting is unavoidable as the amount of newly injected knowledge grows; and 2) model updates are often irreversible. As modern LLMs become increasingly expressive, it is natural to question whether large-scale weight updates are necessary for acquiring a small amount of new knowledge. In this work, we propose a principled framework that models autoregressive language generation as a Markov process over tokens, where model memory is represented by a Markov transition matrix. Under this formulation, incorporating new knowledge/tokens corresponds to extending the state space, and preserving existing transitions guarantees retention of previously learned knowledge. We then prove a sample complexity bound for incorporating new tokens via a token-to-dictionary mapping strategy. In particular, for learning the transition behavior of each new token, the required number of samples scales linearly with the number of existing tokens it is mapped to. To realize this mapping, we propose an embedding-tuning algorithm that requires minimal parameter updates and induces zero forgetting. Experimental results further demonstrate the effectiveness of our method and validate our theoretical findings.

• The paper introduces a Markov process formulation for LLMs that preserves existing knowledge during vocabulary expansion.
• It presents a token-to-dictionary mapping with embedding tuning, reducing sample complexity independent of vocabulary size.
• Experimental results across arithmetic, synthetic, and cross-lingual tasks validate the approach's efficiency and zero forgetting.

Memory as a Markov Matrix: Sample Efficient Knowledge Expansion via Token-to-Dictionary Mapping

Introduction and Motivation

Continual knowledge acquisition remains a major obstacle in scaling LLMs, primarily due to catastrophic forgetting during sequential adaptation. Typical parameter-update strategies—full or parameter-efficient fine-tuning, regularization-based continual learning, or memory-augmented approaches—suffer from either inevitable forgetting as new information accumulates or irreversible changes to preexisting knowledge. "Memory as a Markov Matrix: Sample Efficient Knowledge Expansion via Token-to-Dictionary Mapping" (2605.04308) reconceptualizes language generation as a Markov process, providing a principled framework for knowledge expansion in LLMs that circumvents these issues.

Markovian Formulation of LLMs

The core proposal is to model next-token generation in an LLM as a Markov chain over tokens, where the model’s memory is encoded as a Markov transition matrix. In this perspective:

• The vocabulary defines the discrete state space.
• Next-token probabilities become transition probabilities of the Markov process.
• Incorporating new knowledge or vocabulary is formalized as expanding the state space.

Under this formulation, editing or expanding the vocabulary corresponds to extending the Markov transition matrix, with the preservation of old transition probabilities guaranteeing retention of previously learned knowledge. This connection sharply elucidates the theoretical pathway for eliminating catastrophic forgetting: if new tokens are added without altering existing transition probabilities, previous knowledge is strictly preserved.

Token-to-Dictionary Mapping and Sample Complexity

The primary mechanism for integrating new tokens is a token-to-dictionary mapping. Here, each new token is mapped to a subset or combination of existing tokens, defining its transition behavior. The authors distinguish two regimes:

1. One-to-one mapping: Each new token mirrors the transition distribution of a single existing token. The sample complexity for identifying the correct mapping depends only on the KL-separation between transition distributions of existing tokens and not the vocabulary size.
2. Sparse combination mapping: A new token corresponds to a sparse mixture of $s$ existing tokens. The sample complexity per new token scales with $O(s)$, independent of model dimension or overall vocabulary size.

This yields a sharp theoretical guarantee: the cost of learning a new token is dictated by how many existing tokens it is semantically mapped to, not by global model characteristics.

Embedding Tuning as a Practical Realization

Embedding tuning is proposed as the algorithmic realization of this mapping. In this approach, only the embedding vectors of newly introduced tokens are updated, while the model parameters and embeddings for existing tokens remain frozen. The rationale is that the next-token transition behavior is determined by the embedding itself; updating only the new token's embedding optimally fits the intended transition distribution without impacting those for known tokens.

A formal theoretical analysis establishes that, under mild context and regularity assumptions, embedding tuning requires a sample size proportional to the sparsity of the token-to-dictionary mapping—offering orders-of-magnitude efficiency gains over full-model fine-tuning or even LoRA-style parameter-efficient baselines.

Empirical Results

Experiments rigorously validate the theory, spanning controlled and realistic settings:

1. Arithmetic Operator Learning: Introducing a special token as a new operator in Llama-3.2-3B, embedding tuning achieves accurate operation induction (over 80% accuracy) with as few as 500 training samples, outperforming full model fine-tuning especially in small data regimes and strictly preserving pretrained arithmetic capabilities.

   Figure 1: Accuracy of the special token operation $a \langle \text{spec} \rangle b = a \times b$ improves rapidly with embedding tuning, consistently surpassing full fine-tuning at low sample counts.

2. Synthetic Vocabulary Expansion: Across multiple models and adaptation methods, only embedding tuning eliminates catastrophic forgetting and achieves efficient adaptation as the number of synthetic training samples increases. Competing methods (e.g., LoRA, prompt tuning) exhibit increasing forgetting as the training budget grows.

Figure 2: Left: Decreasing test loss on synthetic vocabulary task as number of training samples increases; Right: Only embedding tuning achieves zero forgetting, while other methods degrade on the original task.

1. Cross-Lingual Vocabulary Expansion: Applying the methodology to real-word datasets ensures new tokens from Spanish, German, and Arabic are integrated with near-zero forgetting of English capabilities, outperforming LoRA, prompt tuning, and full fine-tuning on both adaptation loss and retention measures.

Synthetic Sentence Creation Process

Synthetic sentences are created by systematically replacing real tokens with uninformative, synthetic tokens within otherwise realistic linguistic contexts. This allows for controlled evaluation of adaptation and forgetting mechanisms.

Figure 3: Illustration of replacing real tokens in sentences with synthetic tokens to evaluate vocabulary expansion and knowledge retention.

Implications and Theoretical Contributions

By tightly linking LLM adaptation to Markov process theory, this work provides:

• Formally provable zero forgetting: If no transition probabilities among original tokens are perturbed, knowledge retention is guaranteed by construction.
• Sample complexity bounds unlinked from vocabulary size or model dimension: Learning cost for new tokens is determined by the mapping’s sparsity, sharply reducing data requirements.
• Generality to higher-order Markov transitions: The framework extends to arbitrary autoregressive context windows, with complexity governed by actual linguistic branching factor, not all possible prefix combinations.

Strong empirical results confirm that embedding tuning (only $~0.001\%$ of parameters) suffices for sample-efficient and stable continual learning—even outperforming stronger baselines in both accuracy and knowledge retention.

Limitations and Future Directions

• The current theoretical guarantees assume uniform occurrence rates for new tokens and an absence of transition structure among new tokens.
• Model capacity limits and token frequency biases must be considered for broader generalization.
• Extending the framework to support dynamic transition reweighting and online expansion (including new-to-new token transitions) warrants further exploration.

Conclusion

This work reframes continual vocabulary expansion in LLMs as Markov process state-space augmentation, linking sparse token-to-dictionary mapping to efficient and stable knowledge integration. Embedding tuning, motivated by and aligned with this theory, achieves zero forgetting and strong performance by construction, even in practical, data-constrained cross-lingual settings. The insights and guarantees provided set a foundation for future directions in scalable, stable LLM continual learning, with direct applications in evolving domain adaptation, interactive learning, and dynamic knowledge infusion.