Orthoganality in Linear Attention Issue (Delta Rules!)
Image of Kimi Linear Attention (different than the attention described here).A Quick Note,
This is an explainer series to my daily X posts on "becoming a kernel engineer." I'll usually post daily. You can find them here: https://x.com/dj_sidequests.
I'll keep these "blog posts" mostly unrefined. That is, I will be publishing my notes here on my blog. Feel free to message me asking to clarify something!
The Problem + Brief Intro
1. Useful background intuition:
1) Information reduction. You reduced a 1million nxn, and n >> d, dxd matrix causes information collapse by compressing that much information. We must solve this issue somehow.
2) Orthogonality: the vectors that make up our matrix are perpendicular (90deg separation). We need to look at our key values. And, if 0 => then orthogonal (else not orthogonal).
3) One way to think about this: take the information reduction thing I just said. By compressing several distinct values into a smaller matrix, you lose a lot of information that differentiates the values (e.g. if they are not orthogonal, then your values “collapse” when you multiply it by the query matrix, resulting in you getting the incorrect answer)
2. Issue, 1) By multiplying KV first, you match K and V together. The issue: originally, you run softmax(QK^T) and allow V to adaptively “weigh in” (change the distribution) on the softmax(QK^T). We don’t do that anymore, so V isn’t really “weighing in” anymore.
2) This results in you contributing to the same dxd matrix instead of a NEW n x n matrix that turns into a n x d matrix.
Delta Rule Solution
Compare the old and new value of the key * value matrix = state_matrix * key. then,
State = State_old + beta * (key (value_vector - state_matrix*key))
the beta is a parameter you can optimize. this parameter will tell you how much of the new key-value matrix you should add.
note: we derive state_matrix * key by,
state_matrix = vector * key^T
=> state_matrix * key = vector (since we assume the values are orthogonal, then key^T * key = 1 for the keys of same index, but 0 for keys of differing index, allowing us to remove them from the calculation). What I said could probably didn't make sense, so here's the worked out math by ChatGPT,


Comments
Post a Comment