You Only Cache Once: Decoder-Decoder Architectures for Language Models (2024)

Microsoft researchers Li Dong and colleagues introduced YOCO, a decoder-decoder architecture for large language models that caches key-value pairs only once, in an arXiv paper submitted 8 May 2024 and revised 9 May 2024. YOCO stacks a cross-decoder on a self-decoder that encodes global KV caches reused via cross-attention, reducing GPU memory demands while retaining global attention, and its computation flow allows prefilling to early exit without changing the final output. The authors report YOCO improves inference memory, prefill latency, and throughput by orders of magnitude across context lengths and model sizes, and extends to 1M context length with near-perfect needle retrieval accuracy, with code available at aka.ms/YOCO.

Computer Science Computation and Language Submitted on 8 May 2024 v1 https://arxiv.org/abs/2405.05254v1 , last revised 9 May 2024 this version, v2 Title:You Only Cache Once: Decoder-Decoder Architectures for Language Models View PDF /pdf/2405.05254 HTML experimental https://arxiv.org/html/2405.05254v2 Abstract:We introduce a decoder-decoder architecture, YOCO, for large language models, which only caches key-value pairs once. It consists of two components, i.e., a cross-decoder stacked upon a self-decoder. The self-decoder efficiently encodes global key-value KV caches that are reused by the cross-decoder via cross-attention. The overall model behaves like a decoder-only Transformer, although YOCO only caches once. The design substantially reduces GPU memory demands, yet retains global attention capability. Additionally, the computation flow enables prefilling to early exit without changing the final output, thereby significantly speeding up the prefill stage. Experimental results demonstrate that YOCO achieves favorable performance compared to Transformer in various settings of scaling up model size and number of training tokens. We also extend YOCO to 1M context length with near-perfect needle retrieval accuracy. The profiling results show that YOCO improves inference memory, prefill latency, and throughput by orders of magnitude across context lengths and model sizes. Code is available at this https URL https://aka.ms/YOCO . Submission history From: Li Dong view email /show-email/d1dc91d7/2405.05254 Wed, 8 May 2024 17:57:39 UTC 456 KB \ v1\ /abs/2405.05254v1 v2 Thu, 9 May 2024 14:12:45 UTC 455 KB References & Citations Loading... Bibliographic and Citation Tools Bibliographic Explorer What is the Explorer? https://info.arxiv.org/labs/showcase.html arxiv-bibliographic-explorer Connected Papers What is Connected Papers? https://www.connectedpapers.com/about Litmaps What is Litmaps? https://www.litmaps.co/ scite Smart Citations What are Smart Citations? https://www.scite.ai/ Code, Data and Media Associated with this Article alphaXiv What is alphaXiv? https://alphaxiv.org/ CatalyzeX Code Finder for Papers What is CatalyzeX? https://www.catalyzex.com DagsHub What is DagsHub? https://dagshub.com/ Gotit.pub What is GotitPub? http://gotit.pub/faq Hugging Face What is Huggingface? https://huggingface.co/huggingface ScienceCast What is ScienceCast? https://sciencecast.org/welcome Demos Recommenders and Search Tools Influence Flower What are Influence Flowers? https://influencemap.cmlab.dev/ CORE Recommender What is CORE? https://core.ac.uk/services/recommender arXivLabs: experimental projects with community collaborators arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website. Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them. Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs https://info.arxiv.org/labs/index.html .