Adaptive Memory Decay for Log-Linear Attention

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Amin, Yaxita, Li, Helen Zichen, Zhang, Mengfan, Ayhan, Samet
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909025457668096
author Amin, Yaxita
Li, Helen Zichen
Zhang, Mengfan
Ayhan, Samet
author_facet Amin, Yaxita
Li, Helen Zichen
Zhang, Mengfan
Ayhan, Samet
contents 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 λ 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 λ 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 λ degrades or collapses entirely.
format Preprint
id arxiv_https___arxiv_org_abs_2605_06946
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Adaptive Memory Decay for Log-Linear Attention
Amin, Yaxita
Li, Helen Zichen
Zhang, Mengfan
Ayhan, Samet
Machine Learning
Artificial Intelligence
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 λ 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 λ 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 λ degrades or collapses entirely.
title Adaptive Memory Decay for Log-Linear Attention
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2605.06946