about
You Only Cache Once: Decoder-Decoder Architectures for Language Models (2024) (arxiv.org)
1 point by theanonymousone 22 days ago | hide | past | pdf | discuss on HN

In plain words: The model splits in two: the lower half builds the key-value memory attention layers store, and the upper half reuses it instead of caching its own. It matched Transformers while cutting memory and prefill time by orders of magnitude, and handled 1M-token contexts with near-perfect retrieval.

Abstract · You Only Cache Once: Decoder-Decoder Architectures for Language Models

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 https://aka.ms/YOCO.

Yutao Sun, Li Dong, Yi Zhu, Shaohan Huang, Wenhui Wang, Shuming Ma, Quanlu Zhang, Jianyong Wang, Furu Wei
arXiv:2405.05254 · cs.CL · submitted May 8, 2024 · updated May 9, 2024
abstract · pdf · html

add comment on HN
Also discussed: May 2024 (3 points, 0 comments) · May 2024 (3 points, 1 comment)