[논문 리뷰] Retentive Network: A Successor to Transformer for Large Language Models

Yutao Sun, Li Dong외 6명|arXiv (Cornell University)|2023. 07. 17.

Topic Modeling인용 수 107

한 줄 요약

레텐티브 네트워크(RetNet)는 주의(attention)을 대체하기 위해 다중 스케일 유지 메커니즘을 도입하여 병렬 학습, O(1) 추론, 그리고 선형적인 긴 시퀀스 메모리를 가능하게 하며 Transformers와 경쟁력 있는 성능을 제공합니다.

ABSTRACT

In this work, we propose Retentive Network (RetNet) as a foundation architecture for large language models, simultaneously achieving training parallelism, low-cost inference, and good performance. We theoretically derive the connection between recurrence and attention. Then we propose the retention mechanism for sequence modeling, which supports three computation paradigms, i.e., parallel, recurrent, and chunkwise recurrent. Specifically, the parallel representation allows for training parallelism. The recurrent representation enables low-cost $O(1)$ inference, which improves decoding throughput, latency, and GPU memory without sacrificing performance. The chunkwise recurrent representation facilitates efficient long-sequence modeling with linear complexity, where each chunk is encoded parallelly while recurrently summarizing the chunks. Experimental results on language modeling show that RetNet achieves favorable scaling results, parallel training, low-cost deployment, and efficient inference. The intriguing properties make RetNet a strong successor to Transformer for large language models. Code will be available at https://aka.ms/retnet.

연구 동기 및 목표

• 성능을 희생하지 않으면서 LLM 배치에서 추론 비용과 메모리 사용을 줄이려는 동기를 제시한다.
• 학습 병렬성을 유지하는 Transformer의 후속 아키텍처를 개발한다.
• 학습 및 추론을 최적화하기 위해 병렬, 순환 및 청크 단위 순환 표현을 지원하는 유지 메커니즘을 도입한다.

제안 방법

• 멀티헤드 어텐션을 대체하는 다중 스케일 유지(MSR) 모듈을 제안한다.
• 유지의 이중 표현을 도출한다: 병렬형(훈련 친화적) 및 순환형(추론 친화적).
• 긴 시퀀스를 위한 세 가지 계산 패러다임을 구현한다: 병렬 유지, 순환 유지, 및 청크 단위 순환 유지.
• 표현력 및 학습 안정성을 높이기 위해 게이팅(swish)과 다중 헤드 디케이(gamma)를 도입한다.
• 다중 스케일 헤드로 인한 각 헤드의 분산을 처리하기 위해 GroupNorm을 사용한다.
• MSR + FFN 블록 및 학습/추론 전략과 함께 엔드투엔드 RetNet 아키텍처를 제공한다.
• 확장성, 학습 비용, 추론 지표 전반에서 RetNet을 Transformer 및 효율적인 Transformer 변형과 비교한다.

실험 결과

연구 질문

• RQ1 RetNet가 Transformer와 비교하여 학습 병렬성을 달성하면서 추론 효율성을 유지하거나 향상시킬 수 있는가?
• RQ2 저장소 메모리, 대기시간, 계산 비용을 줄이면서도 유지 메커니즘이 언어 모델링 성능을 유사하게 제공하는가?
• RQ3 병렬, 순환, 청크 단위 순환 표현이 긴 시퀀스 모델링 및 확장성에 어떤 영향을 미치는가?
• RQ4 RetNet으로 대형 모델과 긴 컨텍스트에서 달성할 수 있는 메모리, 처리량, 지연 시간의 이점은 무엇인가?
• RQ5 Transformer에 비해 RetNet이 제로/적은 샷 다운스트림 작업에서 어떻게 성능을 나타내는가?

주요 결과

• RetNet은 Transformer와 비교했을 때 우수한 확장성, 병렬 학습, 저비용 배포 및 효율적인 추론을 달성한다.
• 8k 컨텍스트를 가진 7B 모델의 경우 RetNet은 Transformer의 키-값 캐시 대비 디코딩이 8.4× 빨라지고 메모리의 70%를 절감한다.
• 학습 중에 RetNet은 일반적인 Transformer 대비 25–50%의 메모리를 절약하고 7× 속도를 가속하며, FlashAttention과도 경쟁력이 있다.
• RetNet의 추론 대기 시간은 길이에 의존하지 않으며 배치 크기에 덜 민감하여 디코딩 시 처리량이 높아진다.
• RetNet은 Transformer에 비견되는 언어 모델링 퍼플렉시티를 달성하고 여러 작업에서 제로/적은 샷 학습 성능이 우수하다.
• 게이팅, 그룹 정규화, 다중 스케일 디케이가 성능 향상에 기여하며, 더 큰 헤드 차원이 결과를 개선한다.