tomaarsen / attention_sinks

Extend existing LLMs way beyond the original training length with constant memory usage, without retraining
https://huggingface.co/blog/tomaarsen/attention-sinks
Apache License 2.0
650 stars 41 forks source link

Add benchmarks comparing against Sliding Window Attention #10

Open casper-hansen opened 10 months ago

casper-hansen commented 10 months ago

As the title suggests, we are missing benchmarks on SWA. From my understanding, the only difference is the 4 tokens as sinks that we save.

tomaarsen commented 10 months ago

Hello!

The paper indeed shows 4 forms of attention: a) Dense attention. b) Window attention. c) Window attention with re-calculations. d) Window attention with sink tokens.

And I only benchmark 3 of these: a (transformers), b (windowed) and d (attention_sinks). The only missing one is c: Window attention with re-calculations, i.e. not the one that only differs with the 4 sink keys, that's regular window attention. Window attention with re-calculations seems more involved to implement, though I'm not actually very familiar with it. That's one of the two why it's not implemented in my experiments.

The other reason is that it's (in my opinion) a less interesting comparison than vs dense and regular window attention, as it should perform equivalently but just be a lot slower.

I've been a bit busy this week with work, so I haven't been able to research how to implement this attention approach properly.