Improve LLM Inference efficiency with YOCO (You Only Cache Once) architecture

With every new release of LLMs, we see increase in context window length. Initial models had context length of few thousand tokens and today we are experiencing almost infinite context length that can process a complete book or many hours of video content. However, a significant challenge faced by consumers of Large Language Models is that the inference latency substantially increases with the context length. More the context length, more the compute power required, more the GPU’s, more we pay and more the latency! Architectural intervention is required to overcome this challenge.

In a recent research paper titled “You Only Cache Once (YOCO): Decoder-Decoder Architectures for Language Models https://arxiv.org/abs/2405.05254", the researchers have addressed this key challenge in LLMs.

?How does YOCO address the problem??

• Traditional Approach: In the traditional decoder-only Transformers based LLM architecture, the model processes the input sequence one token at a time. For each token, it needs to consider all previous tokens in the sequence. This requires caching key-value (KV) pairs representing the entire sequence. As the sequence length grows, the cache size and memory usage become significant bottlenecks.

YOCO's Innovation: YOCO proposes a novel decoder-decoder architecture that breaks this cycle of repeated caching. It has two parts:

• Self-decoder: This initial decoder processes the entire input sequence. It efficiently generates global KV caches using self-attention and these caches capture the essence of the entire input sequence.
• Cross-decoder: Subsequent decoder layers reuse the globally cached KV pairs from the self-decoder. They employ a special attention mechanism called "cross-attention" to selectively focus on relevant parts of the cached information for generating the next token in the output sequence.

?Benefits of YOCO:

• Reduced Memory Footprint: By caching key-value pairs only once, YOCO significantly reduces memory demands during inference, especially for long sequences. This allows for processing longer inputs on the same hardware.
• Maintains Global Attention: Despite caching once, YOCO retains the capability of considering the entire input sequence through the cleverly designed cross-attention mechanism in the cross-decoder.
• Improved Efficiency: YOCO demonstrates improvements in inference speed and throughput compared to traditional Transformers.

Overall YOCO achieves significantly better performance compared with Transformers and offers a promising approach for building more efficient and scalable LLMs.

C5i