🤖 AI

Learning to Remember, Learn, and Forget in Attention-Based Models

The paper introduces Palimpsa, a self-attention model that frames in-context learning as a continual learning problem using Bayesian metaplasticity to dynamically balance stability and plasticity, thereby overcoming the fixed capacity and interference limitations of gated linear attention models and significantly improving performance on long-sequence recall and reasoning tasks.

Original authors: Djohan Bonnet, Jamie Lohoff, Jan Finkbeiner, Elidona Shiqerukaj, Emre Neftci

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Djohan Bonnet, Jamie Lohoff, Jan Finkbeiner, Elidona Shiqerukaj, Emre Neftci

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 you are trying to learn a new language by reading a long story. As you read, you need to remember the names of characters and the plot points to understand what's happening next.

The Problem: The "Filing Cabinet" That Gets Full
Current AI models (like the ones powering chatbots) are great at this, but they have a flaw. To remember the story, they usually keep a giant "filing cabinet" of every single word they've read so far. The longer the story, the bigger the cabinet gets. Eventually, this becomes too heavy and slow to carry, especially on smaller devices.

To fix this, scientists created "linear" models. These models use a fixed-size notebook instead of a giant cabinet. They write new notes in the notebook, but if the notebook is full, they have to scribble over the old notes to make space. This causes a problem called catastrophic forgetting: the model accidentally erases important early details (like the main character's name) just to write down the latest sentence.

The Solution: A Smart Notebook Called "Palimpsa"
The authors of this paper propose a new model called Palimpsa. They treat the model's memory like a palimpsest—an ancient piece of parchment where scribes would scrape off old writing to reuse the paper. But instead of just scraping everything off, Palimpsa is smart about what it scrapes.

Here is how it works, using simple analogies:

1. The "Importance Tag" System

Imagine your notebook has a special tag on every page that says, "How important is this?"

Old Linear Models: They treat every page the same. If they need space, they just erase the oldest page, even if it contains the most crucial plot twist.
Palimpsa: It uses a system called Metaplasticity. This is like a "learning rate" for every single piece of information.
- If a fact is important (like the villain's name), the model puts a "Do Not Erase" sticker on it. It becomes very hard to change or overwrite.
- If a fact is stale or unimportant (like a random description of a tree that appeared 1,000 words ago), the model decides it's okay to let that information fade away.

2. Learning, Remembering, and Forgetting

The paper describes Palimpsa as a model that has learned three distinct skills:

Learning: It can absorb new information quickly when it arrives.
Remembering: It protects the "important" information so it doesn't get overwritten by new, less relevant data.
Forgetting: It actively discards "stale" information. This is crucial. If the model never forgot, it would eventually get so full of old, useless data that it couldn't learn anything new. This is called "catastrophic remembering." Palimpsa avoids this by letting go of the old stuff to make room for the new.

3. The "Mamba" Connection

The paper makes a fascinating discovery about a popular model called Mamba2.

Think of Mamba2 as a very fast, efficient runner who is great at short sprints but tends to drop things if the race gets too long.
The authors show that Mamba2 is actually a "special case" of Palimpsa where the "forgetting" switch is turned all the way up. It forgets so aggressively that it doesn't really learn to protect important memories.
The Upgrade: The authors figured out how to take a pre-trained Mamba2 model and "surgically" replace its brain parts with Palimpsa parts. This turns the fast runner into a marathon runner that can remember the start of the race even after 32,000 steps.

What Did They Prove?

The researchers tested this in three ways:

The "Memory Game" (MQAR): They gave the model a list of pairs (like "Key A = Value 1") and asked it to recall them later. Palimpsa was much better at remembering the right answers, especially when the list got very long, because it didn't accidentally overwrite the important keys.
Common Sense: They tested the model on tasks requiring logic and general knowledge (like "If I put a cup on a table, where is the cup?"). Palimpsa versions of the model scored higher than the standard versions, showing they could hold onto the necessary context to answer correctly.
Long Stories: When reading very long texts, Palimpsa didn't lose its ability to understand the story as quickly as the other models did. It could "reach back" further into the text to find the right answer.

The Bottom Line

Palimpsa is a new way for AI to manage its memory. Instead of just filling up a bucket until it spills over, Palimpsa acts like a smart librarian. It knows which books to keep on the shelf forever, which ones to keep for a little while, and which ones to throw away to make room for new arrivals. This allows the AI to handle longer, more complex tasks without getting confused or forgetting the beginning of the story.

Technical Summary: Learning to Remember, Learn, and Forget in Attention-Based Models

Problem Statement
In-Context Learning (ICL) in transformers functions as an online associative memory, enabling high performance on complex sequence tasks. However, gated linear attention models and state space models (SSMs) typically rely on fixed-size memory states to address the computational and memory bottlenecks of standard Transformers (which require linear and quadratic scaling for KV caches). This fixed capacity creates a "stability-plasticity dilemma": the model must learn new information without overwriting critical prior knowledge (catastrophic forgetting) or, conversely, becoming unable to integrate new information due to excessive retention of stale data (catastrophic remembering). Existing metaplasticity methods from neuroscience are often unsuitable for transformers because they rely on distinct task boundaries, growing memory, or lack differentiability. Furthermore, current linear attention models often suffer from limited memory capacity and interference as sequence lengths increase.

Methodology: Palimpsa
The authors propose Palimpsa, a self-attention model that reframes ICL as a continual learning problem solvable through Bayesian metaplasticity. The core insight is that the plasticity of each attention state should be dynamically adjusted based on an "importance state" grounded in a prior distribution that captures accumulated knowledge.

Bayesian Formulation: The attention head is modeled as a Bayesian inference problem where the goal is to minimize a variational free energy objective. The posterior distribution over the memory state $S$ is updated recursively. To prevent catastrophic remembering (where the prior becomes overly concentrated), the authors introduce a weighted posterior that effectively truncates the memory window. This is achieved by discarding a fraction of the accumulated prior weight at each step, controlled by a forgetting gate $\alpha_t = \exp(-A d_t)$ .
Metaplastic Update Rule: By applying Variational Inference (VI) with a diagonal approximation of the covariance matrix, the authors derive a closed-form update rule. The memory state $\mu_t$ $μ_{t}$ and an importance matrix $I_t$ $I_{t}$ (representing the inverse variance or precision of each state) are updated simultaneously:
- $I_t = \alpha_t I_{t-1} + (1 - \alpha_t)I_{prior} + \beta_t \otimes k_t^2$
- $\mu_t = \alpha_t \frac{I_{t-1}}{I_t} \odot \mu_{t-1} + \frac{1}{I_t} \odot [(\beta_t \odot v_t) \otimes k_t]$
  Here, $\beta_t$ acts as an input-dependent importance factor, and $I_t$ dynamically scales the learning rate for each state. High importance (low variance) leads to stability (low plasticity), while low importance allows for faster learning.
Architectural Variants: The authors implement two variants:
- Palimpsa-D: A drop-in replacement for Gated Delta Rule models (Deltanets), adding a projection for $\beta_t$ and a residual term.
- Palimpsa-M: Built upon Mamba2, where the input gate $d_t$ serves a dual role as both the forgetting gate and the input gate.
Theoretical Unification: The paper demonstrates that various gated linear attention models (Longhorn, Deltanets, Mesanet) emerge as specific architectural choices or posterior approximations within this common Bayesian framework. Notably, Mamba2 is shown to be a special case of Palimpsa where the forgetting rate is so high that the importance matrix remains constant ( $I_t \approx I_{prior}$ ), effectively disabling metaplasticity.

Key Contributions

Metaplasticity in ICL: The introduction of Bayesian metaplasticity to mitigate in-context catastrophic forgetting and remembering by dynamically adjusting state plasticity based on uncertainty.
Palimpsa Model: A scalable attention block that prevents catastrophic remembering by gradually releasing outdated knowledge, implemented with custom kernels for efficient GPU training.
Unified Framework: A mathematical derivation showing that existing gated linear attention and state space models are special cases of a common Bayesian framework, revealing Mamba2 as a limit of Palimpsa.
Transformation Pipeline: A method to transform non-metaplastic models (like pre-trained Mamba2) into metaplastic ones via fine-tuning, leveraging the speed of non-metaplastic kernels for pre-training and the memory capacity of metaplasticity for fine-tuning.

Experimental Results
The authors evaluate Palimpsa on synthetic, commonsense reasoning, and long-context benchmarks using 170M, 760M, and 2.7B parameter models.

Synthetic (MQAR): On the Multi-Query Associative Recall benchmark, Palimpsa consistently outperforms non-metaplastic baselines (including Gated Deltanet and Mamba2). The performance gap widens as sequence length increases, demonstrating that metaplasticity effectively mitigates overwriting in complex memorization tasks.
Language Modeling & Commonsense Reasoning: On the FineWeb-Edu dataset, metaplastic variants (Palimpsa-D and Palimpsa-M) achieve superior perplexity and accuracy on tasks like LAMBADA, PIQA, and WinoGrande compared to their non-metaplastic counterparts and standard Transformers. Specifically, at the 760M scale, fine-tuned Palimpsa-D outperforms the Gated Deltanet baseline by ~0.8 points in average accuracy.
Long-Context Performance: In experiments with a 2.7B parameter model, surgically replacing the 8 Mamba2 layers with the weakest forgetting (largest effective memory window) with Palimpsa blocks resulted in a +0.84 point improvement in commonsense reasoning and a +1.1 point improvement on the LongBench benchmark. While the metaplastic model underperformed slightly on short-context retrieval tasks, it demonstrated significantly slower degradation as context length increased, surpassing the baseline at 32k tokens.
Learning Dynamics: Analysis shows that the metaplasticity ratio (differentiation between plastic and consolidated synapses) increases during training, allowing the model to stabilize important information while remaining plastic to new data.

Significance and Claims
The paper claims that Palimpsa provides a principled solution to the stability-plasticity dilemma in fixed-size memory models by integrating concepts from neuroscience (metaplasticity) and Bayesian inference. The significance lies in:

Theoretical Insight: Unifying disparate linear attention architectures under a single Bayesian framework and identifying Mamba2 as a specific, forgetting-dominant limit.
Practical Scalability: Enabling the use of metaplastic models at scale by decoupling pre-training (using fast, non-metaplastic kernels) from fine-tuning (using metaplastic kernels to expand memory capacity).
Performance: Demonstrating that metaplasticity consistently improves memory retention and reasoning capabilities in long-context scenarios without requiring the quadratic memory costs of standard Transformers.

The authors conclude that this approach supports the deployment of performant sequence modeling on resource-constrained systems (edge devices) by optimizing how information is learned, remembered, and forgotten within strict memory bounds.