Skip to content

v-code01/popcountmem

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

3 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

popcountmem

Log-Linear Attention stores its history in a Fenwick tree and reads it back through one scalar per level. The number of independently addressable memory slots at query position t is not the log(t) the paper implies but exactly the popcount of t, each slot is a rank-d linear-attention state, and the newest token's memory periodically collapses to plain linear-attention resolution at positions gated by the head dimension. The recency bias the paper describes is real within a query but not absolute: it is carry-gated by the binary structure of t.

The Fenwick state and the popcount slot count

Log-Linear Attention (Guo, Yang, Goel, Xing, Dao, Kim, arXiv:2506.04761) replaces linear attention's fixed state with a hierarchical set of states over the Fenwick decomposition of the prefix, read out as o_t = sum_l lambda_t^(l) q_t^T S_t^(l), where each S_t^(l) = sum_{s in B_t^(l)} k_s v_s^T aggregates one dyadic bucket and lambda_t^(l) is a data-dependent scalar. The Fenwick decomposition peels the least significant set bit of t, so the bucket sizes are the set bits of t. The count of nonempty buckets, the number of independently addressable memory slots, is therefore

slots(t) = popcount(t) ,

not floor(log2 t) + 1. The paper writes O(log t); that is only the maximum. The exact count swings between 1 at every power of two and floor(log2 t) + 1 just below one. Secondary summaries that quote floor(log2 t) + 1 are stating the bound, not the value.

Each slot is a rank-d state

A bucket state S = sum_s k_s v_s^T of m tokens is a sum of m rank-one terms in d x d, so its rank is min(m, d). The readout o = q^T S = sum_s (q . k_s) v_s lets a query separate tokens inside a bucket through the key alignment, so the resolution is not 1/m. A specific value v_j is recoverable exactly when some q has q . k_s = delta_{sj}, which exists precisely when the keys are independent, that is when the bucket size m <= d:

bucket size 8 32 (= d) 33 64 256
retrieval error 2e-16 8e-13 0.14 0.70 0.94

Below d a bucket keeps its tokens individually; above d it is a rank-d sum from which no query recovers an individual value. The per-bucket scalar lambda cannot change this: it multiplies the already-summed state (a query rescaled by 1/lambda undoes it exactly), so it cannot re-weight tokens the sum has merged.

The power-of-two recency collapse

The newest token sits in a bucket of size lssb(t), the largest power of two dividing t. It is individually recoverable only while lssb(t) <= d, so it collapses into a rank-d sum exactly when t is divisible by the smallest power of two greater than d. For d = 32 that is every multiple of 64:

t 63 64 65 127 128
newest bucket size 1 64 1 1 128
newest token resolves yes no yes yes no

At t = 64, 128, 192, ... the most recent token has the coarsest possible memory, identical to plain linear attention, even though one step earlier and one step later it is stored at full resolution. The collapse period is the smallest power of two exceeding d, so it tracks the head dimension: d = 16 collapses every 32, d = 64 every 128.

What the paper claims and what holds

The paper states an inductive bias: "recent tokens are allocated more fine-grained memory, while distant tokens are compressed more aggressively," and "more recent tokens are retained at high resolution." Within a fixed query that ordering does hold, because the bucket sizes increase from the most recent to the oldest (they are distinct increasing powers of two). What fails is the absolute reading, that a recent token always has fine memory. Its resolution is gated by the carry structure of t and the head dimension, and at the collapse positions the newest token is stored no better than by the linear attention that Log-Linear Attention aims to improve on. This is a qualification of the claim with an exact, periodic, catastrophic-violation regime, not a blanket refutation.

Layout

  • fenwick.py: the Fenwick buckets, the exact slots(t) = popcount(t) slot count, and the newest-bucket = lssb(t) law.
  • capacity.py: the rank-d bucket capacity, the retrieval error crossing at bucket size d, and the proof that lambda cannot re-weight inside a bucket.
  • collapse.py: the newest-token collapse condition lssb(t) > d, the collapse period as the smallest power of two exceeding d, and the slot-count sawtooth.
  • test_popcountmem.py: the popcount slot count, the set-bit bucket sizes, the recoverability crossing at d, the lambda invariance, the collapse condition and period, and the sawtooth.

Reproduce

python fenwick.py
python capacity.py
python collapse.py
python test_popcountmem.py

The retrieval measurements are float64 on CPU; the M4 relevance is that a Log-Linear head at dimension d running on device sees its newest-token retrieval fidelity dip to the linear-attention floor at exactly the multiples of the smallest power of two above d, a sawtooth in t set by the head size.

Note

That Log-Linear Attention uses a Fenwick tree with O(log t) state is stated in the paper. What is here is the exact structure the O(log t) hides: the slot count is exactly popcount(t) (not the floor(log2 t) + 1 bound the paper and its summaries quote), each slot is a rank-d linear-attention state so a bucket resolves its tokens only up to size d, and the newest token's resolution collapses to the linear-attention floor periodically, at every multiple of the smallest power of two greater than d, gated by the binary carry structure of t. The paper frames the power-of-two case as the hierarchy gaining a level; the readout view is that it loses all but one, coarse, level. This qualifies the paper's monotone-recency claim rather than refuting it: the ordering holds within a query, but a recent token does not always have fine memory. The extension is whether the learned per-bucket lambda schedule concentrates on small recent buckets in a way that empirically masks the collapse in a trained model, or whether the periodic dip shows up as a position-dependent retrieval-accuracy sawtooth.

About

Log-Linear Attention's memory slots at query t are exactly popcount(t), not the log(t) the paper implies; each slot is a rank-d linear-attention state, so a token resolves only while its Fenwick bucket size <= d. The newest token's resolution collapses to the linear-attention floor periodically, at multiples of the smallest power of two above d.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages