Attention | Infini-attention

• Related Project: Private
• Category: Paper Review
• Date: 2024-04-15

Leave No Context Behind: Efficient Infinite Context Transformers with Infini-attention

• url: https://arxiv.org/abs/2404.07143
• pdf: https://arxiv.org/pdf/2404.07143
• html: https://arxiv.org/html/2404.07143v1
• abstract: This work introduces an efficient method to scale Transformer-based Large Language Models (LLMs) to infinitely long inputs with bounded memory and computation. A key component in our proposed approach is a new attention technique dubbed Infini-attention. The Infini-attention incorporates a compressive memory into the vanilla attention mechanism and builds in both masked local attention and long-term linear attention mechanisms in a single Transformer block. We demonstrate the effectiveness of our approach on long-context language modeling benchmarks, 1M sequence length passkey context block retrieval and 500K length book summarization tasks with 1B and 8B LLMs. Our approach introduces minimal bounded memory parameters and enables fast streaming inference for LLMs.

TL;DR

• 문제 정의: 다중 언어 대규모 언어모델(MLLM)의 기억력 제한 문제와 해결 방안.
• 주요 방법: Infini-attention을 통한 기억 용량과 계산 비용의 효율적 관리.
• 기여도: 효율적인 장기 기억 처리를 가능하게 하는 새로운 attention 메커니즘 제안.

1. 문제 및 해결 방안

문제: 기존 Transformer 및 LLMs는 attention 메커니즘의 복잡성으로 인해 기억 용량과 계산 시간이 급격히 증가하며, 이로 인해 긴 시퀀스를 처리하는 데 한계가 있습니다.

해결 방안: Infini-attention은 기존 attention 메커니즘에 메모리 압축(compressive memory)를 통합하여 계산 및 저장 공간을 효율적으로 관리할 수 있는 새로운 구조를 제안합니다. 이는 기존 데이터를 버리지 않고 재사용하여 정보를 보다 효과적으로 관리합니다.

2. 선행연구 및 기존 방법

• 선행연구: Transformer와 그 파생 모델들은 attention을 기반으로 하는데, 이는 입력 시퀀스 길이에 따라 메모리와 계산 복잡도가 제곱적으로 증가합니다.
• 기존 방법 문제점: 긴 문맥을 효율적으로 처리하지 못하며, 대용량 모델의 경우 자원 소모가 큽니다.

3. 방법 및 주장

3.1 Infini-attention 메커니즘

Infini-attention은 기존 attention과 메모리 압축를 결합하여 메모리 효율성을 개선하고, 긴 문맥의 정보를 효과적으로 처리할 수 있도록 설계되었습니다.

• Dot-product Attention의 수학적 기초: \(\text{Attention}(Q, K, V) = \text{softmax}\left(\text{QK^T}{\sqrt{d_k}}\right)V\) 수식에서 \(Q, K, V\)는 각각 Query, Key, Value를 나타내며, \(d_k\)는 Key의 차원입니다. 이 수식은 입력 간의 관계를 확률적으로 모델링하여 중요한 정보에 더 많은 가중치를 두도록 합니다.

• 메모리 압축의 역할과 효과: \(M_{s+1} = M_s + \sigma(K)V\) 수식에서 \(M_s\)는 현재 메모리 상태, \(K\)와 \(V\)는 각각 Key와 Value, \(\sigma\)는 활성화 함수입니다. 이 식은 새로운 정보를 메모리에 효율적으로 통합하고, 기존 정보와의 연관성을 고려하여 업데이트합니다.

기본 Attention 메커니즘

기본적인 attention 메커니즘은 입력 간의 관계를 계산하여 중요한 정보에 가중치를 두는 방식입니다. 이는 다음과 같은 수학적 표현으로 나타낼 수 있습니다.

수식에서 \(Q\), \(K\), \(V\)는 각각 Query, Key, Value를 나타내며, \(d_k\)는 Key 벡터의 차원입니다. 이 공식은 Key와 Query 간의 유사도를 계산하여 해당 유사도가 높은 Value에 더 큰 가중치를 부여합니다. 이는 모델이 문맥상 중요한 정보를 효과적으로 추출하도록 돕습니다.

메모리 압축

Infini-attention은 기존의 attention 메커니즘에 메모리 압축를 추가하여, 과거 정보를 보다 효율적으로 활용할 수 있게 합니다. 메모리 압축의 기능은 다음과 같은 수식으로 설명됩니다.

이 식에서 \(M_s\)는 현재 메모리 상태를 나타내며, \(K\)와 \(V\)는 각각 Key와 Value, \(\sigma\)는 활성화 함수입니다. 이 과정은 새로운 정보를 기존 메모리에 효율적으로 통합하며, 이전 정보와의 연관성을 고려하여 메모리를 업데이트합니다.

메모리 업데이트 및 검색

메모리의 업데이트 및 검색 과정은 다음과 같은 식으로 모델링됩니다.

이런 메모리 업데이트는 새로운 입력에 대한 반응을 메모리에 저장하여, 필요할 때 즉시 이 정보를 검색할 수 있게 합니다. 이는 메모리의 효율성을 크게 향상시키며, 긴 문맥 정보를 처리하는 데 필수적입니다.

3.2 근거 및 논리적 연결

• 근거: 기존의 Transformer 모델은 긴 문맥을 처리하기 위해 필요한 메모리가 기하급수적으로 증가하는 문제가 있습니다. Infini-attention은 이를 메모리 압축를 통해 해결하여, 동일한 길이의 문맥을 더 적은 리소스로 처리할 수 있게 합니다.
• 논리적 연결: 메모리 압축는 이전의 정보를 효율적으로 저장하고 필요할 때 즉시 접근할 수 있도록 해, 긴 문맥의 정보를 누락 없이 처리할 수 있게 합니다. 이는 모델이 보다 복잡한 패턴과 문맥을 학습하는 데 도움을 줍니다.

4. 실험 결과 및 평가

• 실험 설정: Infini-attention을 통합한 모델은 기존 모델과 비교하여 향상된 문맥 이해 능력으로 긴 문맥의 언어 모델링 벤치마크에서 기존 모델보다 더 좋은 성능을 보였습니다.
• 성능 평가: Infini-attention은 기존 대비 114배의 메모리 압축 비율을 갖고, 특히 긴 시퀀스 처리에 있어서 우수한 결과를 보입니다.

5. 결론 및 향후 연구 방향

Infini-attention은 긴 문맥의 정보를 효과적으로 처리할 수 있는 새로운 기법으로, 기존의 attention 메커니즘의 한계를 극복하고, 효율적인 장기 기억 처리를 가능하게 함으로써 기존의 메모리 및 계산 비용 문제를 크게 개선합니다.

이 연구는 언어 모델이 더 긴 문맥을 이해하고 처리할 수 있도록 하며, 향후 다양한 언어 처리 작업에 효과적으로 적용될 수 있을 것입니다. 또한, 이 기법은 모델의 일반화 능력을 높이고, 실제 어플리케이션에서의 유용성을 증대시킬 수 있습니다.

앞으로의 연구에서는 더 다양한 언어 및 문맥에 대한 적용을 통해 그 범용성을 검증할 계획이라고 합니다.