https://www.dropbox.com/scl/fi/v1sl9gt0gjjchtq2z3d8l/iclr24_StreamingLLM_Unbound.pdf?rlkey=ekljkgijkqzzb2udnvcbnwbvc&dl=0

Efficient Streaming Language Models with Attention Sinks는 ICLR 2024에 발표된 논문으로, 장시간 대화나 지속적인 텍스트 스트림에서 LLM을 고정된 메모리와 안정적인 품질로 계속 실행하기 위한 StreamingLLM을 제안합니다. 핵심 발견은 Transformer가 의미적으로 중요하지 않은 초기 몇 개 토큰에 비정상적으로 큰 attention을 할당한다는 것이며, 저자들은 이를 attention sink라고 부릅니다.

논문의 핵심 아이디어는 매우 단순합니다.

최근 토큰만 남기는 일반적인 sliding-window KV cache와 달리,
문장의 최초 몇 개 토큰의 KV cache는 영구적으로 보존하고, 나머지 공간에는 가장 최근 토큰들의 KV만 유지한다.

이 방법만으로 별도 fine-tuning 없이 Llama-2, MPT, Falcon, Pythia가 수백만 토큰의 스트림에서 안정적인 perplexity를 유지했으며, sliding-window recomputation보다 최대 22.2배 빠른 decoding을 보입니다.

1. 연구 배경과 문제 정의

1.1 일반적인 autoregressive decoding의 KV cache

Transformer의 각 attention layer에서 과거 토큰 $x_1,\ldots,x_{t-1}$ 의 key와 value를 캐싱하면, 새로운 토큰 $x_t$ 를 생성할 때 과거 토큰을 다시 계산하지 않아도 됩니다.

각 layer의 cache는 대략 다음과 같습니다. $K_{1:t-1}=[k_1,\ldots,k_{t-1}],\qquad V_{1:t-1}=[v_1,\ldots,v_{t-1}]$

하지만 전체 과거 토큰을 보존하는 dense attention에서는 스트림 길이 T가 증가할수록 KV cache 메모리가 O(T)로 증가하고, 전체 생성 비용은 O(T^2)에 가까워집니다.

장시간 대화처럼 입력이 계속 추가되는 경우에는 다음 두 문제가 동시에 발생합니다.

KV cache 메모리의 지속적 증가
훈련 context length를 넘었을 때 모델 품질 붕괴

예를 들어 Llama-2가 4K context로 학습되었다면, 단순히 4K보다 긴 position을 넣는 것만으로는 안정적인 추론이 보장되지 않습니다. 논문은 이러한 문제를 “긴 문서 전체를 한 번에 이해하는 문제”라기보다는, 무한히 지속되는 스트림에서 최근 문맥을 안정적으로 사용하는 문제로 정의합니다.

1.2 세 가지 기존 inference 방식

논문의 Figure 1은 네 가지 방식을 비교합니다.

① Dense attention

모든 과거 토큰의 KV를 보관합니다.

$\mathcal{C}_t=\{KV_1,\ldots,KV_{t-1}\}$

장점: 훈련 길이 안에서는 모든 문맥에 접근 가능
단점:
- cache 크기가 계속 증가
- 긴 입력에서는 메모리 부족
- 훈련 길이를 넘으면 positional extrapolation 문제 발생
총 계산량: 대략 O(T^2)

Llama-2-13B로 65K 토큰 PG-19 텍스트를 평가한 예에서는 perplexity가 약 5641까지 붕괴합니다.

② Window attention

가장 최근 L개의 KV만 유지합니다.

$\mathcal{C}_t= \{KV_{t-L},\ldots,KV_{t-1}\}$

cache 크기: O(L)
전체 decoding 비용: O(TL)
메모리와 token당 latency는 일정

표면적으로는 가장 자연스러운 streaming 방식입니다. 그러나 cache가 가득 차서 최초 토큰이 eviction되는 순간 perplexity가 폭증합니다.

Figure 1의 Llama-2-13B 실험에서는 PPL이 5158 수준으로 무너집니다.

③ Sliding window with re-computation

각 새 토큰을 생성할 때마다 최근 L개의 원래 토큰을 다시 모델에 입력하여 KV를 재계산합니다.

이는 최근 window를 마치 독립적인 정상 입력처럼 처리하므로 성능은 좋습니다. 그러나 매 토큰마다 길이 L의 Transformer forward를 다시 수행해야 합니다.

$\text{total cost}\approx O(TL^2)$

성능: 안정적
메모리: 고정
속도: 매우 느림

논문에서는 이를 품질 측면의 사실상 oracle baseline으로 사용합니다.

④ StreamingLLM

StreamingLLM은 다음 두 종류의 KV를 유지합니다. $\mathcal{C}_t= \underbrace{\{KV_1,\ldots,KV_s\}}_{\text{attention sinks}} \cup \underbrace{\{KV_{t-w+1},\ldots,KV_t\}}_{\text{recent window}}$

여기서:

s: 보존할 초기 sink token 수, 기본적으로 4
w: 최근 토큰 수
전체 cache size L=s+w

따라서 메모리와 연산량은 window attention과 동일한 수준인 O(L), O(TL)이지만, 품질은 recomputation 방식에 가깝습니다.

2. 관련 연구

논문은 long-context LLM 연구를 세 가지 범주로 구분합니다.

2.1 Length extrapolation

Length extrapolation은 짧은 길이로 학습한 모델이 테스트 시 더 긴 sequence에 일반화하도록 만드는 연구입니다.

RoPE

RoPE는 query와 key를 position-dependent rotation으로 변환합니다. $q_i’=R(i)q_i,\qquad k_j’=R(j)k_j$

따라서 dot product가 상대 위치 i-j에 의존하게 됩니다.

$(q_i’)^\top k_j’ =q_i^\top R(j-i)k_j$

RoPE는 상대 위치 정보를 자연스럽게 표현하지만, 훈련에서 보지 못한 매우 큰 position index에서는 rotation phase가 훈련 분포를 벗어나므로 성능이 저하될 수 있습니다.

ALiBi

ALiBi는 attention logit에 거리에 비례하는 선형 bias를 추가합니다.

$a_{ij} = \frac{q_i^\top k_j}{\sqrt{d}} -m_h(i-j)$

ALiBi는 RoPE보다 length extrapolation이 양호할 수 있지만, 저자들의 MPT 실험에서는 훈련 길이보다 훨씬 긴 입력에서 역시 붕괴했습니다.

논문은 RoPE와 ALiBi 같은 positional encoding만으로는 사실상 무한 길이의 streaming을 달성하기 어렵다고 봅니다.

2.2 Context window extension

이 범주는 모델이 한 번에 attend할 수 있는 context 자체를 확장합니다.

대표적인 연구는 다음과 같습니다.

FlashAttention, FlashAttention-2
Longformer, BigBird, Reformer 같은 sparse/approximate attention
Position Interpolation
NTK-aware RoPE scaling
YaRN
long-context fine-tuning

이 방법들은 4K context를 32K 또는 128K로 늘릴 수 있지만, context는 여전히 유한합니다.

StreamingLLM은 context window를 확장하는 방법이 아닙니다. 예를 들어 cache가 4K라면 현재 token은 다음 정보만 볼 수 있습니다.

최초 sink token 몇 개
가장 최근 약 4K 토큰

그 사이에 eviction된 수십만 토큰은 볼 수 없습니다.

따라서 StreamingLLM이 “4M tokens를 처리한다”는 것은 4M 토큰 전체를 동시에 기억한다는 뜻이 아니라, 4M 토큰 동안 모델을 중단 없이 안정적으로 실행한다는 뜻입니다. 논문도 context extension과 streaming을 명시적으로 구분합니다.

2.3 Long-context utilization

Context window가 길어져도 모델이 모든 위치의 정보를 효과적으로 사용한다는 보장은 없습니다.

대표적인 문제는 Lost in the Middle입니다.

앞부분 정보는 비교적 잘 사용
끝부분 정보도 비교적 잘 사용
중간 정보의 활용률은 낮음

StreamingLLM 역시 long-context utilization 문제를 직접 해결하지 않습니다. 오히려 최근 문맥만 유지하기 때문에, 오래전에 등장한 정보는 별도 memory 또는 retrieval이 없으면 사용할 수 없습니다.

즉, 세 범주의 차이는 다음과 같습니다.

범주	목표	StreamingLLM과 관계
Length extrapolation	학습 길이보다 긴 스트림에서 안정성 유지	직접적 목표
Context extension	한 시점에 볼 수 있는 토큰 수 증가	별도 방법과 결합 가능
Long-context utilization	긴 context에서 필요한 정보를 실제로 활용	직접 해결하지 않음

2.4 Quantization outlier 및 “do-nothing attention”

저자들은 attention sink 현상을 quantization outlier 연구와 연결합니다.

일부 attention head는 실제로 가져올 정보가 없어도 softmax 때문에 반드시 어떤 value들의 가중합을 출력해야 합니다. 이때 특정 token이 사실상 “버리는 attention mass”를 받는 장소가 될 수 있습니다.

이 관점은 다음 연구들과 관련됩니다.

SmoothQuant
Quantizable Transformers
Softmax-off-by-One

특히 Softmax-off-by-One은 attention denominator에 상수 1을 추가합니다. $\operatorname{Softmax}_1(x)_i = \frac{e^{x_i}} {1+\sum_j e^{x_j}}$

이렇게 하면 실제 token에 할당되는 attention weight의 합이 1보다 작을 수 있습니다. $\sum_i \operatorname{Softmax}_1(x)_i < 1$

즉, attention head가 필요할 경우 일부 probability mass를 아무 token에도 할당하지 않는, 일종의 null attention을 구현할 수 있습니다.

3. Attention sink 현상

3.1 실험적 관찰

논문의 Figure 2에서는 Llama-2-7B에 길이 16인 문장 256개를 입력하고 layer/head별 평균 attention logit을 시각화합니다.

관찰 결과는 다음과 같습니다.

Layer 0–1: 최근 token에 집중하는 local attention
그보다 깊은 layer: 거의 모든 head에서 초기 token에 강한 attention
초기 token의 의미와 무관하게 비슷한 현상 발생

즉, 초기 token이 중요한 의미 정보를 제공해서가 아니라, attention computation을 안정화하는 구조적 기능을 수행한다는 주장입니다.

3.2 Window attention이 붕괴하는 직접적 이유

일반적인 attention은 다음과 같습니다.

$\alpha_{ij} = \frac{\exp(z_{ij})} {\sum_{k\le i}\exp(z_{ik})}, \qquad z_{ij}=\frac{q_i^\top k_j}{\sqrt{d}}$

초기 sink token의 logit이 매우 큰 경우를 생각해보면,

$z_{i1}\gg z_{ij},\qquad j>1$

정상적인 실행에서는 sink token이 denominator의 상당 부분을 차지합니다. $D_{\text{normal}} = \exp(z_{i1})+ \sum_{j=2}^{N}\exp(z_{ij})$

그런데 window attention이 초기 token을 제거하면 denominator가 다음처럼 변합니다. $D_{\text{evicted}} = \sum_{j=2}^{N}\exp(z_{ij})$

따라서 남은 token들의 attention weight가 비정상적으로 재정규화됩니다.

$\alpha’_{ij} = \frac{\exp(z_{ij})}{D_{\text{evicted}}}$

이는 단순히 “초기 token의 정보가 사라지는 것”보다 큰 문제입니다. 모든 남은 value에 대한 가중치 분포가 동시에 바뀌기 때문입니다.

이 변화가 layer마다 누적되면서 hidden state가 훈련 시의 activation distribution에서 벗어나고, 결국 perplexity가 폭증합니다. 저자들은 최초 token이 cache에서 제거되는 시점과 perplexity surge가 일치함을 보여줍니다.

3.3 왜 초기 token인가?

저자들의 설명은 causal autoregressive training의 visibility asymmetry입니다.

초기 token $x_1$ 은 거의 모든 후속 prediction에서 보입니다.

$x_1 \rightarrow x_2,x_3,\ldots,x_T$

반면 후반 token $x_j$ 는 j 이후의 token에서만 보입니다.

$x_j \rightarrow x_{j+1},\ldots,x_T$

따라서 초기 token은 학습 중 훨씬 많은 query와 상호작용하며, 모델이 일관된 sink로 활용하기 가장 쉬운 위치가 됩니다.

3.4 의미 때문인가, 위치 때문인가?

이를 검증하기 위해 저자들은 Llama-2-13B의 최초 네 token을 원래 token이 아닌 네 개의 newline token \n으로 교체합니다.

PG-19의 65K token book에서 결과는 다음과 같습니다.

Cache 구성	PPL
0+1024: 최근 1024개만	5158.07
4+1020: 원래 초기 4개 + 최근 1020개	5.40
$4_{\text{newline}}+1020$	5.60

원래 초기 token 대신 의미 없는 newline token을 넣어도 perplexity가 거의 완전히 복구됩니다. 따라서 핵심은 초기 token의 semantic content가 아니라 해당 위치에 sink 역할을 수행하는 KV가 존재하는 것입니다.

다만 엄밀하게 말하면, 임의의 새로운 newline token을 inference 중에 삽입한 것은 아닙니다. 원래 초기 위치의 token을 newline으로 바꾸어 모델을 forward하고, 그 결과로 생성된 초기 KV를 계속 보존한 것입니다.

4. 방법론: StreamingLLM

4.1 Rolling KV cache with attention sinks

전체 cache budget이 L, sink token 수가 s라면 다음을 보존합니다. $\mathcal{C}_t= \{1,\ldots,s\} \cup \{t-(L-s)+1,\ldots,t\}$

예를 들어 cache size가 8이고 s=4라면:

sink: token 0, 1, 2, 3
rolling cache: 가장 최근 4개 token

시간이 지나도 token 0–3은 eviction되지 않고, rolling 부분만 FIFO 방식으로 갱신됩니다.

논문의 5페이지 Figure 4는 cache가 다음처럼 이동하는 과정을 보여줍니다.

[0,1,2,3,4,5,6,7]

새 token이 들어오면 오래된 rolling token을 제거하여:

[0,1,2,3,5,6,7,8]

이후: [0,1,2,3,6,7,8,9]

와 같은 형태로 유지합니다. 초기 4개는 고정되고 최근 window만 이동합니다.

4.2 중요한 구현 세부사항: cache 내부 position 재할당

StreamingLLM에서 단순히 원래 absolute position을 유지해서는 안 됩니다.

예를 들어 실제 원문 position이 다음과 같다고 합시다.

[0,1,2,3,6,7,8]

그리고 token 9를 생성하려고 합니다.

원래 위치를 그대로 사용하면 query와 cache position은 다음과 같습니다.

[0,1,2,3,6,7,8,9]

하지만 StreamingLLM은 cache를 연속된 sequence처럼 재배치합니다.

[0,1,2,3,4,5,6,7]

즉:

sink token: position 0–3
rolling tokens: position 4–6
현재 query: position 7

이렇게 해야 모델이 훈련 중 보았던 짧고 연속적인 상대 거리 분포를 유지할 수 있습니다.

4.3 RoPE에서의 처리

RoPE는 key에 position-dependent rotation이 적용되므로, 이미 회전된 key를 cache하면 eviction 후 position을 다시 압축하기 어렵습니다.

일반적인 key는 다음과 같습니다.

$k_j^{\text{rope}}=R(p_j)k_j$

StreamingLLM에서는 rolling cache token의 position이 계속 바뀌므로, 논문은 RoPE 적용 전 key를 보존합니다.

$k_j^{\text{raw}}=W_Kh_j$

그 후 현재 cache 내 위치 $\tilde p_j$ 를 이용해 다시 회전합니다.

$k_j^{\text{stream}} = R(\tilde p_j)k_j^{\text{raw}}$

query에도 현재 cache 기준 position을 적용합니다.

$q_t^{\text{stream}} = R(\tilde p_t)W_Qh_t$

이 방식은 rolling cache의 각 token에 대해 key rotation을 다시 수행해야 하지만, 전체 Transformer layer를 재계산하는 것보다 훨씬 저렴합니다.

4.4 ALiBi에서의 처리

ALiBi는 key 자체를 회전하지 않고 attention score에 거리 bias를 더합니다. $z_{ij} = \frac{q_i^\top k_j}{\sqrt d} -m_h(i-j)$

따라서 cache 내에서 압축한 연속 position을 이용해 bias만 다시 계산하면 됩니다.

즉, 실제 원문 거리의 “점프”를 사용하지 않고 cache 내 contiguous distance를 사용합니다. $\tilde z_{ij} = \frac{q_i^\top k_j}{\sqrt d} -m_h(\tilde i-\tilde j)$

논문은 RoPE보다 ALiBi와의 결합이 구현상 더 직접적이라고 설명합니다.

4.5 계산 복잡도

전체 스트림 길이를 T, cache size를 L이라고 하면:

방법	KV 메모리	총 attention 계산량	품질
Dense attention	O(T)	O(T^2)	훈련 길이 초과 시 붕괴
Window attention	O(L)	O(TL)	sink eviction 후 붕괴
Sliding-window recomputation	O(L)	O(TL^2)	안정적
StreamingLLM	O(L)	O(TL)	안정적

StreamingLLM의 장점은 window attention과 동일한 asymptotic efficiency를 유지하면서, recomputation에 가까운 품질을 얻는다는 것입니다.

5. 전용 Sink Token을 사용한 pre-training

5.1 기존 모델이 여러 sink token을 요구하는 이유

Llama-2 같은 기존 모델에서는 초기 token 1개만 보존해서는 성능이 완전히 회복되지 않고, 일반적으로 약 4개가 필요합니다.

저자들은 그 이유를 훈련 sample마다 동일한 첫 token이 항상 존재하지 않았기 때문이라고 설명합니다.

Llama-2에 <s>가 존재하더라도, chunking 절차 때문에 모든 training chunk의 position 0이 동일한 <s>가 되는 것은 아닙니다. 따라서 모델은 특정 token 하나가 아니라 여러 초기 위치에 sink 역할을 분산시킵니다.

5.2 Learnable sink token

저자들은 모든 training sample 앞에 동일한 learnable placeholder token을 추가합니다. $[\texttt{<sink>},x_1,x_2,\ldots,x_n]$

이 token은 별도의 lexical meaning을 가질 필요가 없고, attention mass를 받아내는 전용 register 역할을 합니다.

훈련 loss는 일반 token에 대해 기존과 동일하게 계산하며, sink token 자체를 예측 대상으로 사용할 필요는 없습니다.

이렇게 학습하면 여러 초기 token을 보존할 필요 없이, inference에서 다음 cache만으로 안정적 streaming이 가능합니다.

$\{\texttt{<sink>}\} \cup \{\text{recent tokens}\}$

5.3 Zero Sink와 Softmax-off-by-One

논문은 세 가지 160M 모델을 비교합니다.

Vanilla
- 표준 softmax
Zero Sink
- Softmax-off-by-One
- all-zero key/value를 암묵적으로 하나 추가한 것과 동등
Learnable Sink
- 모든 sample 앞에 학습 가능한 sink token 추가

Softmax-off-by-One은 다음과 같습니다.

$\alpha_i = \frac{e^{z_i}} {1+\sum_j e^{z_j}}$

이를 key k_0=0, value v_0=0인 가상 token이 있다고 해석할 수 있습니다.

$q^\top k_0=0,\qquad e^{q^\top k_0}=1$

이 token의 value가 0이므로 해당 token에 attention을 할당해도 output에 실제 정보가 더해지지 않습니다.

5.4 Pre-training 비교 결과

PG-19에서 cache 구성별 PPL은 다음과 같습니다.

모델	0+1024	1+1023	2+1022	4+1020
Vanilla	27.87	18.49	18.05	18.05
Zero Sink	29214	19.90	18.27	18.01
Learnable Sink	1235	18.01	18.01	18.02

여기서 x+y는 초기 token x개와 최근 token y개를 보존한다는 의미입니다.

해석은 다음과 같습니다.

Vanilla는 한 개 sink만으로 충분하지 않음
Zero Sink도 다른 초기 token을 추가로 필요로 함
Learnable Sink는 한 개의 sink token만 유지해도 최적 수준의 PPL
sink token이 제거된 0+1024에서는 Learnable Sink 모델도 붕괴하므로, 모델이 실제로 해당 token에 강하게 의존하도록 학습된 것임

6. 실험 설정

6.1 모델

다양한 architecture와 positional encoding에 대한 일반성을 확인하기 위해 다음 모델을 사용합니다.

RoPE 계열

Llama-2: 7B, 13B, 70B
Falcon: 7B, 40B
Pythia: 약 2.8/2.9B, 6.9B, 12B

ALiBi 계열

MPT: 7B, 30B

기본 설정에서는 초기 4개 token을 sink로 보존합니다.

6.2 Baseline

Dense attention
Window attention
Sliding window with re-computation
StreamingLLM

Sliding-window recomputation은 속도는 느리지만 최근 window를 매번 정상적인 independent sequence로 다시 인코딩하므로 품질 oracle에 가깝습니다.

6.3 데이터와 평가

Long-text language modeling

PG-19 test set
100권의 긴 책을 연결
metric: token-level perplexity

Cache size:

Llama-2: 2048
Falcon, Pythia, MPT: 1024

이는 각 모델의 pretraining window의 절반 수준으로 설정되었습니다.

Streaming question answering

ARC-Easy
ARC-Challenge
StreamEval

Efficiency

Llama-2-7B, 13B
단일 NVIDIA A6000
Hugging Face Transformers 구현
cache size 256–4096
per-token latency와 GPU memory 측정

7. 실험 결과

7.1 Long-text language modeling

Figure 3에서는 약 20K token 문서에 대한 PPL 변화를 보여줍니다.

Dense attention

pretraining context length를 넘은 이후 품질 저하
RoPE/ALiBi의 length extrapolation 한계가 나타남

Window attention

cache가 채워질 때까지는 정상
최초 token이 eviction되는 순간 PPL 급증
단순히 오래된 의미 정보가 없어져서 서서히 나빠지는 것이 아니라, 명확한 phase transition 형태

Sliding-window recomputation

전체 구간에서 안정적인 PPL
매우 높은 계산 비용

StreamingLLM

recomputation baseline과 거의 동일한 PPL
window attention과 동일한 수준의 cache 크기 및 decoding complexity

논문은 다양한 모델에서 이 경향이 일관되게 나타난다고 보고합니다.

7.2 4 million token streaming

PG-19의 100개 책을 연결하여 4M token 이상에서 language modeling을 수행합니다.

평가 모델은 다음과 같습니다.

Llama-2-7B/13B/70B
Falcon-7B/40B
Pythia-2.8B/6.9B/12B
MPT-7B/30B

StreamingLLM은 전체 스트림에서 PPL이 안정적으로 유지됩니다. 책 경계에서 PPL이 변동하지만, sequence length가 증가함에 따라 발산하는 현상은 없습니다.

그러나 이 결과는 다음과 같이 해석해야 합니다.

4M token 전부를 context로 사용한 것이 아니라,
4M token 동안 초기 sink와 제한된 최근 window만 보면서 안정적으로 language modeling한 것이다.

따라서 이는 4M-token context understanding 결과가 아닙니다.

7.3 Sink token 개수 ablation

여러 모델에서 초기 token 수를 바꾼 결과입니다.

모델	Window	1 sink	2 sinks	4 sinks	8 sinks
Falcon-7B	17.90	12.12	12.12	12.12	12.12
MPT-7B	460.29	14.99	15.00	14.99	14.98
Pythia-12B	21.62	11.95	12.09	12.09	12.02
Llama-2-7B	3359.95	11.88	10.51	9.59	9.54

주요 해석:

모델에 따라 1개로 충분한 경우도 있음
Llama-2에서는 1–2개로 완전 복구되지 않음
4개에서 대부분 포화
8개로 늘려도 추가 개선은 미미

따라서 4개는 모든 모델에 적용하기 위한 보수적 default입니다.

7.4 Sink-token pretraining의 일반 성능

160M parameter 모델 두 개를 동일한 조건으로 학습합니다.

Pythia-160M recipe
deduplicated Pile
8× NVIDIA A6000
batch size 256
143K training steps
Vanilla vs Learnable Sink Token

Training loss curve는 거의 동일합니다.

7개 zero-shot benchmark 결과는 다음과 같습니다.

모델	ARC-C	ARC-E	HellaSwag	LAMBADA	OBQA	PIQA	WinoGrande
Vanilla	18.6	45.2	29.4	39.6	16.0	62.2	50.1
+Sink	19.6	45.6	29.8	39.9	16.6	62.6	50.8

Sink token을 추가해도 일반적인 downstream 성능이나 convergence를 훼손하지 않았고, 수치상으로는 모든 benchmark에서 소폭 상승했습니다. 다만 차이가 작고 통계적 유의성은 제시되지 않았으므로, 성능 향상보다는 성능 손상이 없었다고 해석하는 것이 타당합니다.

7.5 Attention visualization

논문의 7페이지 Figure 7은 sink token이 없는 모델과 있는 모델의 attention map을 비교합니다.

Vanilla model

아래 layer: local attention
깊은 layer: 여러 초기 token에 attention 집중

Sink-token model

모든 layer/head에서 전용 sink token에 명확한 attention 집중
다른 초기 token이 받는 attention은 감소

즉, sink token이 단순히 streaming PPL을 개선한 것뿐 아니라, 실제로 redundant attention mass를 집중적으로 흡수하는 구조적 역할을 학습했음을 보여줍니다.

7.6 ARC streaming QA

ARC-Easy와 ARC-Challenge의 모든 question-answer pair를 하나의 긴 스트림으로 연결합니다. 각 answer 위치에서 정확도를 측정합니다.

방법	Llama-2-7B-Chat ARC-E/C	13B ARC-E/C	70B ARC-E/C
One-shot	71.25 / 53.16	78.16 / 63.31	91.29 / 78.50
Dense	OOM	OOM	OOM
Window	3.58 / 1.39	0.25 / 0.34	0.12 / 0.32
StreamingLLM	71.34 / 55.03	80.89 / 65.61	91.37 / 80.20

해석:

Dense attention은 전체 stream의 KV를 저장하지 못해 OOM
Window attention은 sink가 eviction된 후 거의 random 수준으로 붕괴
StreamingLLM은 각 sample을 독립적으로 처리한 one-shot baseline과 동등하거나 약간 높은 정확도

StreamingLLM이 one-shot보다 일부 높게 나온 것은 방법론적으로 핵심적인 개선이라기보다 prompt/history 차이와 평가 variance에 따른 결과일 가능성이 큽니다.

7.7 StreamEval

기존 LongEval은 긴 document를 제공한 뒤 한 번 질문하는 방식입니다. 저자들은 streaming setting에 더 적합한 StreamEval을 구성합니다.

구조는 다음과 같습니다.

line별로 key-value 형태의 내용을 계속 입력
10개 line마다 query
정답은 현재 위치에서 20 line 이전에 등장
전체 stream이 최대 약 120K tokens까지 증가

따라서 StreamEval은 임의로 오래된 정보를 요구하는 것이 아니라, 최근 window 안의 정보를 반복적으로 검색하는 능력을 평가합니다.

결과:

StreamingLLM은 120K token까지 합리적인 정확도 유지
Dense attention은 pretraining context length 부근에서 실패
Window attention은 cache size를 넘은 뒤 실패
LongChat-32K와 Llama-2-7B-32K-Instruct에도 StreamingLLM을 결합 가능
context-extended model을 사용하면 rolling window 자체를 크게 하여 더 넓은 최근 문맥을 유지할 수 있음

7.8 Cache size ablation

Cache size를 늘리면 더 많은 최근 context를 볼 수 있으므로 PPL이 지속적으로 감소할 것처럼 보이지만, 실제 결과는 그렇지 않습니다.

모델	작은 cache	중간 cache	큰 cache	최대 cache
Falcon-7B	13.61	12.84	12.34	12.84
MPT-7B	14.12	14.25	14.33	14.99
Pythia-12B	13.17	12.52	12.08	12.09
Llama-2-7B	9.73	9.32	9.08	9.59

Cache를 늘려도 단조롭게 개선되지 않습니다.

이는 다음 가능성을 시사합니다.

모델이 긴 context를 효율적으로 활용하지 못함
더 오래된 local context가 prediction에 noise를 제공
positional distribution이 pretraining과 달라짐
PG-19의 다음-token prediction에서는 매우 긴 local context가 반드시 유용하지 않음

이 결과는 StreamingLLM이 context utilization 문제를 해결하지 않는다는 점을 뒷받침합니다.

7.9 Efficiency

단일 A6000에서 Llama-2-7B와 13B를 사용하여 cache size별 per-token latency를 비교합니다.

Llama-2-7B 예시

대략적인 Figure 10 수치는 다음과 같습니다.

Cache size	StreamingLLM	Recomputation
256	31 ms	약 63 ms
512	31 ms	약 103 ms
1024	35 ms	약 223 ms
2048	45 ms	약 523 ms
4096	65 ms	약 1411 ms

4096 cache에서: $\frac{1411}{65}\approx 21.7$

모델과 설정 전체에서 최대 22.2× speedup을 보고합니다.

Scaling 차이

StreamingLLM: $\text{latency}\propto L$

Recomputation: $\text{latency}\propto L^2$

메모리는 두 방식이 비슷합니다. 둘 다 cache 또는 recomputation window의 길이가 L로 제한되기 때문입니다.

논문은 StreamingLLM의 속도 증가가 단순한 kernel optimization이 아니라, 최근 window 전체의 Transformer computation을 반복하지 않는 알고리즘적 복잡도 차이에서 발생함을 강조합니다.

8. 핵심 공헌

8.1 Attention sink 현상의 체계적 규명

초기 token에 대한 높은 attention이 semantic importance가 아니라, softmax attention이 불필요한 probability mass를 처리하는 방식과 관련될 수 있음을 보여줍니다.

특히 다음 증거를 결합합니다.

layer/head attention visualization
최초 token eviction 시점의 PPL surge
초기 token을 newline으로 바꾸어도 성능 회복
sink token 수 ablation
전용 sink token pretraining 후 attention 집중 변화

8.2 별도 학습 없는 streaming inference

기존 pretrained model의 weight를 수정하지 않고 KV cache management와 positional indexing만 바꾸어 적용할 수 있습니다.

이는 다음 측면에서 실용적입니다.

fine-tuning 불필요
모델 architecture 변경이 작음
RoPE와 ALiBi 모두 지원
cache memory가 sequence length와 무관하게 고정
기존 serving framework에 통합 가능

8.3 Context length와 generation length의 분리

일반적으로 context length와 생성 가능한 전체 sequence length가 동일시되기 쉽습니다.

StreamingLLM은 이를 분리합니다.

Attention span: 한 token이 실제로 참고할 수 있는 범위
Total stream length: 모델이 중단 없이 처리할 수 있는 전체 길이

StreamingLLM에서는 attention span이 고정되어도 total stream length는 사실상 무제한일 수 있습니다.

8.4 Pretraining 설계에 대한 시사점

모든 sample 앞에 전용 sink/register token을 넣으면, 모델이 임의의 lexical token 여러 개를 sink로 사용하는 대신 하나의 명시적인 null-like token으로 attention을 집중시킬 수 있습니다.

이는 이후의 다음 연구들과도 개념적으로 연결됩니다.

ViT register tokens
attention outlier suppression
quantization-friendly Transformer
null attention
massive activation 및 outlier feature 연구

9. 한계와 비판적 분석

9.1 진정한 long-term memory가 아니다

가장 중요한 한계입니다.

StreamingLLM은 eviction된 과거 token의 KV를 완전히 삭제합니다. 따라서 오래전 정보가 최근 window에 다시 등장하지 않는 한 접근할 수 없습니다.

예를 들어 100K token 전의 사용자가 이름을 말했고 현재 cache가 4K라면, 별도 memory system이 없을 경우 그 이름을 회상할 수 없습니다.

따라서 다음 표현은 구분해야 합니다.

“4M tokens 동안 안정적으로 실행 가능”: 맞음
“4M-token context를 이해하거나 검색 가능”: 아님

9.2 StreamEval이 방법에 유리하게 설계됨

StreamEval에서는 정답이 항상 20 line 이전에 위치합니다. 이는 rolling window 안에 정답이 있도록 의도된 설정입니다.

따라서 이 실험은 다음을 잘 평가합니다.

장시간 실행 중 최근 정보 접근
cache가 계속 이동할 때의 안정성

그러나 다음은 평가하지 않습니다.

수만 token 전 정보 recall
long-range multi-hop reasoning
전역 document summarization
멀리 떨어진 evidence 통합

따라서 일반적인 long-context benchmark 성능으로 확장 해석해서는 안 됩니다.

9.3 Attention weight가 causal importance를 의미하지는 않음

논문은 초기 token에 높은 attention이 있다는 사실과 해당 KV를 제거했을 때 모델이 붕괴한다는 사실을 보여줍니다. 이는 강한 causal evidence입니다.

다만 다음은 완전히 분리되어 있지 않습니다.

key가 softmax denominator를 안정화하는 역할
value가 residual stream에 전달하는 역할
특정 layer/head에서만 중요한지
sink token의 key와 value 중 어느 쪽이 핵심인지

보다 정밀한 실험은 다음과 같을 수 있습니다.

sink K만 보존, V는 0으로 설정
sink V만 보존, K는 대체
head별 sink ablation
layer별 sink removal
attention-logit patching
sink token의 residual contribution 분석

논문의 Softmax 설명상으로는 주된 기능이 key/logit을 통한 denominator 안정화일 가능성이 크지만, 기존 모델에서는 value contribution도 배제할 수 없습니다.

9.4 “Softmax가 합이 1이어서 sink가 생긴다”는 설명의 불완전성

Softmax normalization은 sink 발생의 필요조건에 가까운 설명이지만, 충분조건은 아닙니다.

모든 attention model이 반드시 초기 token에 극단적인 sink를 형성해야 하는 것은 아닙니다. 실제 sink 형성에는 다음 요인이 함께 작용할 수 있습니다.

causal visibility asymmetry
BOS 및 document-boundary token 분포
positional encoding
residual stream과 layer normalization
attention output projection
head specialization
training data chunking 방식
model이 “no-op” attention을 구현할 수 있는 다른 경로의 부재

따라서 더 정확한 해석은 다음과 같습니다.

Softmax가 각 head의 attention mass를 1로 정규화하기 때문에 no-op 경로가 필요하며, causal training에서 모든 후속 위치에 노출되는 초기 token이 그 역할을 학습하기 쉬워 attention sink가 형성된다.

9.5 Sink token pretraining 실험의 규모

전용 sink token 학습은 160M 모델에서만 검증됩니다.

이는 개념 검증으로는 충분하지만, 다음은 확인되지 않았습니다.

7B–70B scale에서도 하나의 sink로 충분한가
instruction tuning과 RLHF 후에도 sink 역할이 유지되는가
GQA/MQA architecture에서도 동일한가
최신 long-context model에서도 효과가 있는가
sink token이 quantization이나 pruning에 어떤 영향을 주는가

9.6 Perplexity 중심 평가

PG-19 next-token PPL은 streaming stability를 평가하기에는 적절하지만, 실제 응용 능력을 충분히 평가하지 못합니다.

추가로 필요한 평가는 다음과 같습니다.

장기 대화 consistency
tool-use agent trajectory
streaming summarization
entity/state tracking
continual code completion
multi-document memory
retrieval-augmented streaming

ARC streaming과 StreamEval을 포함했지만, 대부분 최근 문맥 기반의 비교적 단순한 retrieval/QA입니다.

10. 이 논문의 가장 중요한 해석

이 논문의 핵심은 “긴 context를 이해하는 새로운 attention”이라기보다 다음과 같습니다.

기존 LLM의 sliding-window inference가 실패하는 원인은 단순히 오래된 정보가 제거되기 때문만이 아니다.
모델이 softmax attention의 안정화를 위해 사용하는 초기 sink token까지 제거하기 때문에 activation distribution이 붕괴한다.

따라서 다음 두 기능을 분리합니다.

Semantic context
- 최근 token들이 담당
Attention normalization anchor
- 초기 sink token들이 담당

StreamingLLM은 이 두 기능에 서로 다른 cache 정책을 적용합니다.

$\text{KV cache} = \text{permanent normalization anchors} + \text{temporary semantic context}$

이 단순한 구조가 논문의 가장 독창적인 통찰입니다.

11. 요약 평가

장점

매우 단순하지만 효과가 큰 발견과 알고리즘
별도 fine-tuning 없이 기존 모델에 적용 가능
RoPE와 ALiBi 모델에서 폭넓게 검증
attention visualization, token replacement, ablation, pretraining 실험을 통한 다각적 근거
window attention 수준의 계산량으로 recomputation 수준의 품질
최대 22.2×의 실질적인 speedup
serving system에 직접 적용 가능한 높은 실용성

한계

전체 과거를 기억하는 long-context 방법은 아님
StreamEval이 recent-context streaming에 편향됨
sink의 layer/head별 causal mechanism은 충분히 규명되지 않음
전용 sink-token pretraining은 160M 규모에서만 검증
softmax 기반 설명은 직관적이지만 완전한 mechanistic explanation은 아님

한 문장으로 정리

StreamingLLM은 최초 몇 개 token을 영구적인 attention normalization anchor로 보존하고 최근 token만 rolling cache로 관리함으로써, 유한한 context window를 가진 LLM을 고정 메모리로 무한히 오래 실행할 수 있게 하는 inference 방법입니다.

** Efficient Streaming Language Models with Attention Sinks (ICLR 2024)