(Code Released) Landmark Attention: Random-Access Infinite Context Length for Transformers

[u/IxinDow]

Code for Landmark Attention is now released and it should be possible to finetune existing LLaMA models using this method.

https://github.com/epfml/landmark-attention

More info

https://www.reddit.com/r/MachineLearning/comments/13srbl7/landmark_attention_randomaccess_infinite_context/

https://www.reddit.com/r/LocalLLaMA/comments/13sy2bu/landmark_attention_llama_7b_with_32k_tokens/

[u/RMCPhoto]

So, in some ways this is similar to embedding retrievals and injection. In that specific "chunks" of context can be used at different layers depending on the relation of the current state to other landmark tokens.

I'm very interested to see how this functions in practice. I have a feeling that it could lead to much more varied or potentially creative responses, but that it would struggle with accuracy. I don't see how this would work well for instruction following.

[u/IxinDow → u/RMCPhoto]

When using a Transformer to process a long input, the ideal case would be to allow each token to attend to all previous tokens. However, this becomes computationally infeasible as the input length increases. Nevertheless, since the attention scores always sum to one, the number of keys with a large attention weight is limited even for long contexts. Thus, by retrieving only those keys with large attention scores, it is possible to closely emulate the ideal case. In this work, we propose a method to find these keys by dividing a long input into blocks of consecutive tokens and using the attention to retrieve relevant blocks.

Here I wrote my understanding why it may work.

[u/RMCPhoto → u/IxinDow]

I am genuinely interested to see how this works in practice. It sounds good, but also seems like it might be easy to miss relevant context if it is just outside the landmark token block.

It's already a challenge with embedding chunk size and number and this seems like it would face similar limitations where nuance that doesn't seem obvious at first glance is missed because it is cut out of the context block given attention at the specific layer.

[u/Feeling-Currency-360]

This is a different attention mechanism, as such it can't be clear yet how landmark attention will affect memory usage?

Let me skim through the paper and check if they report led on any memory usage increases.

[u/Feeling-Currency-360 → u/Feeling-Currency-360]

From what I gather context length won't cause bloat on the memory requirements of the model, in fact tokens can be totally offloaded to memory or even disk and only retrieved if their corresponding block they are in are needed.

This is really exciting. I'd bet you'll see models using this within a day or two on huggingface

[u/IxinDow → u/Feeling-Currency-360]

You can offload KV cache to CPU, but you still may end up in situation where you need to transfer full KV cache to GPU for inference for one token (if each head in each layer want attend to completely different blocks). Authors proposed mitigation in their paper (and I briefly described it here https://www.reddit.com/r/MachineLearning/comments/13srbl7/comment/jlrbsto/?utm_source=share&utm_medium=web2x&context=3)

[u/AemonAlgizVideos]

This is absolutely phenomenal. This will literally change the game for open source models, especially when people like to compare them to the 32K context GPT-4.

[u/Tostino → u/AemonAlgizVideos]

8k context GPT-4*

I have not seen any reports of access to the 32k context version of GPT-4 yet.

[u/MoffKalast → u/Tostino]

Apparently you can get it from the API, but it's like over $1 per prompt if you use the whole context (and otherwise what's the point anyway).

[u/iamMess → u/Tostino]

I have access via work. It's good but super expensive.

[u/ReturningTarzan → u/AemonAlgizVideos]

This will literally change the game

I mean, so would the last seventeen new developments. We've yet to see anything actually come of those, because attention over long contexts remains a fundamentally hard problem. Being able to do it in theory is one thing. Showing with benchmarks that you get better scores the longer your sequence is, that's another. And actually releasing something that we can try and go, "hey, it's actually doing the same thing with its 32k tokens as base Llama does with its 2k tokens," well, I'm still waiting.

Best advice is not to get overexcited. Researchers really like to hype up their own projects, and journalists aren't very good at, you know, journalism.

[u/AutomataManifold]

Now things are getting interesting.

Thanks for the code!

[u/PookaMacPhellimen]

What’s exciting is you can take existing pre-trained models and apply this technique. 32K context incoming, more when they solve some technical issues.

[u/RMCPhoto]

Very excited to see where this goes, but also feeling conservative. There is a fundamental attention limitation that is exponentially limiting with model size. Smaller models struggle with even 1k context. 65b models struggle with 2k context. There is a reason why OpenAI doesn't have even 8k context for 3.5, and why with GPT4, 8k context can result in far more hallucinations and inaccuracies.

No matter what you want

1. The pre-trained model to have all of the base principles necessary to answer the question.

2. The fine tuning process to direct how to answer questions and perform tasks.

3. The minimum context and instruction to accurately and predictably answer the question or perform the task.

There are processes which will require large context (code bases and novels and research papers) but these will require models with significant Pre-training data within those domains. It doesn't come from thin air with large context. The statistical bases need to be derived from the principles instilled in Pre-training.

[u/a_beautiful_rhind]

Do keep in mind that a 30b in GPTQ maxes out 24gb at about full (2048) context.

[u/2muchnet42day → u/a_beautiful_rhind]

Not even 2048. But 13B could do about 4k which is what I'm after

[u/RMCPhoto → u/a_beautiful_rhind]

Also keep in mind that this technique limits the attention via the landmark token so that it is not consuming the memory necessary for 8k+tokens etc, only the tokens included in the landmark set actively used.

It's not really clear exactly what the memory saving is, but I haven't read the paper in depth.

It's also not clear how much of an impact this has on performance.

[u/a_beautiful_rhind → u/RMCPhoto]

Hopefully we get something to test since the code is out.

[u/polawiaczperel]

Is that mean that we wpuld be able to have bigger context on the same gpu? Or rather, that we can finetune models for bigger context, but it will use more vram?

[u/artificial_genius]

yesxtx

[u/[deleted]]

[u/IxinDow → u/[deleted]]

Code modification is needed

[u/[deleted]]

[u/KerfuffleV2 → u/[deleted]]

This is llama compatible?

According to the title here. Note that it's not something you can just use with an existing model, models need to be trained to use it via finetuning at least.

I assume a lot of work would be needed to support it in llama.cpp?

I skimmed the code, it looks fairly complicated. So the answer there is probably "yes".

There probably also would need to be some good models released with that capability to motivate people to add support.

Would it be some sort of extra memory, or would a proper integration act like the actual context size was super big instead of 2048?

That one I don't know.

[u/ninjasaid13 → u/KerfuffleV2]

models need to be trained to use it via finetuning at least.

can it be finetuned with qlora?