Full-Context Teacher Distillation for Native Memory
A student that cannot reread its full history risks losing useful context, so training first builds segmented language modeling and then matches a full-context teacher while history is available only through evolving memory.
What it is
- Segmented continued training Train on one text segment at a time; earlier segments are available through memory rather than being reread as text.
- Teacher matching A full-context model supplies prediction targets while the memory-based student sees the current input and its carried memory.
- Answer-focused tuning During task training, distillation loss is applied to response positions, while history and queries help form memory and context.
Why it matters
This separates learning to carry history from learning to answer in a way that resembles a model with the whole history in view, while keeping the student’s access to earlier text mediated by memory.
Before getting into the training stages, here is the basic idea: first teach the model to keep doing language modeling as a sequence is split into pieces, with the past carried in memory. Then use a model that can still read the whole history as a teacher, and train the memory-based model to make similar predictions—especially on the answers people care about.
-
Context budget: A full-context model can consult earlier text directly, but the memory-based student cannot. During continued pre-training, the student sees the current segment and gets earlier information through the memory left by previous segments. The canonical example is a long document split into consecutive chunks: after one chunk is processed, its useful information has to travel in memory for the next chunk’s predictions. The final chunk is kept even when it is shorter than the chosen segment length. This is not the same as asking the model to reread the entire document at each step.
-
A behavioral target: Learning to predict the next token with memory gives the student a language-modeling skill, but the downstream goal is to answer questions using history. The training setup therefore pairs historical context, a query, and a reference response. Context is processed chronologically into memory; then the query and observed response prefix are used to compare the student with a teacher that can still see the text history. The teacher’s predictions are the target, not merely the reference answer tokens.
-
Where the loss goes: The teacher and student are evaluated on the same teacher-generated response trajectory under teacher forcing. That means the student is not supplying its own next prefix during this comparison. At a response position, the teacher’s distribution is trimmed to its top- tokens and renormalized there, while the student still assigns probabilities over the full vocabulary; in plain terms, it is penalized for probability it places outside the teacher’s selected set too. During post-training, direct loss is placed on response tokens, not on history, query, or padding. During continued pre-training, by contrast, cross-entropy on ground-truth next tokens teaches continuation with memory.
The stages answer different questions: can memory support ordinary continuation across segments, and can it support answers that resemble a full-context model? I like that separation because it makes the training signal legible, though matching teacher predictions on fixed trajectories does not guarantee the student will behave identically on its own rollouts. The remaining question is what the memory manages to retain for a query it has not yet seen.
Read the deep dive
Why People Wanted It
A model that reads a long history at answer time can use that history directly, but keeping it all as text defeats the point of a compact native memory. The real question is not simply how to store history; it is how to train a model to answer as though it still had the full history when it can only consult its evolving memory.
The method separates that problem into two jobs: learn to continue language modeling across segments using memory, then teach the memory-based student to match a full-context backbone on response distributions.
The Core Move
First train the student to carry useful history across text segments; then train its memory parameters to make the student’s next-token predictions resemble a frozen full-context teacher’s predictions.
The description gives the objective in words: the teacher keeps the textual history, the student receives the current segment or query and its incoming memory, and training compares their next-token behavior. No equation is needed to state this training recipe.
Why Segmented Language Modeling Teaches Memory Use
Now the version for readers who want more of the machinery: continued pre-training (CPT) makes earlier segments available only through the incoming memory state, so predicting the next segment requires using what the memory retained.
A sequence is divided into consecutive, non-overlapping segments. Within a segment, the student performs causal language modeling with that segment as its textual context; earlier segments are represented by memory. After the segment, the native storage procedure updates memory for the next step. Memory begins at zero for each independent sequence, continues across its segments, and the final segment is kept even if it is shorter than the specified segment length. The direct targets are ground-truth next tokens, trained with cross-entropy. History can shape memory without itself receiving direct loss.
This stage establishes language modeling under the memory constraint. But learning to predict ordinary continuation tokens does not by itself make the student’s answers match a full-context model—so how is that gap targeted?
Why Top-Token Distillation Is Off-Policy
Here is the part that gets technical: the teacher and student are evaluated on fixed teacher-generated trajectories under teacher forcing, rather than on histories generated by the student’s own policy.
The frozen backbone is the teacher and retains the textual history. The student sees the current input and the same observed response prefix, while preceding history arrives as its memory state. At each supervised position, the teacher’s target is reduced to its top- tokens; the teacher distribution is renormalized on that selected support, while the student distribution remains normalized over the full vocabulary. The forward KL therefore penalizes student probability assigned outside the teacher’s selected tokens as well as mismatches within them. The objective can also include next-token cross-entropy on the observed target, with nonnegative weights; at least one of the KL and CE weights is nonzero. The paper specifies these operations in prose and notation, but the recipe here needs no additional formula.
The top-token target economizes on the teacher distribution without turning the student into a closed-vocabulary model. Yet broad next-token supervision may spend effort on tokens that are not the response the task needs. What changes when training focuses on answers?
Why Response Masks Focus Post-Training
Now the deeper training distinction: post-training carries historical context through memory, but applies its KL supervision to response positions rather than to the history or query.
Memory parameters start from the CPT checkpoint. For each instance, context is processed chronologically into memory; then the query and observed response prefix are supplied to the student, while the teacher retains the corresponding text history. Both are teacher-forced on the same observed prefix, and the teacher’s distributions—not the reference response tokens alone—provide the distillation targets. The response mask excludes history, query, and padding from direct token loss. This makes the response distribution the learning signal while the memory-building history still shapes the state. The unified context-query-answer format also accommodates heterogeneous memory-related tasks.
The distinction is important: reference responses determine the evaluated prefixes, while teacher predictions determine what the student is asked to match. How much of the teacher’s behavior can a compact memory reproduce when only those answer positions drive the objective?
Minimal Causal Code
This sketch shows the CPT state handoff. The storage operation is intentionally left as a named function: its internal update is a separate mechanism, not specified by this training recipe.
memory = zero_memory()
for segment in consecutive_segments(sequence):
logits = student(segment, memory)
loss = cross_entropy(logits, next_token_targets(segment))
loss.backward()
optimizer.step()
optimizer.zero_grad()
memory = store_segment(segment, memory)
The essential causal link is visible: a segment is predicted using the incoming memory, then processed into the state passed to the next segment.
What You Gain, and What You Lose
CPT teaches language modeling when prior text is no longer in the student’s context window; post-training aligns response predictions with a teacher that can still see the full history. Response masks concentrate direct distillation on answers, and top-token targets reduce the teacher distribution to selected support.
The tradeoff is that the student is trained against fixed teacher-generated trajectories, not its own rollouts. The teacher’s full textual history is a richer input than the student’s memory, so matching its distributions remains a learning target rather than a guarantee of identical behavior. And the memory still has to retain what later queries need: this training recipe supervises its usefulness indirectly through prediction and distillation.
Teach memory first to carry the past, then teach it to predict answers like a model that never had to forget the full history.