←back to thread

Multi-Token Attention

(arxiv.org)
152 points fzliu | 4 comments | 02 Apr 25 22:20 UTC | HN request time: 0.224s | source
Show context
kouteiheika ◴[03 Apr 25 03:56 UTC] No.43564572[source]▶
This is another potential improvement to the transformer architecture from Facebook (the other one that comes to mind is this one from same authors: https://arxiv.org/abs/2405.18719), but note that it comes with a major problem that might not be obvious at first glance: it's just not usable in practice without a ton of work. It modifies the innards of the attention mechanism, so it is incompatible with Flash Attention (or any other optimized attention library), and you do not want to train anything beyond toy models without Flash Attention (the performance hit is just way too big).

There's pytorch's FlexAttention which could maybe make this practical, but currently it's just way too buggy.

replies(2): >>43564707 #>>43567487 #
albertzeyer ◴[03 Apr 25 10:22 UTC] No.43567487[source]▶
Why do you say FlexAttention is too buggy? I have heard about a lot of successful usages of it, and never heard about any such problems.

Also note, depending on your model dimensions and sequence lengths, often the attention computation plays only a minor role (maybe 10% overall or so), and the MLP computation dominates.

replies(1): >>43570474 #
1. kouteiheika ◴[03 Apr 25 14:44 UTC] No.43570474[source]▶
Last time I tried it I encountered both showstopper bugs (it was completely obviously broken) and subtle correctness bugs (it looked like it was working, but since I'm paranoid I have unit tests for everything and numerically the errors were too big compared to what you'd get with eager attention or Flash Attention), and it was too slow for my taste compared to Flash Attention so I just dropped it. And I wasn't even doing anything super exotic with it.

Maybe it's better now, but I'd still consider using FlexAttention without a corresponding unit test checking its accuracy against an equivalent eager implementation completely irresponsible.

replies(1): >>43578074 #
2. gessha ◴[04 Apr 25 03:33 UTC] No.43578074[source]▶
What unit tests do you use for nn modules and how do you come up with them?
replies(2): >>43578654 #>>43580802 #
3. porridgeraisin ◴[04 Apr 25 05:30 UTC] No.43578654[source]▶
I would like to know too
4. kouteiheika ◴[04 Apr 25 11:40 UTC] No.43580802[source]▶
Unit tests which test random inputs across different sizes (e.g. with different number of heads, head sizes, embedding dimensions, etc.) and compare two different implementations' output to each other (e.g. attention implemented manually in an eager fashion vs a bunch of accelerated attention libraries).

Also more integration-like tests where I take an already pretrained model, load it using an established library (e.g. Huggingface Transformers) and I also load the very same checkpoint into my reimplementation (where I vary the implementation, e.g. swap the attention implementation) and compare the outputs. Funnily enough, I recently even found a bug in HF's Transformers this way when I updated to a newer version and my previously matching output was not matching anymore.