Cross-Layer Attention

• Related Project: Private
• Category: Paper Review
• Date: 2024-05-25

Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

• url: https://arxiv.org/abs/2405.12981
• pdf: https://arxiv.org/pdf/2405.12981
• html https://arxiv.org/html/2405.12981v1
• abstract: Key-value (KV) caching plays an essential role in accelerating decoding for transformer-based autoregressive large language models (LLMs). However, the amount of memory required to store the KV cache can become prohibitive at long sequence lengths and large batch sizes. Since the invention of the transformer, two of the most effective interventions discovered for reducing the size of the KV cache have been Multi-Query Attention (MQA) and its generalization, Grouped-Query Attention (GQA). MQA and GQA both modify the design of the attention block so that multiple query heads can share a single key/value head, reducing the number of distinct key/value heads by a large factor while only minimally degrading accuracy. In this paper, we show that it is possible to take Multi-Query Attention a step further by also sharing key and value heads between adjacent layers, yielding a new attention design we call Cross-Layer Attention (CLA). With CLA, we find that it is possible to reduce the size of the KV cache by another 2x while maintaining nearly the same accuracy as unmodified MQA. In experiments training 1B- and 3B-parameter models from scratch, we demonstrate that CLA provides a Pareto improvement over the memory/accuracy tradeoffs which are possible with traditional MQA, enabling inference with longer sequence lengths and larger batch sizes than would otherwise be possible

Contents

• Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
  • TL;DR

TL;DR

대규모 언어모델의 키-값 캐시 메모리 최적화: Cross-Layer Attention 기법

• 대규모 언어모델(LLMs)에서 키-값(KV) 캐시의 메모리 사용을 줄이기 위해 Cross-Layer Attention(CLA) 기법을 제안합니다.
• CLA는 여러 계층 간 KV 캐시를 공유함으로써 필요한 메모리 양을 감소시키고 모델의 효율성을 향상시킵니다.
• 실제 실험을 통해 CLA가 기존의 어텐션 메커니즘과 비교하여 메모리 사용량을 절감하면서도 정확도를 유지하거나 향상시키는 것을 확인하였습니다.

서론

최근 대규모 언어모델(LLMs)의 사용이 증가함에 따라, 모델의 효율성과 확장성을 개선할 수 있는 방법에 대한 연구가 활발히 이루어지고 있습니다. 특히, 트랜스포머 기반 모델에서 주목받는 부분 중 하나는 키-값(KV) 캐시의 메모리 사용량입니다. KV 캐시는 시퀀스 길이와 배치 크기에 비례하여 확장되므로, 큰 시퀀스 길이를 처리할 때 메모리 사용량이 크게 증가합니다. 이는 배치 크기를 제한하고, 비용이 많이 드는 기술을 사용하여 메모리 부담을 줄이는 등의 문제로 이어집니다.

선행 연구

이전 연구들은 다양한 방법으로 KV 캐시의 메모리 사용량을 줄이기 위한 방법을 제안했습니다. 예를 들어, 낮은 Precision로 KV 활성화를 저장하거나 중요하지 않은 KV 캐시 항목을 제거하는 방법, 그리고 어텐션 메커니즘 내의 쿼리 헤드들 간에 키와 값 공유 등이 있습니다. 그러나 이런 접근 방식들은 각각의 단점을 가지고 있으며, 새로운 모델 아키텍처가 필요한 상황입니다.

연구 방법

Cross-Layer Attention (CLA)

KV 캐시의 크기를 줄이는 새로운 차원의 접근 방식을 제안합니다. CLA는 트랜스포머 아키텍처에 변형을 가하여, 여러 계층에 걸쳐 KV 활성화를 공유함으로써 KV 캐시의 크기를 줄입니다. 이 방식은 특히 메모리 부담이 큰 대규모 모델에 유용하며, 메모리 사용 효율성을 향상시키면서도 모델의 정확도를 유지하거나 향상시킬 수 있습니다.

실험 설정

다양한 구성의 CLA를 통해 정확도와 메모리 사용량을 측정하기 위해 광범위한 사전 훈련 실험을 수행했습니다. 이 실험들은 다양한 아키텍처 하이퍼파라미터, 학습률 및 모델 크기에 걸쳐 수행되었습니다.

결과

CLA는 기존의 Multi-Query Attention (MQA) 및 Grouped-Query Attention (GQA) 아키텍처에 비해 정확도/메모리 파레토 개선을 가능하게 합니다. 특히, 1B 및 3B 파라미터 스케일에서 CLA와 MQA를 결합한 경우, MQA 베이스 라인 대비 KV 캐시 크기를 2배 줄일 수 있었으며, 이는 복잡도에 거의 영향을 주지 않고 성능을 유지하거나 향상시켰습니다.

결론 및 미래의 연구 방향

CLA는 대규모 언어모델에서 KV 캐시의 메모리 요구 사항을 줄이면서도 모델 성능을 유지하거나 향상시킬 수 있는 효과적인 방법입니다. 향후 연구에서는 더 긴 시퀀스와 더 큰 배치 처리를 통해 CLA의 효율성을 평가하는 것이 중요할 것입니다.