Adaptive Memory Decay for Log-Linear Attention

TL;DR

This paper introduces a method to adaptively learn memory decay parameters in log-linear attention models, enhancing their ability to recall relevant information in long sequences.

Contribution

It proposes a lightweight, input-dependent decay mechanism for log-linear attention that improves long-range memory recall without increasing complexity.

Findings

Input-dependent decay outperforms fixed decay in associative recall and language modeling.

Largest improvements occur in long-range memory tasks where fixed decay fails.

The method preserves log-linear complexity with negligible parameter overhead.

Abstract

Sequence models face a fundamental tradeoff between memory capacity and computational efficiency. Transformers achieve expressive context modeling at quadratic cost, while linear attention and state-space models run in linear time by compressing context into a fixed-size hidden state, inherently limiting recall. Log-linear attention navigates this tradeoff by organizing memory across a Fenwick tree hierarchy, growing its hidden state logarithmically with sequence length at log-linear compute cost. However, its memory decay parameter {\lambda} is fixed and independent of the input, assigning uniform weights across all hierarchy levels regardless of the content, which introduces unnecessary rigidity. We propose learning {\lambda} directly from the input via a lightweight two-layer MLP, producing per-token, per-level decay that adapts to content rather than position. A softplus activation lets each Fenwick tree level scale independently, avoiding the inter-level competition that softmax introduces. This modification preserves log-linear complexity exactly and adds negligible parameter overhead. We evaluate on associative recall, selective copying, and language modeling, finding that input-dependent decay consistently outperforms the baseline, with the largest gains in long-range memory settings where baseline {\lambda} degrades or collapses entirely.