🤖 machine learning

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

This article proposes a sample-efficient framework for continuous knowledge expansion in large language models by representing memory as a Markov transition matrix and employing a token-to-vocabulary mapping strategy with minimal embedding updates to avoid catastrophic forgetting.

Original authors: Kaustubh Pethkar, Ziyang Xiong, Zuofeng Shang, Yingcong Li

Published 2026-05-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Kaustubh Pethkar, Ziyang Xiong, Zuofeng Shang, Yingcong Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Large Language Model (LLM) as a vast, highly organized library of stories. Every time the model reads a sentence, it is as if a librarian looks at the last word and guesses which word comes next.

Normally, if you want to teach this librarian a brand-new word (such as a new scientific term or a slang expression), you must retrain the entire library. However, here lies the problem: when you retrain the entire library to learn the new word, the librarian often forgets how to tell the old stories correctly. This is known as "catastrophic forgetting."

This article proposes a clever, low-effort way to add new words without damaging the old ones. Here is the breakdown with simple analogies:

1. The Library as a "Traffic Map" (The Markov Idea)

The authors view the language model not as a complex brain, but as a traffic map.

The Nodes: Every word in the dictionary is a city on the map.
The Roads: The connections between words are roads. If you are in the city "The," certain roads lead to "cat," "dog," or "sun."
The Memory: The model's "memory" is merely the map of these roads. It knows that "The" usually leads to "cat" with a certain probability.

2. The Problem: Adding a New City

If you want to teach the model a new word (let's call it "Zorp"), you are essentially adding a new city to the map.

The Old Way (Full Fine-Tuning): You try to redraw the entire map to include "Zorp." In doing so, you accidentally delete or alter the roads between the old cities. The librarian forgets that "The" leads to "cat" and instead starts saying that "The" leads to "Zorp."
The Article's Way: Instead of redrawing the entire map, you simply add a sign for "Zorp." You say: "Hey, when you see 'Zorp,' treat it exactly the way you treat 'The' or 'Cat'."

3. The Solution: The "Token-to-Dictionary" Map

The article proposes a strategy called Token-to-Dictionary Mapping.

Imagine you have a new, strange word "Zorp." Instead of teaching the model the meaning of "Zorp" from scratch, you tell the model: "When you see 'Zorp,' just act as if it were the word 'Star'."
The model does not need to learn new roads for "Zorp." It only needs to know that "Zorp" points to the same destination as "Star."
Since the model does not need to change the existing roads (the transition probabilities between old words), it never forgets the old stories.

4. The "Sample Efficiency" (Why It Is Fast)

The authors prove mathematically that this is incredibly efficient.

The Analogy: If you want to learn a new word, you do not need to read the entire library again. You only need to read a few examples of how this new word is used in sentences.
The Math: The number of examples you need depends on how many existing words you map the new word to. If you map "Zorp" to only 5 common words, you need only a tiny amount of data to learn it. You do not need to worry about the size of the entire library (which might have 100,000 words).

5. The Experiment: Teaching Mathematics and Made-Up Words

The researchers tested this with two scenarios:

The Magic Operator: They taught a model a special symbol (like ⟨spec⟩) that meant "multiply."
- Result: The model learned very quickly to multiply with the new symbol.
- The Gain: Crucially, the model still knew how to add numbers. With other methods, learning the new symbol caused the model to forget how to add. This method kept the addition skills perfect.
Made-Up Words: They invented 100 nonsense words (like "glor" and "zorp") and inserted them into sentences.
- Result: The model learned to use these invented words correctly without forgetting how to speak English.
- Comparison: Other methods (such as standard fine-tuning or LoRA) caused the model to forget English while learning the invented words. The new method caused not a single instance of forgetting.

Summary

Consider this method as adding a new room to a house without tearing down the walls of the existing rooms.

Old Way: To add a new room, you rebuild the entire foundation. The old rooms crack and fall apart.
New Way: You build the new room and attach a door that leads directly into an existing hallway. The old hallway remains exactly as it was, and the new room fits perfectly inside.

The article claims this is a mathematically proven way to expand a language model's vocabulary with very few data points while suffering absolutely no loss of prior knowledge.

Technical Summary: Memory as a Markov Matrix

Problem Statement

Large Language Models (LLMs) face a critical challenge in continuous learning: integrating new knowledge (e.g., new vocabulary, entities, or domains) without triggering catastrophic forgetting. Existing approaches such as full fine-tuning, regularization, or replay-based methods often lead to the degradation of previously learned behaviors once new knowledge is introduced. Furthermore, these methods typically rely on large-scale weight updates, which are often irreversible and computationally expensive. The central question this work addresses is: How can LLMs integrate a small amount of new knowledge without suffering catastrophic forgetting?

Methodology

1. Markov Formulation of Language Generation

The authors propose a principled framework that models autoregressive language generation as a Markov process over tokens.

State Space: Tokens correspond to the states of the Markov chain.
Transition Matrix: The model's memory is represented by a Markov transition matrix, where the probability of the next token $p_\theta(v | u)$ defines the transition dynamics from token $u$ to $v$ .
Knowledge Expansion as State Expansion: The integration of new tokens ( $U$ ) is formalized as an expansion of the state space from $V$ to $V \cup U$ .
Zero-Forgetting Condition: To prevent forgetting, the transition probabilities between already existing tokens ( $V \times V$ ) must remain invariant. New tokens are integrated by defining their transition distributions ( $q(u)$ ) over the existing vocabulary $V$ , without altering the original $V \to V$ transitions.

2. Token-to-Vocabulary Mapping

To realize this expansion without modifying the original transition structure, the authors introduce a Token-to-Vocabulary Mapping Strategy:

Concept: New tokens are not learned as independent entities with unique transition dynamics. Instead, each new token $u \in U$ is mapped to a sparse combination of existing tokens in $V$ .
Embedding Tuning: This mapping is implemented via an Embedding Tuning Algorithm. Only the embedding vectors corresponding to the newly introduced tokens are updated (or extended). The original token embeddings and the model's transition computation graph remain frozen.
Mechanism: By updating only the embedding of the new token, the model learns to represent the new token as a linear combination of existing semantic structures (embeddings), effectively reusing the learned transition dynamics of the base vocabulary.

3. Theoretical Analysis

The work derives bounds on the sample complexity for learning these new tokens under the mapping strategy:

One-to-One Mapping: If a new token behaves identically to a single existing token, the sample complexity scales with the KL-divergence margin ( $\gamma$ ) between token distributions, independent of the total vocabulary size $T$ .
Sparse Combination: In the more general case where a new token is a sparse combination of $s$ existing tokens, the required number of samples scales as $\tilde{O}(s)$ . Crucially, this complexity depends on the sparsity of the representation ( $s$ ) rather than the total vocabulary size ( $T$ ) or the model dimensionality.
Higher-Order Chains: The framework can be extended to higher-order Markov chains (context-dependent generation) by reconstructing the state space, where the sample complexity depends on the effective branching factor of natural language rather than the exponential growth of the vocabulary.

Main Contributions

Markov Formulation: The work formalizes language generation as a Markov process and interprets model memory as a transition matrix. It frames vocabulary expansion as a state expansion problem, where forgetting is eliminated by preserving original transition dynamics.
Theoretical Guarantees: It provides guarantees on sample complexity, demonstrating that learning new vocabulary is efficient when new tokens can be sparsely mapped to existing tokens. The costs scale linearly with the number of mapped tokens ( $s$ ) rather than the total vocabulary size.
Embedding Tuning Algorithm: The authors propose a practical algorithm that updates only the embeddings of new tokens. This method induces Zero-Forgetting by construction, as it leaves the original token-to-token transition probabilities untouched.
Empirical Validation: Extensive experiments validate the theory in three scenarios:
- Arithmetic Operators: Learning a new token for multiplication showed that embedding tuning achieved high accuracy with minimal data while preserving addition capabilities (Zero-Forgetting), whereas full fine-tuning led to catastrophic forgetting of addition.
- Synthetic Vocabulary: On synthetic nonsense words, embedding tuning achieved strong performance with Zero-Forgetting on the WikiText benchmark, while LoRA and full fine-tuning caused significant forgetting.
- Cross-Lingual Expansion: When expanding the vocabulary for Spanish, German, and Arabic, embedding tuning maintained near Zero-Forgetting in English while achieving the best adaptation to the target languages.

Results and Significance

The work demonstrates that large-scale weight updates are unnecessary for acquiring small amounts of new knowledge. By treating new tokens as sparse extensions of the existing embedding vocabulary, the proposed method achieves:

Sample Efficiency: High performance with significantly fewer training samples compared to full fine-tuning.
Zero-Forgetting: Theoretical and empirical evidence confirms that updating only the embeddings of new tokens fully preserves the model's original capabilities.
Scalability: The approach is computationally efficient and requires updates to only a tiny fraction of parameters (the embeddings of new tokens).

The authors claim that this work bridges the gap between theoretical sample complexity and practical continuous learning. By modeling LLMs as Markov processes, they provide a formal justification for why embedding tuning prevents forgetting and offer a principled alternative to heuristic, parameter-efficient fine-tuning methods. The framework suggests that the "cost" of learning new knowledge is determined by the semantic sparsity of the new concepts relative to the existing vocabulary, rather than the sheer size of the model.