Learning to Remember, Learn, and Forget in Attention-Based Models
This addresses memory limitations in attention-based models for sequence processing, offering a theoretical and practical improvement, though it builds incrementally on existing gated linear attention methods.
The paper tackled the problem of fixed memory capacity and interference in gated linear attention models for in-context learning by proposing Palimpsa, a self-attention model that treats it as a continual learning problem, resulting in consistent outperformance on benchmarks like MQAR and Commonsense Reasoning tasks.
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.