논문리뷰 | Streaming LLM

https://github.com/microsoft/vidur

GitHub - microsoft/vidur: A large-scale simulation framework for LLM inference

 

 

1. Abstract

 

대형 언어 모델(LLM)은 대화 시스템, 문서 요약, 코드 완성, 질의 응답 등 다양한 자연어 처리 응용 프로그램에서 필수적입니다. 그러나 LLM이 더 긴 시퀀스를 안정적으로 생성하도록 하기에는 한계가 있으며, 이는 주로 사전 훈련 시 설정된 주의(attention) 창 크기로 인해 발생합니다.

 

 

Figure1

 

Figure 1: StreamingLLM과 기존 방법 비교

길이 LLL의 텍스트로 사전 학습된 언어 모델은 TTT번째 토큰을 예측합니다 (T≫LT \gg LT≫L).
(a) Dense Attention은 O(T2)O(T^2)O(T2)의 시간 복잡도를 가지며 캐시 크기가 증가합니다. 텍스트 길이가 학습된 길이를 초과하면 성능이 저하됩니다.
(b) Window Attention은 최근 LLL개의 토큰의 키-값(KV) 상태를 캐시하지만, 초기 토큰의 KV가 삭제되면 성능이 급격히 저하됩니다.
(c) Sliding Window with Re-computation은 새 토큰이 들어올 때마다 최근 LLL개의 토큰에 대해 KV 상태를 재구성하며 긴 텍스트에서 잘 작동하지만, O(TL2)O(TL^2)O(TL2) 복잡도로 인해 매우 느립니다.
(d) StreamingLLM은 주의 집중 계산의 안정성을 위해 “attention sink” (초기 몇 개의 토큰)을 유지하며, 최신 토큰과 결합하여 효율적이고 안정적인 성능을 제공합니다.

 

윈도우 어텐션(Beltagy et al., 2020)(그림 1 b)이라고 알려진 직관적인 접근법은 가장 최근 토큰의 KV 상태에 대해서만 고정 크기의 슬라이딩 윈도우를 유지합니다. 이는 캐시가 초기에 채워진 후 일정한 메모리 사용량과 디코딩 속도를 보장하지만, 시퀀스 길이가 캐시 크기를 초과하면 모델이 붕괴됩니다. 예를 들어 첫 번째 토큰의 KV를 제거하는 것만으로도 문제가 발생할 수 있습니다(그림 3 참조). 다른 전략은 재계산을 통한 슬라이딩 윈도우(그림 1 c)로, 각 생성된 토큰에 대해 최근 토큰의 KV 상태를 다시 구축합니다. 이 방법은 성능이 우수하지만, 윈도우 내에서 이루어지는 2차 어텐션 계산으로 인해 상당히 느리기 때문에 실제 스트리밍 애플리케이션에 사용하기에는 적합하지 않습니다.

 

윈도우 어텐션의 실패를 이해하기 위해, 우리는 자기회귀 언어모델(LLM)에서 흥미로운 현상을 발견했습니다. 그림 2에서 볼 수 있듯이, 언어 모델링 작업의 관련성과 상관없이 초기 토큰에 상당한 어텐션 점수가 할당됩니다. 우리는 이러한 토큰을 "어텐션 싱크"라고 부릅니다. 이들은 의미적 중요성이 부족함에도 불구하고 상당한 어텐션 점수를 받습니다. 이는 Softmax 연산 때문이라고 생각합니다. Softmax는 모든 문맥 토큰에 대한 어텐션 점수의 합이 1이 되도록 요구하기 때문입니다. 따라서 현재 쿼리가 이전 토큰과 강한 매치를 갖지 않더라도, 모델은 이 불필요한 어텐션 값을 어딘가에 할당해야 합니다. 초기 토큰이 싱크 토큰이 되는 이유는 직관적입니다. 자기회귀 언어 모델링의 특성상 초기 토큰은 모든 후속 토큰에 노출되므로, 어텐션 싱크로 활용되기 쉽습니다.

Figure 2

 

이러한 통찰을 바탕으로, 우리는 StreamingLLM이라는 간단하고 효율적인 프레임워크를 제안합니다. 이를 통해 유한 어텐션 윈도우로 학습된 LLM이 추가 fine-tuning 없이도 무한 길이의 텍스트를 처리할 수 있습니다. StreamingLLM은 어텐션 싱크 토큰이 높은 어텐션 값을 갖는다는 사실을 활용합니다. 이러한 토큰을 유지하면 어텐션 점수 분포를 정상에 가깝게 유지할 수 있습니다. 따라서 StreamingLLM은 단순히 어텐션 싱크 토큰의 KV(초기 4개 토큰으로 충분)와 슬라이딩 윈도우의 KV를 함께 유지하여 어텐션 계산을 앵커링하고 모델 성능을 안정화시킵니다. StreamingLLM을 통해 Llama-2-[7, 13, 70]B, MPT-[7, 30]B, Falcon-[7, 40]B, Pythia-[2.9, 6.9, 12]B 등의 모델이 400만 토큰 이상을 안정적으로 모델링할 수 있습니다. 유일한 실용적 기준선인 재계산을 통한 슬라이딩 윈도우와 비교했을 때, StreamingLLM은 최대 22.2배 속도 향상을 달성하여 LLM의 스트리밍 사용을 실현했습니다.

 

그림 2에서 Llama-2-7B 모델의 평균 어텐션 로짓을 시각화하면 다음과 같은 관찰 사항이 있습니다:

1. 하위 2개 레이어(레이어 0, 1)의 어텐션 맵은 "지역적" 패턴을 보이며, 최근 토큰에 더 많은 어텐션이 할당됩니다.
2. 하위 2개 레이어를 넘어가면, 모델이 모든 레이어와 헤드에 걸쳐 초기 토큰에 상당한 어텐션을 할당합니다.

이러한 관찰을 통해 어텐션 싱크 가설을 확인할 수 있었습니다. 언어 모델은 스트리밍 배포를 위해 단 하나의 어텐션 싱크 토큰만 요구하도록 사전 학습될 수 있습니다. 구체적으로, 모든 학습 샘플의 맨 앞에 추가적인 학습 가능한 토큰을 배치하면 이를 전용 어텐션 싱크로 활용할 수 있습니다. 1억 6천만 개 매개변수의 언어 모델을 처음부터 학습시키면, 이 단일 싱크 토큰을 추가함으로써 스트리밍 경우에도 모델 성능을 유지할 수 있습니다. 이는 기존 모델의 경우 동일한 수준의 성능을 달성하기 위해 여러 개의 초기 토큰을 어텐션 싱크로 재도입해야 하는 것과는 대조됩니다.

 

마지막으로, StreamingLLM은 LLM의 문맥 길이를 확장하지 않고도 KV 캐시 내 토큰에서 일관성 있는 텍스트를 효율적으로 생성할 수 있다는 점을 강조합니다. 이는 최소한의 메모리 사용과 과거 데이터 의존도로 지속적인 운영 요구사항에 잘 부합합니다. 또한 StreamingLLM은 최근 문맥을 증가시키는 방법과 보완적으로 사용할 수 있습니다.

 

---

2. 관련 연구

 

언어 모델을 긴 텍스트에 적용하는 연구는 주로 3가지 초점 분야에서 이루어졌습니다:

1. 길이 외삽(Length Extrapolation): 언어 모델이 훈련된 것보다 긴 텍스트를 처리할 수 있게 하는 것
2. 문맥 창 확장(Context Window Extension): 언어 모델의 어텐션 창 크기를 늘리는 것
3. 긴 텍스트 활용 개선(Improving LLMs' Utilization of Long Text): 언어 모델의 장기 메모리와 긴 문맥 활용을 향상시키는 것

이 세 분야가 서로 관련되어 보이지만, 한 분야의 진전이 다른 분야로 이어지지 않는다는 점에 주목할 필요가 있습니다. 예를 들어, 언어 모델의 문맥 크기를 늘린다고 해서 그를 넘어서는 성능 향상이 보장되지 않으며, 두 접근법 모두 긴 문맥의 효과적인 활용을 보장하지 않습니다.

StreamingLLM 프레임워크는 주로 첫 번째 분야, 즉 길이 외삽에 초점을 맞추고 있습니다. 이는 언어 모델이 훈련 창 크기를 크게 넘어서는, 심지어 무한 길이의 텍스트를 처리할 수 있게 하는 것을 목표로 합니다. 우리는 언어 모델의 어텐션 창 크기를 늘리거나 장기 메모리와 긴 문맥 활용을 개선하지 않습니다. 후자 두 분야는 우리의 초점과 직교하지만, 향후 우리의 기술과 통합될 수 있는 가능성이 있습니다.

 

길이 외삽(Length Extrapolation)은 짧은 텍스트로 학습된 언어 모델이 더 긴 텍스트를 처리할 수 있게 하는 것을 목표로 합니다. 주된 연구 방향은 트랜스포머 모델에 대한 상대적 위치 인코딩 방법을 개발하는 것입니다. 이를 통해 모델이 훈련 창 크기를 넘어서 작동할 수 있게 합니다. 그 중 하나가 Rotary Position Embeddings(RoPE)(Su et al., 2021)입니다. RoPE는 각 어텐션 레이어의 쿼리와 키를 변형하여 상대적 위치 정보를 통합합니다. 이 방법은 유망해 보였지만, 후속 연구(Press et al., 2022; Chen et al., 2023)에 따르면 훈련 창 크기를 넘어서는 텍스트에서 성능이 저하되는 것으로 나타났습니다. 또 다른 접근법인 ALiBi(Press et al., 2022)는 쿼리-키 어텐션 점수에 거리 기반 편향을 도입하여 상대적 위치 정보를 제공합니다. 이는 외삽 성능 향상을 보였지만, MPT 모델에 대한 우리의 테스트에서 텍스트 길이가 훈련 길이를 크게 초과하면 여전히 문제가 발생했습니다. 그러나 현재 방법론으로는 무한 길이 외삽을 달성하지 못하고 있어, 기존 언어 모델로는 스트리밍 애플리케이션에 적합하지 않습니다.

 

문맥 창 확장(Context Window Extension)은 LLM의 문맥 창을 확장하여 한 번의 순방향 통과에서 더 많은 토큰을 처리할 수 있게 하는 것에 초점이 맞춰져 있습니다. 주요 연구 방향은 훈련 효율성 문제에 집중되어 있습니다. 훈련 중 어텐션 계산의 복잡도가 2차 함수이기 때문에, 장문맥 LLM을 개발하는 것은 계산 능력과 메모리 문제 모두를 해결해야 하는 과제입니다. 이에 대한 해결책으로는 FlashAttention과 같은 시스템 최적화(Dao et al., 2022; Dao, 2023)나, 근사 어텐션 방법(Zaheer et al., 2020b; Beltagy et al., 2020; Wang et al., 2020; Kitaev et al., 2020)과 같이 모델 품질을 효율성에 희생하는 방식 등이 제안되었습니다.

최근에는 RoPE를 이용해 사전 훈련된 LLM을 확장하는 연구(Chen et al., 2023; kaiokendev, 2023; bloc97, 2023; Peng et al., 2023)가 활발하게 진행되고 있습니다. 이는 위치 보간과 fine-tuning을 통해 이루어집니다. 그러나 이러한 기술들은 LLM의 문맥 창을 제한적으로만 확장할 수 있어, 무한한 입력을 처리해야 한다는 이 논문의 주된 관심사를 충족시키지 못하고 있습니다.

Figure 3

 

그림 3에서 관찰할 수 있는 내용은 다음과 같습니다:

1. 밀집 어텐션(Dense Attention)은 입력 길이가 사전 훈련 어텐션 창 크기를 초과하면 실패한다.
2. 윈도우 어텐션(Window Attention)은 입력 길이가 캐시 크기를 초과하면 붕괴된다. 즉, 초기 토큰들이 제거되면 문제가 발생한다.
3. StreamingLLM은 안정적인 성능을 보이며, 재계산을 통한 슬라이딩 윈도우 기준선과 거의 유사한 perplexity를 나타낸다.

긴 텍스트의 효과적인 활용 개선(Improving LLMs' Utilization of Long Text)은 LLM이 입력으로 받은 긴 문맥을 단순히 활용하는 것이 아니라 더 잘 포착하고 활용할 수 있게 최적화하는 것을 목표로 합니다. Liu et al.와 Li et al.의 연구에 따르면, 앞서 언급한 두 가지 방향의 성공이 반드시 긴 문맥의 효과적인 활용으로 이어지는 것은 아닙니다. 이러한 장문맥의 효과적인 활용은 여전히 해결해야 할 과제입니다.

본 연구는 가장 최근 토큰들을 안정적으로 활용할 수 있게 하여, LLM의 원활한 스트리밍 적용을 목표로 합니다. 이는 앞서 언급된 다른 두 분야와는 직교하는 접근법입니다.

 

---

3. StreamingLLM

 

3.1 The Faiure of window Attention and Attention Sinks 

윈도우 어텐션 기법은 추론 과정에서 효율성을 제공하지만, 언어 모델링의 perplexity가 매우 높아지는 문제가 있습니다. 결과적으로 모델의 성능이 스트리밍 애플리케이션 배포에 적합하지 않습니다. 이 섹션에서는 '어텐션 싱크(attention sink)' 개념을 사용하여 윈도우 어텐션의 실패를 설명하고, 이를 StreamingLLM의 영감으로 삼았습니다.

 

Figure 3

 

그림 3은 20K 토큰 텍스트에 대한 언어 모델링의 perplexity를 보여줍니다. 텍스트 길이가 캐시 크기를 초과하여 초기 토큰들이 제외될 때 perplexity가 급증하는 것을 명확히 볼 수 있습니다. 이는 초기 토큰들이 예측 대상 토큰과의 거리와 상관없이 LLM의 안정성 유지에 매우 중요하다는 것을 시사합니다.

 

Llama-2-7B와 다른 모델들의 모든 레이어와 헤드의 어텐션 맵을 시각화한 결과(그림 2), 하위 두 레이어를 제외한 모든 레이어와 헤드에서 모델이 지속적으로 초기 토큰들에 집중하는 것을 발견했습니다. 이는 명확한 시사점을 제공합니다: 이러한 초기 토큰들의 KV를 제거하면 어텐션 계산에서 SoftMax 함수(수식 1)의 분모에서 상당한 부분이 제거됩니다. 이러한 변화는 정상적인 추론 환경에서 기대되는 것과는 매우 다른 어텐션 점수 분포를 야기합니다.

수식 1

 

언어 모델링에서 초기 토큰의 중요성에 대해 두 가지 가능한 설명이 있습니다:

 

(1) 토큰의 의미론적 특성이 매우 중요하거나

(2) 모델이 토큰의 절대적 위치에 대한 편향을 학습합니다.

 

이러한 가능성들을 구분하기 위해, 우리는 실험을 수행했으며(표 1), 처음 네 개의 토큰을 줄바꿈 토큰 "\n"으로 대체했습니다. 관찰 결과는 모델이 여전히 이러한 초기 줄바꿈 토큰들을 상당히 강조한다는 것을 보여줍니다. 더욱이, 이들을 다시 도입하면 언어 모델링 Perplexity가 원래의 초기 토큰들을 가졌을 때와 비슷한 수준으로 복원됩니다. 이는 시작 토큰들의 의미론적 가치보다는 절대적 위치가 더 큰 중요성을 가진다는 것을 시사합니다.

 

LLM이 초기 토큰을 어텐션 싱크로 주목함. 모델이 언어 모델링과의 의미론적 관련성과 관계없이 초기 토큰들에 불균형적으로 집중하는 이유를 설명하기 위해, 우리는 "Attention Sink" 개념을 도입합니다. SoftMax 함수(방정식 1)의 특성상 주목받는 모든 토큰들이 0의 값을 가지는 것이 불가능합니다. 이는 현재 임베딩이 예측을 위한 충분한 자체 포함 정보를 가지고 있더라도, 모든 레이어의 모든 헤드에서 다른 토큰들로부터 어떤 정보를 집계해야 함을 의미합니다. 결과적으로, 모델은 불필요한 어텐션 값들을 특정 토큰들에 덤프하는 경향이 있습니다. 양자화 이상치 영역에서도 비슷한 관찰이 이루어졌으며(Xiao 외, 2023; Bondarenko 외, 2023), 이에 대한 잠재적 해결책으로 SoftMax-Off-by-One(Miller, 2023)이 제안되었습니다.

---

 

Table 1

윈도우 어텐션은 긴 텍스트에서 성능이 좋지 않습니다. 초기 4개 토큰을 최근 1020개 토큰과 함께 재도입했을 때(4+1020) Perplexity가 복원됩니다. 원래의 초기 4개 토큰을 줄바꿈 토큰 "\n"으로 대체했을 때도(4"\n"+1020) 비슷한 수준의 펄플렉서티 복원을 달성합니다. 캐시 설정 x+y는 x개의 초기 토큰을 y개의 최근 토큰과 함께 추가하는 것을 나타냅니다. 펄플렉서티는 연결된 PG19 테스트 세트의 첫 번째 책(65K 토큰)에서 측정되었습니다.

 

 

table 2

Table 2

StreamingLLM에서의 효과. (1) 윈도우 어텐션(0+y)은 펄플렉서티가 급격히 증가합니다. (2) 하나 또는 두 개의 초기 토큰을 도입하는 것은 모델 Perplexity를 완전히 복원하지 못하며, 이는 모델이 첫 번째 토큰만을 어텐션 싱크로 사용하지 않는다는 것을 보여줍니다. (3) 4개의 초기 토큰을 도입하는 것이 일반적으로 충분하며, 추가적인 토큰 도입은 수확 체감을 보입니다. 캐시 설정 x+y는 x개의 초기 토큰을 y개의 최근 토큰에 추가하는 것을 나타냅니다. 펄플렉서티는 PG19 테스트 세트의 400K 토큰에서 평가되었습니다.

 

1. 기본 개념 설명:

• 펄플렉서티(perplexity): 언어 모델의 성능을 측정하는 지표입니다. 값이 낮을수록 모델의 성능이 좋다는 의미입니다.
• 토큰(token): 텍스트를 작은 단위로 나눈 것입니다. 단어나 부분 단어가 될 수 있습니다.
  
  2. 실험의 주요 발견:

a) 윈도우 어텐션 문제:

• 긴 텍스트를 처리할 때 최근 토큰만 보는 방식(윈도우 어텐션)은 성능이 좋지 않았습니다.
• 이는 마치 긴 글을 읽을 때 앞부분을 전혀 기억하지 못하는 것과 비슷합니다.

b) 초기 토큰의 중요성:

• 처음 4개의 토큰이 특별히 중요한 것으로 나타났습니다.
• 재미있는 점은, 이 초기 토큰들을 단순한 줄바꿈 문자("\n")로 바꿔도 비슷한 성능을 보였다는 것입니다.
• 이는 토큰의 실제 의미보다는 "위치"가 더 중요하다는 것을 보여줍니다.

c) 최적의 토큰 수:

• 1-2개의 초기 토큰만으로는 충분하지 않았습니다.
• 4개의 초기 토큰이 가장 효과적이었습니다.
• 4개 이상 추가해도 성능 향상이 미미했습니다.
  
  3. 실용적 의미:

• 이 발견은 LLM이 긴 텍스트를 처리할 때 메모리를 더 효율적으로 사용할 수 있게 해줍니다.
• 모든 텍스트를 기억하는 대신, 초기 부분과 최근 부분만 집중적으로 보면서도 좋은 성능을 낼 수 있다는 것을 보여줍니다.

이것은 마치 사람이 긴 책을 읽을 때, 처음 부분과 최근에 읽은 부분을 잘 기억하는 것과 비슷한 패턴을 보여준다고 할 수 있다.

---

 

왜 Llama-2, MPT, Falcon, Pythia와 같은 다양한 자기회귀 LLM들이 다른 토큰들이 아닌 초기 토큰들을 그들의 어텐션 싱크로 일관되게 주목할까? 우리의 설명은 간단합니다: 자기회귀 언어 모델링의 순차적 특성으로 인해, 초기 토큰들은 모든 후속 토큰들에 의해 보여질 수 있는 반면, 후반부 토큰들은 제한된 수의 후속 토큰들에게만 보여질 수 있습니다. 결과적으로, 초기 토큰들은 불필요한 어텐션을 포착하는 어텐션 싱크의 역할을 하도록 더 쉽게 훈련됩니다.

 

우리는 LLM이 일반적으로 하나가 아닌 여러 개의 초기 토큰들을 어텐션 싱크로 활용하도록 훈련된다는 점을 주목했습니다. 그림 2에서 보여지듯이, 어텐션 싱크로서 4개의 초기 토큰을 도입하는 것이 LLM의 성능을 복원하기에 충분합니다. 반면에, 단 하나 또는 두 개를 추가하는 것은 완전한 복구를 달성하지 못합니다. 우리는 이러한 패턴이 사전 훈련 중에 모든 입력 샘플에서 일관된 시작 토큰을 포함하지 않았기 때문에 발생한다고 믿습니다. Llama-2가 각 단락에 "<s>" 토큰을 접두사로 붙이긴 하지만, 이는 텍스트 청킹 이전에 적용되어 대부분 무작위 토큰이 0번째 위치를 차지하게 됩니다. 이러한 균일한 시작 토큰의 부재로 인해 모델은 여러 초기 토큰들을 어텐션 싱크로 사용하게 됩니다. 우리는 모든 훈련 샘플의 시작 부분에 안정적인 학습 가능한 토큰을 통합함으로써, 일관된 스트리밍을 보장하기 위한 여러 초기 토큰들의 필요성을 제거하면서 단일하게 전담 어텐션 싱크로 작용할 수 있다는 가설을 세웠습니다. 우리는 섹션 3.3에서 이 가설을 검증할 것입니다.

 

 

3.2 Rolling KV Cache with Attention Sinks

 

이미 훈련된 LLM에서 LLM 스트리밍을 가능하게 하기 위해, 우리는 모델 미세조정 없이 윈도우 어텐션의 펄플렉서티를 복구할 수 있는 간단한 방법을 제안합니다. 현재의 슬라이딩 윈도우 토큰들과 함께, 우리는 어텐션 계산에서 몇 개의 시작 토큰들의 KV를 재도입합니다. StreamingLLM의 KV 캐시는 개념적으로 그림 4와 같이 두 부분으로 나눌 수 있습니다:

(1) 어텐션 싱크(4개의 초기 토큰들)가 어텐션 계산을 안정화시킵니다.

(2) 롤링 KV 캐시는 언어 모델링에 중요한 가장 최근의 토큰들을 유지합니다.

StreamingLLM의 설계는 다재다능하며 RoPE(Su 외, 2021)와 ALiBi(Press 외, 2022)와 같은 상대적 위치 인코딩을 사용하는 모든 자기회귀 언어 모델에 원활하게 통합될 수 있습니다.

Figure 4 : the KV cache of Streaming LLM

StreamingLLM의 KV 캐시 구조

1. 어텐션 싱크 (Attention Sink)
   • 캐시의 시작 부분에 위치합니다.
   • 일반적으로 처음 4개의 토큰을 포함합니다.
   • 이 토큰들은 항상 캐시에 유지되며, 어텐션 계산을 안정화하는 역할을 합니다.
2. 슬라이딩 윈도우 (Sliding Window)
   • 어텐션 싱크 다음에 위치합니다.
   • 가장 최근의 L-4개 토큰을 포함합니다 (L은 총 캐시 크기).
   • 새로운 토큰이 생성될 때마다 이 윈도우가 이동합니다.
3. 총 캐시 크기
   • 어텐션 싱크(4개 토큰)와 슬라이딩 윈도우(L-4개 토큰)를 합쳐 총 L개의 토큰을 저장합니다.
4. 캐시 업데이트 메커니즘
   • 새로운 토큰이 생성되면 슬라이딩 윈도우의 가장 오래된 토큰이 제거됩니다.
   • 새로운 토큰의 KV는 슬라이딩 윈도우의 끝에 추가됩니다.
   • 어텐션 싱크 토큰들은 항상 그대로 유지됩니다.
5. 효율성
   • 이 구조는 일정한 메모리 사용량을 유지하면서도 무한한 길이의 입력을 처리할 수 있게 합니다.
   • O(TL) 시간 복잡도를 가집니다 (T는 총 토큰 수, L은 캐시 크기).

이 구조를 통해 StreamingLLM은 효율적인 메모리 사용과 안정적인 성능을 동시에 달성할 수 있습니다. 어텐션 싱크가 어텐션 계산을 안정화하는 역할을 하면서, 슬라이딩 윈도우는 최근 컨텍스트를 유지하여 모델이 긴 시퀀스에 대해서도 일관된 성능을 발휘할 수 있게 합니다.

 

StreamingLLM이 토큰들의 상대적 거리를 결정하고 위치 정보를 추가할 때, 원본 텍스트 내의 위치가 아닌 캐시 내의 위치에 집중합니다. 이러한 구분은 StreamingLLM의 성능에 매우 중요합니다. 예를 들어, 현재 캐시(그림 4)가 토큰 [0, 1, 2, 3, 6, 7, 8]을 가지고 있고 9번째 토큰을 디코딩하는 과정에 있다면, 할당되는 위치는 원본 텍스트에서의 위치인 [0, 1, 2, 3, 6, 7, 8, 9]가 아닌 [0, 1, 2, 3, 4, 5, 6, 7]입니다.

RoPE와 같은 인코딩의 경우, 우리는 회전 변환을 도입하기 이전에 토큰들의 키를 캐시합니다. 그런 다음 각 디코딩 단계에서 롤링 캐시의 키들에 위치 변환을 적용합니다. 반면, ALiBi와의 통합은 더 직접적입니다. 여기서는 '점프하는' 바이어스 대신 연속적인 선형 바이어스가 어텐션 점수에 적용됩니다. 캐시 내에서 위치 임베딩을 할당하는 이 방법은 StreamingLLM의 기능에 매우 중요하며, 사전 훈련된 어텐션 윈도우 크기를 넘어서도 모델이 효율적으로 작동하도록 보장합니다.

 

*회전 변환을 도입한다?

캐시 내에서 위치 임베딩을 할당한다는 것은 0,1,2,3,6,7,8,9의 위치 캐시 값을 가지고 있어야 하는 한 방법 아닌가?

그렇다면 각각 수의 위치 KV값이 이미 사전에 학습되어있다는 전제하게 위의 방식이 가는한거 아닌가?

 

 

3.3  Pretraining LLM with Attentino Sinks

 

섹션 3.1에서 상세히 설명한 바와 같이, 모델이 여러 초기 토큰들에 과도하게 주목하는 주요한 이유는 과도한 어텐션 점수를 덜어낼 수 있는 지정된 싱크 토큰의 부재 때문입니다. 이로 인해, 모델은 의도치 않게 전역적으로 보이는 토큰들, 주로 초기 토큰들을, 어텐션 싱크로 사용합니다. 잠재적인 해결책은 불필요한 어텐션 점수의 저장소 역할을 할 "싱크 토큰"이라고 표시되는 전역 학습 가능한 어텐션 싱크 토큰을 의도적으로 포함하는 것일 수 있습니다. 대안으로는 기존의 SoftMax 함수를 SoftMax-off-by-One(Miller, 2023)과 같은 변형으로 대체하는 것이 있습니다.

 

*싱크 토큰

 

기본 attention과 사전 훈련 시 0 토큰 또는 학습 가능한 sink 토큰을 추가한 경우의 비교. 안정적인 스트리밍 perplexity를 보장하려면 기본 모델은 여러 초기 토큰이 필요합니다. Zero Sink는 약간의 향상을 보이지만, 여전히 다른 초기 토큰이 필요합니다. 반면 학습 가능한 Sink Token으로 훈련된 모델은 sink 토큰만 추가해도 안정적인 스트리밍 perplexity를 유지합니다. 캐시 구성 x+y는 x개의 초기 토큰과 y개의 최신 토큰을 추가하는 것을 의미합니다.Perplexity는 PG19 테스트 세트의 첫 번째 샘플에서 평가되었습니다.

 

기존 SoftMax 대신 모든 문맥적 토큰에 대한 주의 점수 합산이 필요하지 않은 SoftMax1을 사용하면 효과적일 수 있습니다. SoftMax1은 주의 계산에서 Key와 Value가 모두 0인 토큰을 추가하는 것과 같습니다. 이를 "Zero Sink"라고 부릅니다.

모델 성능 검증을 위해, 1억 6천만 개의 파라미터를 가진 세 가지 언어 모델을 동일한 설정으로 사전 학습했습니다. 첫 번째는 표준 SoftMax attention(기본), 두 번째는 SoftMax1(Zero Sink), 세 번째는 학습 가능한 placeholder 토큰(Sink Token)을 추가한 모델입니다. Table 3에서 보듯, Zero Sink는 일부 개선을 보이지만 여전히 다른 초기 토큰에 의존합니다. Sink Token 도입은 attention 메커니즘을 안정화하는 데 매우 효과적이며, 이 Sink Token과 최신 토큰을 함께 사용하면 모델 성능이 안정적으로 유지됩니다.

 

---

4. Experiments

 

우리는 최근 주목받는 네 가지 모델 계열(Llama-2, MPT, PyThia, Falcon)로 StreamingLLM을 평가했습니다. 이 중 Llama-2, Falcon, PyThia는 RoPE 위치 인코딩을, MPT는 ALiBi를 사용합니다. 다양한 모델 선택은 연구 결과의 타당성과 강인성을 보장합니다. 우리는 StreamingLLM을 기존의 Dense Attention, Window Attention, Sliding Window 방식과 비교했습니다. 이후 실험에서 별도 언급이 없는 한, StreamingLLM은 기본적으로 네 개의 초기 토큰을 Attention Sink로 사용합니다.

 

4.1 Language Modeling On Long Texts across LLM Families and Scales

우리는 PG19 테스트 세트(긴 책 100권으로 구성)를 사용하여 StreamingLLM의 언어 모델링 perplexity를 평가했습니다. Llama-2 모델의 캐시 크기는 2048, Falcon, Pythia, MPT 모델은 1024로 설정하였으며, 이는 시각적 명확성을 위해 사전 학습 창 크기의 절반으로 설정된 것입니다.

 

Figure 3에서 StreamingLLM은 sliding window with re-computation 방식의 oracle baseline과 동등한 perplexity 성능을 20K 토큰 길이의 텍스트에서도 유지하는 것을 보여줍니다. 반면, Dense Attention은 입력 길이가 사전 학습 창을 초과하면 성능이 떨어지며, Window Attention은 캐시 크기를 초과할 때 초기 토큰이 제거되어 성능이 저하됩니다.

 

Figure 5에서는 StreamingLLM이 Llama-2-[7,13,70]B, Falcon-[7,40]B, Pythia-[2.8,6.9,12]B, MPT-[7,30]B 등 다양한 모델 계열과 규모에 걸쳐 400만 개 이상의 토큰으로 이루어진 매우 긴 텍스트도 안정적으로 처리할 수 있음을 보여줍니다.

 

Figure 5

 

Figure 5: 다양한 LLM 계열과 규모로 구성된 400만 개의 토큰으로 이루어진 초장기 텍스트에서 StreamingLLM의 언어 모델링 perplexity를 보여줍니다. perplexity는 전체 텍스트에서 안정적으로 유지됩니다. PG19의 100권의 책을 연결한 테스트 세트를 사용하여 언어 모델링을 수행했으며, 책이 전환되는 지점에서 perplexity의 약간의 변동이 발생합니다.

 

Figure 7

 

Figure 7: 256개의 문장(각 16개 토큰 길이)에 대해 평균 attention logit을 시각화하여, sink token 없이(pre-trained, 왼쪽)와 sink token을 사용해 사전 훈련된 경우(오른쪽)를 비교합니다. 두 경우 모두 동일한 레이어와 헤드를 보여줍니다. 주요 관찰 사항은 다음과 같습니다:

(1) sink token이 없을 경우, 낮은 레이어에서 국소적 attention을 보이며, 깊은 레이어에서는 초기 토큰에 대한 attention이 증가합니다. (2) sink token이 있을 경우, 모든 레이어에서 이를 향한 명확한 attention이 나타나며, 불필요한 attention을 모읍니다.

(3) sink token이 있을 경우 다른 초기 토큰에 대한 attention이 감소하여, sink token을 지정하는 것이 스트리밍 성능 향상에 유리함을 보여줍니다.

*attention logit

 

 

4.2 Result of Pre-Training with Sink Tokens

 

스트리밍 LLM의 개선을 검증하기 위해, 우리는 매 샘플 시작 부분에 sink token을 추가한 모델과 기존 설정을 따른 모델을 비교하며, 각각 1억 6천만 개의 매개변수를 가진 언어 모델 두 개를 동일한 조건에서 훈련했습니다. Pythia-160M 코드베이스를 사용해 NVIDIA A6000 GPU 8대로 구성된 서버에서, 중복이 제거된 Pile 데이터셋으로 훈련을 진행했습니다(Gao et al., 2020). 배치 크기를 256으로 줄인 것을 제외하고는 Pythia의 모든 학습 설정을 유지하며, 두 모델을 143,000 스텝 동안 훈련했습니다.

 

Figure 6: 모델의 사전 훈련 손실 곡선. sink token을 포함한 모델과 포함하지 않은 모델 모두 유사한 수렴 경향을 보입니다.

 

Table 4

 

Table 4: 7개의 NLP 벤치마크에서의 제로샷 정확도(%) 비교. 이 벤치마크에는 ARC-[Challenge, Easy], HellaSwag, LAMBADA, OpenbookQA, PIQA, Winogrande가 포함됩니다. 사전 훈련 시 sink token을 추가해도 모델 성능에 부정적인 영향을 주지 않습니다.

 

 

수렴 및 일반 모델 성능. 사전 훈련에 sink token을 포함해도 모델의 수렴과 여러 NLP 벤치마크에서의 성능에는 부정적인 영향이 없습니다. Figure 6에서 볼 수 있듯이 sink token이 포함된 모델은 일반 모델과 유사한 수렴 경향을 보입니다. ARC, HellaSwag, LAMBADA 등 7개 NLP 벤치마크에서 평가한 결과(Table 4 참조), sink token이 포함된 모델의 성능은 기존 방식으로 훈련된 모델과 유사하게 나타났습니다.

 

스트리밍 성능. Table 3에 따르면, 기존 방식과 sink token을 사용한 모델의 스트리밍 perplexity가 다릅니다. 일반 모델은 안정적인 perplexity를 유지하기 위해 다수의 토큰이 필요하지만, sink token을 사용한 모델은 단일 토큰으로도 만족스러운 성능을 달성합니다.

 

Table 5

 

Table 5: ARC-[Easy, Challenge] 데이터셋에서의 정확도(%)입니다.

질문들은 실시간 채팅 환경을 모방하기 위해 이어서 연결되어 답변되었습니다.

Dense baseline은 Out-of-Memory (OOM) 오류로 실패합니다.

Window attention은 정확도가 낮습니다.

StreamingLLM은 샘플별로 한 번에 처리하는 baseline과 비슷한 결과를 보입니다.

Window attention과 StreamingLLM은 캐시 크기 1024를 사용합니다.

 

Figure 9 : 샘플 데이터 100개 이상의 평균 정확도. SteamEval benchmark의 수행 결과

 

 

Attention Visualization. 그림 7은 sink 토큰이 있는 모델과 없는 모델의 어텐션 맵을 비교합니다. sink 토큰이 없는 모델은 Llama-2-7B(그림 2)와 유사하게 초기 레이어에서는 로컬 어텐션을, 더 깊은 레이어에서는 초기 토큰에 집중하는 모습을 보입니다. 반면, sink 토큰으로 훈련된 모델은 레이어와 헤드에 걸쳐 일관되게 sink에 집중하며, 이는 효과적인 어텐션 오프로드 메커니즘을 나타냅니다. 초기 토큰에 대한 어텐션이 감소하고 sink에 대한 집중이 강하게 나타나는 이 특성은 sink 토큰이 모델의 스트리밍 성능을 향상시키는 이유를 설명합니다.

 

 

 

4.3 Results on Streaming Question answering with Instruction-Tuned Models

 

StreamingLLM의 실제 적용 가능성을 보여주기 위해, 우리는 실제 시나리오에서 일반적으로 사용되는 지침 튜닝된 LLM을 사용하여 다중 라운드 질문-답변을 시뮬레이션합니다. 먼저 ARC-[Challenge, Easy] 데이터셋에서 모든 질문-답변 쌍을 이어붙여 Llama-2-[7,13,70]B-Chat 모델에 연속적인 스트림을 입력하고, 각 답변 위치에서 모델의 완료 결과를 정확히 일치하는 기준으로 평가합니다. 표 5에 나타난 것처럼, dense attention은 Out-of-Memory (OOM) 오류를 일으키며 이 설정에 적합하지 않음을 보여줍니다. window attention 방법은 효율적으로 작동하지만, 입력 길이가 캐시 크기를 초과할 때 무작위 출력으로 인해 정확도가 낮습니다. 반면, StreamingLLM은 스트리밍 형식을 효율적으로 처리하여 샘플별로 한 번에 처리하는 baseline 정확도와 일치하는 뛰어난 성능을 보입니다.

 

StreamingLLM에 더 적합한 시나리오를 강조하기 위해, 우리는 LongEval (Li et al., 2023) 벤치마크에서 영감을 받은 StreamEval 데이터셋을 소개합니다. 그림 8에 나타난 것처럼, LongEval의 긴 범위에서 단일 쿼리를 사용하는 방식과 달리, 우리는 새로운 정보의 10줄마다 모델에 쿼리를 던집니다. 각 쿼리의 답변은 항상 20줄 이전의 내용을 반영하며, 이는 질문이 최근 정보와 관련된 실제 사례를 나타냅니다.

 

 

그림 9에서 보듯이, StreamingLLM을 사용하는 LLM은 입력 길이가 120K 토큰에 접근해도 합리적인 정확도를 유지합니다. 반면, dense attention과 window attention은 각각 사전 훈련된 텍스트 길이와 KV 캐시 크기에서 실패합니다. 또한, 우리는 두 개의 컨텍스트 확장 모델인 LongChat-7b-v1.5-32k (Li et al., 2023)와 Llama-2-7B-32K-Instruct (Together, 2023)을 사용하여 StreamingLLM이 컨텍스트 확장 기법을 보완할 수 있음을 보여줍니다. StreamingLLM에서의 컨텍스트 확장은 스트리밍 LLM의 최대 캐시 크기를 확장하여 더 넓은 로컬 정보를 포착할 수 있게 하는 것을 의미합니다.

 

Figure 10

Figure 10: 슬라이딩 윈도우 접근법과 재계산 기준선, StreamingLLM 간의 토큰당 디코딩 지연 시간과 메모리 사용량을 캐시 크기(어텐션 윈도우 크기)에 대해 비교한 그래프입니다. StreamingLLM은 토큰당 최대 22.2배의 속도 향상을 보여주며, 재계산 기준선과 유사한 메모리 사용량을 유지합니다.

 

4.4 Abalation Studies

초기 토큰 수. 표 2에서는 최근 토큰과 함께 다양한 초기 토큰 수를 추가하는 것이 스트리밍 당혹도(perplexity)에 미치는 영향을 분석합니다. 그 결과, 하나 또는 두 개의 초기 토큰만 추가하는 것은 충분하지 않으며, 초기 토큰 4개를 추가하는 것이 적절한 임계값으로 보이고, 그 이후의 추가는 미미한 영향을 미친다고 나타났습니다. 이 결과는 StreamingLLM에서 4개의 초기 토큰을 어텐션 싱크로 도입한 선택을 정당화합니다.

 

캐시 크기. 표 6에서는 캐시 크기가 StreamingLLM의 당혹도에 미치는 영향을 평가합니다. 직관과는 달리, 캐시 크기를 늘려도 언어 모델링 당혹도가 일관되게 낮아지지 않습니다. 이러한 불일치는 이러한 모델들이 받은 전체 컨텍스트를 최대한 활용하지 못할 수 있는 잠재적인 한계를 보여줍니다. 향후 연구는 이 모델들이 광범위한 컨텍스트를 더 잘 활용할 수 있도록 개선하는 방향으로 진행되어야 할 것입니다.

 

4.5 Effieiency Results

 

우리는 StreamingLLM의 디코딩 지연 시간과 메모리 사용량을 슬라이딩 윈도우 재계산 방식과 비교합니다. 이 방법은 품질이 허용 가능한 유일한 기준선입니다. 두 방법 모두 Huggingface Transformers 라이브러리(Wolf et al., 2020)를 사용하여 구현되었으며, NVIDIA A6000 GPU 한 대에서 Llama-2-7B 및 Llama-2-13B 모델로 테스트되었습니다. 그림 10에 나타난 것처럼, 캐시 크기가 증가함에 따라 StreamingLLM의 디코딩 속도는 선형적으로 증가합니다. 반면, 슬라이딩 윈도우 재계산 기준선은 디코딩 지연 시간이 이차적으로 증가합니다. 결과적으로 StreamingLLM은 토큰당 최대 22.2배의 인상적인 속도 향상을 달성합니다. 지연 시간이 줄어듦에도 불구하고, StreamingLLM은 재계산 기준선과 일관된 메모리 사용량을 유지합니다.

 

표 6: 캐시 크기가 StreamingLLM 성능에 미치는 영향. StreamingLLM에서 캐시 크기를 증가시켜도 당혹도(perplexity)가 일관되게 감소하지 않으며, 이는 이 모델들이 제공된 컨텍스트를 완전히 활용하지 못할 수 있음을 나타냅니다. 캐시 구성 x+y는 x개의 초기 토큰과 y개의 최근 토큰을 추가하는 것을 의미합니다. 당혹도는 결합된 PG19 테스트 세트의 400K 토큰을 기준으로 평가되었습니다.

 

 

 

 

---

 

5. Conclusion

 

스트리밍 애플리케이션에 LLM을 배포하는 것은 시급한 과제지만, 긴 텍스트에 대한 성능 저하와 효율성 문제로 인해 어려움이 따릅니다. Window attention은 부분적인 해결책을 제공하지만, 초기 토큰이 제외될 때 성능이 급격히 떨어집니다. 이러한 토큰이 “attention sinks” 역할을 한다는 점을 인식하여, 우리는 "StreamingLLM" 을 도입했습니다. 이 프레임워크는 간단하고 효율적이며, LLM이 별도의 파인 튜닝 없이 무제한의 텍스트를 처리할 수 있게 합니다. 최근 토큰과 함께 attention sinks를 추가하여, StreamingLLM은 최대 400만 토큰의 텍스트도 효율적으로 모델링할 수 있습니다. 또한, 전용 sink 토큰으로 모델을 사전 훈련하면 스트리밍 성능을 향상시킬 수 있음을 보였습니다. StreamingLLM은 LLM의 사전 훈련 윈도우 크기와 실제 텍스트 생성 길이를 분리하여, LLM의 스트리밍 배포를 위한 길을 열어줍니다.