Enabling Robust In-Context Memory and Rapid Task Adaptation in Transformers with Hebbian and Gradient-Based Plasticity

TL;DR

This work investigates whether explicit synaptic plasticity can enable robust, intra-sequence adaptation in Transformer decoders by adding fast-weight components updated online through two biologically inspired rules: neuromodulated Hebbian plasticity and gradient-based plasticity. Across copying, cue–reward, few-shot regression, and one-shot image classification benchmarks, gradient-based plasticity excels at long-horizon credit assignment, while Hebbian plasticity shines under sparse supervision, with static weights often sufficient when information content is low. The authors provide a formal framework that couples a fast-weight memory with an outer-loop meta-trained backbone, revealing a functional taxonomy: dense feedback favors global gradient-like updates, whereas local associative updates are more effective for sparse, event-driven tasks. These findings illuminate how explicit plasticity can complement attention, offering interpretable adaptation dynamics and practical guidelines for deploying adaptive Transformers in real-world settings. The work also outlines directions to scale plasticity to larger models and to extend the approach to attention parameters, aiming to bridge emergent in-context learning with mechanistic, biologically inspired adaptation.

Abstract

Large language models display in-context learning as an emergent effect of scale, but they rely on static weights during inference. In contrast, biological systems continually adapt via synaptic plasticity. We investigate whether explicit, biologically inspired plasticity can endow Transformers with faster in-sequence adaptation. To this end, we augment decoder-only Transformers with fast-weight modules updated either by (i) a neuromodulated Hebbian rule or (ii) the gradient-based plasticity mechanism of Duan et al. (2023). Across copying, regression, and few-shot classification tasks (CIFAR-FS, Omniglot), Hebbian plasticity consistently achieves lower loss and stronger few-shot generalization, while gradient-based updates perform best on long-horizon credit assignment. When associations are short and linearly separable, static weights suffice, defining a clear boundary condition for when plasticity helps. Analysis of learned modulatory signals reveals that gradient-based rules maintain large, persistent updates, whereas Hebbian plasticity is sharply gated around salient events. Together, these results show that explicit plasticity complements attention by enabling rapid, task-specific adaptation, and clarify when different plasticity mechanisms are most effective.

Figures (2)

• Figure 1: Aggregate validation metrics across tasks. Top: Accuracy-style objectives (higher is better) for copying recall and the two classification datasets. Bottom: Loss-style objectives (lower is better) for cue–reward query loss and few-shot regression mean-squared error. Bars show the mean across three seeds; whiskers denote one standard deviation.
• Figure 2: Neuromodulation and plastic-weight norms per epoch for copying (left) and CIFAR-FS (right). Solid curves show the mean over three seeds; shaded regions denote one standard deviation. Gradient plasticity sustains high $\eta(t)$ throughout the episode, whereas Hebbian updates fire in short bursts around support tokens.