Learning to Remember, Learn, and Forget in Attention-Based Models
Abstract
In-Context Learning (ICL) in transformers acts as an online associative memory and is believed to underpin their high performance on complex sequence processing tasks. However, in gated linear attention models, this memory has a fixed capacity and is prone to interference, especially for long sequences. We propose Palimpsa, a self-attention model that views ICL as a continual learning problem that must address a stability-plasticity dilemma. Palimpsa uses Bayesian metaplasticity, where the plasticity of each attention state is tied to an importance state grounded by a prior distribution that captures accumulated knowledge. We demonstrate that various gated linear attention models emerge as specific architecture choices and posterior approximations, and that Mamba2 is a special case of Palimpsa where forgetting dominates. This theoretical link enables the transformation of any non-metaplastic model into a metaplastic one, significantly expanding its memory capacity. Our experiments show that Palimpsa consistently outperforms baselines on the Multi-Query Associative Recall (MQAR) benchmark and on Commonsense Reasoning tasks.
Lay Summary
When artificial intelligence models read long documents or hold extended conversations, they rely on a temporary working memory to keep track of the context. To operate efficiently, newer AI models restrict this memory to a fixed size. However, this creates a strict dilemma: to learn something new, the model must overwrite something old. If it overwrites too freely, it suffers from "catastrophic forgetting," losing critical past details. If it becomes too rigid, it completely stops learning. We solve this problem by introducing Palimpsa, an AI architecture inspired by how the brain dynamically adapts its neural connections. Instead of treating all memories equally, Palimpsa uses a probability-based approach to track the "importance" of every piece of information it has encountered. By actively protecting crucial knowledge from being erased—while allowing stale details to safely fade away—Palimpsa vastly expands the effective capacity of fixed-size AI memories. Our experiments demonstrate that this approach significantly improves the model's ability to reason and reliably recall information over long sequences.