Cortex: Consolidating Context into Parametric Memory
Jintian Zhang, Feng He, Xinnian Liang, Ziwen Xu, Haoming Xu, Haojing Huang, Shuangzhi Wu, Ningyu Zhang
Cortex consolidates context-induced behavioral gains into parametric memory, enabling persistent performance without source context.
How can we reliably convert transient, context-dependent model behaviors into permanent, internal model parameters without degrading general capabilities?
Large language models can exploit new information provided in context, but these performance gains vanish as soon as the supporting context is removed. Cortex internalizes these gains by generating prospective queries to anticipate future needs and distilling context-conditioned outputs into compact, mergeable feed-forward network updates. Across diverse backbones, this approach recovers nearly 90% of context-induced gains while retaining over 99% of original non-target capabilities after 50 sequential updates.
Paper Primer
The paper formalizes Memory Consolidation (MC) as a unified objective: given a model and a context that improves performance, the goal is to update the model so it retains those gains without the context, while preserving its original capabilities. This addresses the "memory failure" where models perform well only when provided with external documentation or instructions, but lose that utility immediately upon removal.
Cortex implements this through a three-part mechanism: prospective rehearsal, FFN-based plasticity, and context-to-memory distillation. It generates prospective queries to simulate future use cases, then trains either low-rank synaptic updates or structural nonlinear channel expansions to align the student model's output distribution and hidden states with the context-conditioned teacher.
Cortex achieves high consolidation rates while maintaining near-perfect retention of original model capabilities.
In a 50-step sequential factual editing task on a 14B model, Cortex-Structural recovered 88.9% of context-induced gains (CR) while retaining 99.8% of non-target performance (RR). Outperforms AlphaEdit by 66.2 percentage points in consolidation rate.
Structural expansion is more resilient to interference than synaptic updates in continual settings.
Without query replay, Structural retained 40.5% of consolidated knowledge after 25 subsequent edits, compared to 8.0% for Synaptic updates. Structural expansion provides a more stable substrate for long-term memory accumulation.
Why is this approach better than simply fine-tuning the model on the context?
Standard fine-tuning often leads to catastrophic forgetting of original capabilities. Cortex explicitly optimizes for both consolidation and retention, using hidden-state alignment and query replay to ensure that new knowledge is integrated without disrupting existing parametric memory.
What is the scope of this method — does it only work for factual updates?
The testbed covers five task families: agent experience, novel API learning, persona steering, knowledge suppression, and factual editing. It is designed to handle declarative, procedural, and behavioral changes across both dense and mixture-of-experts backbones.
Cortex demonstrates that context-induced behavioral gains can be reliably "baked" into model parameters, shifting the paradigm from transient in-context prompting to persistent, cumulative parametric memory.
Abstract
Identifies the loss of context gains in LLMs and frames the need for memory consolidation.
Large language models can leverage surrounding context to boost performance, yet these improvements vanish once the context is removed. This transient behavior motivates the need for a mechanism that internalizes context‑induced gains while preserving the model’s original capabilities.
The Challenge of Contextual Memory
We define memory consolidation, its motivation, and the proposed Cortex approach.
Experience can be reused beyond the situation where it is acquired, so agents need mechanisms that let past experience keep shaping future behavior. The complementary learning systems (CLS) framework explains how rapid hippocampal encoding and slower neocortical consolidation achieve this, as illustrated in Figure 1.
**Figure 1. Biological inspiration and the memory consolidation objective.** (a) Hippocampal–neocortical interactions support integration of newly acquired information into long-term cortical memory. (b) MC aims to transform a context-conditioned model $M$ into an updated model $M'$ that retains the context-induced gains without the original context $C$, while preserving non-target capabilities. The highlighted FFN region represents the learned changes in CORTEX.
Memory Consolidation is the process of turning a context‑induced performance boost into a permanent change to the model’s parameters.
We measure success with two metrics: Consolidation Rate (CR), the fraction of the context‑induced gain recovered after $C$ is removed, and Retention Rate (RR), the proportion of the original performance retained on unrelated queries. These metrics are evaluated on the Memory Consolidation Testbed, which covers declarative knowledge, procedural knowledge, and behavioral control across five task families.
Cortex implements MC through three steps: (1) generate prospective queries that anticipate where the context will be useful; (2) apply low‑rank updates to feed‑forward networks (FFNs) – either Synaptic (additive) or Structural (non‑linear channel expansion) – with a shared expansion for pretrained MoE models; and (3) train the new parameters with a context‑conditioned teacher while freezing the pre‑update model to limit interference.
The central gap is that transient context yields useful behavior, but without consolidation those gains disappear once the context is gone.
Formalizing Memory Consolidation
Formulating the performance gap and consolidation objectives the paper seeks to address.
We consider a language model $M$ that receives a temporary context $C=\{c_1,\dots,c_n\}$ and is queried with $q$ to produce a response $y$. Performance is measured by a scoring function $r$ over the model’s output distribution.
For each context $C$ we split queries into two sets: $Q^{+}_{C}$, whose answers should change because of $C$, and $Q^{-}_{C}$, whose answers should stay the same. For example, if $C$ encodes “A is B’s father”, then “Who is B’s father?” belongs to $Q^{+}_{C}$ while “What is A’s name?” belongs to $Q^{-}_{C}$.
When $\Delta_{\text{ctx}}$ is substantially positive we call $C$ non‑trivial; the model exploits $C$ but loses the benefit once $C$ is removed, indicating a memory failure. If the context‑free performance is already low, a near‑zero gap signals a capability failure rather than a memory issue.
By definition, CR = 1 means the updated model fully recovers the original context‑induced gain; RR = 1 means it perfectly preserves the baseline retain score. Either metric can exceed one if the update overshoots the original performance.
We consider three consolidation protocols. In the single‑context setting ($n=1$) a lone context is distilled into $M'$. In batch consolidation ($n>1$) multiple contexts are processed together. In the sequential protocol, contexts arrive over time; each new $C_t$ updates $M_{t-1}$ to $M_t$, and past contexts cannot be revisited.
The Cortex Consolidation Mechanism
Cortex distills context-induced behaviors into persistent FFN parameters using prospective rehearsal and context-to-memory distillation.
Cortex consolidates context-dependent behaviors into model parameters by distilling the frozen teacher's context-aware behavior into a student's newly introduced FFN parameters. This process targets both the recovery of context-induced gains and the preservation of non-target capabilities.
Instead of merely memorizing past inputs, the model generates "prospective queries" that anticipate how the context might be used in future tasks, creating a synthetic dataset that captures the context's functional value.
Generate $q_1 = \text{"What is the weather in London?"}$ and $y_1 = \text{"search(London)"}$.
Generate $q_2 = \text{"How do I check the time?"}$ and $y_2 = \text{"I cannot help with that."}$.
Filter $F$ retains $(C, q_1, y_1)$ as a successful tool-use example, discarding the irrelevant $q_2$.
Prospective rehearsal focuses the model's limited consolidation capacity on functional tool-use behaviors rather than the entire raw context.
FFNs act as the model's key-value memory; Cortex targets these layers for plasticity, either by updating existing projections or by adding new nonlinear channels to expand capacity.
The student model learns to mimic the teacher's context-aware behavior by minimizing the divergence between their output distributions and internal representations, effectively "distilling" the context into the student's parametric memory.
**Figure 2** Rehearsal and consolidation workflow of Cortex. (a) Rehearsal construction for a single context (1), concatenated contexts (2), or separately processed contexts whose examples are subsequently merged (3). (b) Single, batch, and sequential consolidation. In the sequential setting, blue modules denote earlier writes and the red module denotes the current write. (c) Context-to-memory distillation: the frozen teacher receives the source context, while the student omits it and trains only the newly introduced FFN or MoE parameters. Output-distribution KL and MSE on aligned states after the FFN residual update supervise the transfer of contextual behavior and representations into parametric memory.
Evaluating Memory Consolidation
The testbed defines five tasks and three consolidation settings to evaluate memory consolidation and retention.
The testbed defines three consolidation settings that cover distinct update regimes.
Single consolidation processes one context at a time; batch consolidation merges compatible contexts; sequential consolidation incorporates arriving contexts through successive updates.
A suite of five tasks that probes how well a model can retain context‑induced gains (consolidation) while preserving unrelated abilities (retention).
How does this testbed differ from prior memory‑benchmark suites?
Prior suites typically evaluate either retention or forgetting in isolation. Here, each task jointly reports consolidation rate (how much of the context gain is recovered) and retention rate (how well unrelated abilities survive), and the testbed explicitly varies the consolidation protocol (single, batch, sequential) to stress different update regimes.
**Figure 4** Task construction in the Memory Consolidation Testbed. Panels (a–e) illustrate agent experience consolidation, novel API learning, persona steering, knowledge suppression, and factual editing. Source information is organized into a context $C$, with target queries $Q_C^+$ and, where applicable, retention queries or criteria $Q_C^-$. Colors distinguish memory types; task headers indicate the intended operation, scope, and supported consolidation settings. These settings describe protocol support, rather than complete experimental coverage of every task–setting combination. In panel (c), the illustrated adherence question and fluency criterion refer to Power-seeking; Psychopathy uses behavior-matching accuracy and has no reported RR.
Experimental Results
Key quantitative outcomes of Cortex across single, batch, and continual settings.
Cortex converts contextual gains into persistent model capabilities, extending both acquisition and retention beyond what plain prompting can achieve.
Cortex retains 88.9 % consolidation rate after 50 successive edits, far surpassing the 28.5 % achieved by the strongest baseline (ROME).
Structural Cortex reaches 88.9 % CR and 99.8 % RR; Synaptic reaches 80.6 % CR and 97.2 % RR, while ROME stalls at 28.5 % CR with comparable RR.
**Figure 5.** Single and batch consolidation across tasks and backbones. Each panel compares consolidation and retention for one task-backbone setting; labels give (CR, RR). Insets enlarge crowded regions, and bottom strips report CR for methods without a complete RR measurement. RR uses fluency for steering and locality for single editing; batch editing additionally includes general-capability evaluations. Background colors group task families, and extreme negative CR values are compressed for display.
**Figure 6. Continual consolidation over 50 factual updates.** (a) Final editing and general-capability accuracies; heatmap colors are normalized within each metric. (b) Final CR–RR trade-offs and parameters modified or added per edit; the hatched WISE bar is a matrix-size estimate. (c) CR for the fixed first ten facts, aligned by subsequent edits (top): at age $u$, fact $j$ is evaluated at checkpoint $j + u$. CR combines efficacy, generalization, and portability; the band shows the sample standard deviation of the three run-level means. The bottom comparison averages over all 50 facts, evaluated immediately after their own edits or at checkpoint 50. (d) Final-checkpoint repetition-triggered stopping and length-limit termination within a 512-token observation window, using 700 responses per run from original and rephrased questions; the stacked categories are mutually exclusive, and dots show individual-run totals. Other error bars and numerical $\pm$ values also denote sample standard deviations across three runs.
**Figure 7.** Design choices governing consolidation and retention. (a,b) LoRA controls trained on Cortex data. (c,d) FFN-only versus FFN-and-attention updates. (e) Mean original-question and rephrasing accuracy for the first 25 facts, aligned by subsequent edits; WISE is an external reference. (f,g) Hidden-state MSE. (h,i) MoE sharing at matched total or per-token active expansion-weight budgets, excluding biases. (j,k) Layer-median channel-update CV over training windows. (l) Expert exposure divided by its layer's mean within each seed, then averaged across seeds; circles mark branches below 1% of their layer mean in every seed. (m,n) Continual CR and RR against added parameters per edit, in millions. Small and large markers denote training seeds and their means, except in (a,b); bands show $\pm 1$ sample standard deviation. Q2.5 denotes Qwen2.5-14B-Instruct, Q3-4B-I denotes Qwen3-4B-Instruct, and OSS denotes GPT-OSS.
Related Work
We position memory consolidation among related adaptation and memory approaches.
Memory consolidation (MC) enables models to internalize context‑induced behaviors for future use without the original context. We relate MC to neighboring problems such as retrieval‑augmented generation, task‑specific editing, and test‑time training, and then compare Cortex with existing approaches.
FT-M updates all model parameters on a small dataset to imprint a new behavior, then freezes the model for downstream use.
These methods locate the MLP layer responsible for a fact and apply a low‑rank update that forces the desired answer while minimally disturbing other weights.
WISE shards the parameter space into random subspaces and merges edits by routing activations, reducing interference between successive edits.
**Figure 8.** Fact-level trajectories of the continual baselines. Rows show FT-M, ROME, MEMIT, AlphaEdit, and WISE. Columns show efficacy, generalization, portability, and locality. Within each panel, rows are update indices and columns are fact indices; cells average raw accuracy over three seeds on a common 0–100% scale. Gray cells correspond to facts not yet edited.
Conclusion
We recap our consolidation framework, testbed results, and outline promising future research avenues.
Memory Consolidation was formalized as internalizing context‑induced gains while preserving other capabilities, and a five‑task testbed was introduced to measure both consolidation and retention. Across scales and architectures, Cortex consistently occupies the empirical consolidation–retention Pareto frontier, and future work will explore cross‑task continual consolidation, large‑batch regimes, pruning‑merge of structural expansions, and hypernetworks that predict structural parameters to amortize consolidation cost.
Technical Details and Ablations
A FFN‑local context correction and shared MoE expansion We study how a fixed FFN update can represent the effect of a context across multiple computation states. Throughout, parameters are fixed during inference. Goldwaser et al. [29] establish input‑dependent parameter updates that recover context‑conditioned outputs under architectural controllability conditions. Cortex instead learns a shared functional correction through context‑to‑memory distillation. The results below characterize its representation on finite‑state sets and the routing dependence of its shared MoE expansion.
A.1 Context correction and hidden‑state supervision Fix a finite collection of aligned teacher–student computation states. For a block whose FFN output is added directly to the residual stream, write its student output before the current update as
$$ \bar h_i = v_i + F(u_i),\qquad u_i = N(v_i), $$
where $v_i$ is the attention sublayer’s residual output and $N$ is the fixed pre‑FFN normalization. Let $h_i^{T}$ be the corresponding teacher target, and set
$$ \delta_i = h_i^{T} - \bar h_i,\qquad \Delta = [\delta_1,\dots,\delta_n]. \tag{9} $$
The correction includes the entire discrepancy in the block output, including the residual difference induced by attention. An additive FFN update $R$ matches these targets exactly when $R(u_i)=\delta_i$ for every $i$.
This view directly motivates hidden‑state MSE. At a fixed layer, retain the current upstream updates and disable only its added branch when defining $\bar h_i$. Then
$$ \|h_i^{S}-h_i^{T}\|_2^{2}= \|R(u_i)-\delta_i\|_2^{2}. \tag{10} $$
Thus intermediate‑state alignment supervises the contextual effects to be consolidated, while output KL supervises the resulting predictive behavior. During joint training, the inputs and required corrections also vary with upstream parameters.
For a block with post‑FFN normalization $P$, the corresponding expressions are
$$ h_i^{S}= v_i + P\bigl(F(u_i)+R(u_i)\bigr),\qquad \bar h_i = v_i + P\bigl(F(u_i)\bigr), \tag{11} $$
$$ \operatorname{Reff}(u_i)=P\bigl(F(u_i)+R(u_i)\bigr)-P\bigl(F(u_i)\bigr). \tag{12} $$
With $\delta_i = h_i^{T}-\bar h_i$, the same loss identity holds with $\operatorname{Reff}$ in place of $R$. Alignment therefore supervises the correction after the architecture’s normalization. The linear least‑squares result below concerns the raw branch output; it is not a closed‑form solution to the post‑normalized loss. For the direct‑addition case $P=I$, the two corrections coincide.
For structural plasticity,
$$ R(u)=D_s\phi_s(u),\qquad \phi_s(u) = (U_s u)\odot \phi(G_s u), \tag{13} $$
where $\phi$ is the backbone’s channel activation, with $U_s, G_s\in\mathbb{R}^{k\times d_h}$ and $D_s\in\mathbb{R}^{d_h\times k}$. These parameters implement a width expansion
$$ W'_U = \begin{bmatrix}W_U\\U_s\end{bmatrix},\quad W'_G = \begin{bmatrix}W_G\\G_s\end{bmatrix},\quad W'_D = \begin{bmatrix}W_D\\D_s\end{bmatrix}. \tag{14} $$
Consequently $F'(u)=F(u)+R(u)$, with the original weight blocks unchanged. Any subsequent normalization receives the same input before and after folding, so it preserves this implementation equivalence.
For biased affine projections, concatenate the gate and up biases along with their new rows, and add the two output biases: $b'_D = b_D + b_{D,s}$. The same identity holds for any shared channel‑wise gate‑up operator, including GPT‑OSS’s clipped, scaled gate and shifted up activation. For a fixed‑feature fit with an output bias, use the augmented feature matrix $[\Phi^\top\;1]^\top$.
**Proposition A.1 (Shared correction with fixed features).** Fix $U_s, G_s$ and any raw branch targets $\Delta$; for a direct‑addition block these are the corrections in (9). Let $\Phi = [\phi_s(u_1),\dots,\phi_s(u_n)]$. There exists a single $D_s$ such that $D_s\Phi = \Delta$ iff
$$ \operatorname{row}(\Delta)\subseteq \operatorname{row}(\Phi). \tag{15} $$
In general
$$ \min_{D_s}\|D_s\Phi-\Delta\|_F^{2}= \|\Delta(I-\Phi^{\dagger}\Phi)\|_F^{2}, \qquad D_s = \Delta\Phi^{\dagger} \tag{16} $$
attains this minimum.
*Proof.* Every row of $D_s\Phi$ lies in $\operatorname{row}(\Phi)$. Since $P=\Phi^{\dagger}\Phi$ is the orthogonal projector onto that space, $\Delta P$ is the closest such matrix to $\Delta$. It is attained by $D_s=\Delta\Phi^{\dagger}$, and equals $\Delta$ precisely under (15).
A.2 Finite‑state realizability of structural plasticity The feature map in (13) is itself trainable. The next result gives a sufficient condition under which its features can separate an arbitrary finite set, without requiring linearly independent inputs.
**Theorem A.2 (Finite‑set interpolation by a structural expansion).** Let $u_1,\dots,u_n\in\mathbb{R}^{d_h}$ be non‑zero and satisfy $u_i\neq u_j$ and $u_i\neq -u_j$ whenever $i\neq j$. Let
$$ \phi(s)=s\,\sigma(\eta s + \zeta s^{3}),\qquad \sigma(s)=\frac{1}{1+e^{-s}}, $$
with $\eta>0,\;\zeta\ge 0$. For any $\delta_1,\dots,\delta_n\in\mathbb{R}^{d_h}$ and any $k\ge n$, there exist fixed parameters in (13) such that $R(u_i)=\delta_i$ for all $i$. The activation family includes SiLU ($\eta=1,\;\zeta=0$) and tanh‑approximate GELU ($\eta=2\sqrt{2/\pi},\;\zeta=0.044715\eta$).
*Proof sketch.* Choose a vector $a$ outside the finitely many hyperplanes orthogonal to the $u_i$ and $u_i\pm u_j$. Define $z_i = a^\top u_i$; these are non‑zero with distinct absolute values. For $t>0$ set
$$ f_i(t)=z_i\,\phi(tz_i)=t z_i^{2}\,\sigma(\eta t z_i + \zeta t^{3}z_i^{3}). $$
The functions $f_1,\dots,f_n$ are linearly independent; consequently the vectors $(f_1(t),\dots,f_n(t))$ span $\mathbb{R}^n$ as $t$ varies. Selecting $t_1,\dots,t_n>0$ so that the matrix $\Phi_{ji}=f_i(t_j)$ is invertible, and setting the rows of $U_s$ to $a^\top$ and the rows of $G_s$ to $t_j a^\top$ yields the feature matrix $\Phi$. Then $D_s = \Delta\Phi^{-1}$ interpolates all targets. For $k>n$ extra output columns can be set to zero.
The theorem provides a finite‑set capacity guarantee. For a smaller chosen width, Proposition A.1 characterizes the correction represented by the learned features. Repeated inputs with identical targets can be merged before applying the theorem. Its hypotheses concern the inputs actually read by the expansion, after any pre‑FFN normalization.
*Native GPT‑OSS channels.* With the biases and clipped activation used by GPT‑OSS, a finite‑set construction only requires pairwise distinct inputs. Write the native channel operator as
$$ \Psi_{\alpha,L}(g,u)=\bigl(\operatorname{clip}_{[-L,L]}(u)+1\bigr)\odot \bigl(\min(g,L)\odot\sigma(\alpha\min(g,L))\bigr), \qquad \alpha,L>0, $$
with element‑wise clipping and multiplication. The branch is
$$ S(x)=D_s\Psi_{\alpha,L}(G_s x + b_{G,s},\,U_s x + b_{U,s}) + b_{D,s}. $$
Choosing $a$ so that the scalars $z_i=a^\top u_i$ are distinct, ordering them increasingly, and selecting thresholds $\tau_1<z_1$ and $\tau_j\in(z_{j-1},z_j)$ for $j\ge2$, one can make the feature matrix approach an upper‑triangular matrix with non‑zero diagonal, guaranteeing invertibility for some finite $\lambda$. Hence arbitrary targets can be interpolated using $n$ new channels.
**Corollary A.3 (FFN‑local matching of model outputs).** Consider a deterministic Transformer whose final FFN output is added directly to the residual stream. Freeze all existing parameters and fix a finite set of aligned prediction states. If the student inputs to that FFN satisfy the hypotheses of Theorem A.2, a structural expansion at that FFN can match arbitrary teacher final residual states on this set. If teacher and student share the final normalization and output head, their logits and next‑token distributions then agree on the set.
The same construction can match intermediate states by proceeding from lower to upper layers: after earlier updates have been fixed, recompute the current student inputs and the corrections in (9), and apply the theorem whenever its hypotheses hold. The corrections compensate for the teacher’s continued access to context at every layer. This provides a representation rationale for FFN‑local hidden‑state distillation; performance on unseen states depends on the learned approximation and rehearsal coverage.
Injection point. The model‑output corollary concerns direct residual addition. For a block of the form $v + N_{\text{post}}(F(u))$, a target $h_i^{T}$ instead requires $h_i^{T}-v\in\operatorname{range}(N_{\text{post}})$. If these range conditions hold for all states and their FFN inputs satisfy Theorem A.2, choose preimages $w_i$ and interpolate the pre‑normalization corrections $w_i - F(u_i)$. The fixed‑feature proposition and the MoE identities below apply to the corresponding module outputs regardless of this subsequent normalization.
For $N_{\varepsilon,\gamma}(w)=\gamma\odot w/\sqrt{\|w\|_2^{2}/d_h}+\varepsilon$ with $\varepsilon>0$ and non‑zero scale coordinates, write $q_i = (h_i^{T}-v_i)\odot$.
Section: Proof Necessity follows from $\Delta$ = LH. For sufficiency take L = $\Delta$H†, which satisfies LH = $\Delta$ and has rank at most rank($\Delta$). A rank factorization implements L = BA.