AK-MOMENTUM: AN ACTIVATION-AWARE GRADIENT MEMORY

QueryRetainWrite
t=8
‖x^t‖2=1
δt=

EMA MOMENTUM

The same decay in every direction
EMA momentum: memory readout across fixed input directionsThe axes never rotate. At each unit input direction q in the positive quadrant, the colored field's radius equals the stored response M q. Colors identify stored values after cancellation. The thin outer arc is the unit input scale.
u1
u2
x^t
Mt=βMt−1+(1−β)gt
Mean retention per eigenbasis entry
ρj=β

AK-MOMENTUM

Extra forgetting along the current key
AK-Momentum: memory readout across fixed input directionsThe axes never rotate. At each unit input direction q in the positive quadrant, the colored field's radius equals the stored response M q. Colors identify stored values after cancellation. The thin outer arc is the unit input scale.
u1
u2
x^t
Mt=Mt−1(βI−ηx^tx^t⊤)+ηδtx^t⊤
Mean retention per eigenbasis entry
ρj=β−ηλ^j
Each value has a color; memory carries it forward.Radius: stored responseβ=0.9,η=0.4
Timestep 8, write. Current unit input is 58.5 degrees from u1. EMA retains 0.9 in all directions. AK retains 0.5 along the key and 0.9 perpendicular to it. Color labels each value. The field radius gives the stored response in every direction of the illustrated positive input quadrant. One scalar output and two input coordinates. Positive unit inputs are sampled more often near u1. The axes stay fixed. Across every unit direction q in the positive input quadrant, the field radius is the stored response M q. This example has positive responses throughout that quadrant. The field is a response plot, not a covariance ellipse or a plot of eigendirections of the rectangular buffer. Matching colors denote the same illustrative scalar error. Opposite signed contributions of the same color cancel first. Cross-color cancellation is allocated proportionally for drawing; colors are explanatory labels, not a unique decomposition of a scalar sum. The full numerical buffer is preserved independently of this display. EMA retains 0.9 in every direction. AK retains 0.5 along the current key and 0.9 perpendicular to it.