SepLLM: Accelerate Large Language Models by Compressing One Segment into One Separator

🔼 이 그림은 ‘Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. …‘라는 입력이 주어졌을 때, 여러 층에 걸친 어텐션 점수를 시각화한 것입니다. 쉼표(’,’)와 마침표(’.’)와 같은 구분자 토큰들이 상당히 높은 어텐션 점수를 받는 것을 확인할 수 있는데, 이는 구분자들이 문맥 정보를 효율적으로 압축하고 전달하는 역할을 한다는 것을 시사합니다.

read the caption

Figure 2: The visualization of attention scores among different layers given the input “Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. …”. Note that the separator tokens like “,” and “.” contribute massive attentions.

🔼 이 그림은 SepLLM의 전반적인 패러다임을 보여줍니다. 왼쪽은 ‘ABC,DE.FG\n’ 입력에 대한 훈련 또는 사전 채우기 단계의 어텐션 마스크를 나타냅니다. 오른쪽은 생성 단계에서의 KV 캐시 관리를 보여줍니다. SepLLM은 특정 레이어에서 각 토큰이 이전 레이어에서 출력된 이전 토큰의 hidden state의 일부만 볼 수 있도록 제한합니다. 이 부분집합에는 초기 몇 개 단어(attention sinks), 현재 토큰 이전의 모든 구분 기호 토큰, 현재 토큰에 가장 가까운 n개의 토큰이 포함됩니다. 훈련 또는 사전 채우기 단계에서는 입력 컨텍스트의 모든 토큰에 해당하는 모든 쿼리 벡터와 모든 키 벡터를 곱할 필요가 없습니다. 그림 3의 마스크 행렬에서 강조 표시된 요소에 해당하는 쿼리 키 쌍의 벡터만 곱하면 됩니다. 생성 단계에서는 새 토큰을 생성할 때 초기 토큰, 구분 기호 토큰 및 인접 토큰에 대한 KV 캐시만 유지됩니다. 따라서 SepLLM의 KV 캐시는 훨씬 작고 메모리가 덜 필요합니다.

read the caption

Figure 3: The overall paradigm of SepLLM. The left side illustrates the attention mask in the training or pre-filling stage given the input “ABC,DE.FG\n\absent𝑛\backslash n\ italic_n”. The right side illustrates the KV cache management in the generation stage.

🔼 SepLLM의 스트리밍 적용을 위한 프레임워크를 보여주는 그림입니다. KV 쌍은 네 개의 캐시 블록(초기 캐시, 구분자 캐시, 과거 윈도우 캐시, 로컬 윈도우 캐시)에 저장됩니다. 각 행은 반복을 나타내고, 열은 각 캐시 블록의 상태를 보여줍니다. 런타임 사용량(Sizerun)이 최대 용량(c)에 도달하면 SepLLM은 과거 윈도우 캐시에 있는 구분자 토큰의 KV 캐시를 구분자 캐시로 이동하고 다른 KV 캐시는 삭제합니다. 구분자 캐시가 최대 용량(s)에 도달하면 로컬 윈도우 캐시와 과거 윈도우 캐시의 크기가 주기적으로 변화하며, 평균적으로 전체 KV 캐시 사용량은 최대 용량(c)보다 작게 유지됩니다.

read the caption

Figure 4: Overall framework of the proposed SepLLM tailored for streaming applications. The KV pairs are storaged in four cache blocks (displayed as four columns), and are updated in each iteration (shown in a single row). Once the runtime usage S⁢i⁢z⁢er⁢u⁢n𝑆𝑖𝑧subscript𝑒𝑟𝑢𝑛Size_{run}italic_S italic_i italic_z italic_e start_POSTSUBSCRIPT italic_r italic_u italic_n end_POSTSUBSCRIPT reach the max capacity c, SepLLM move KV caches of separator tokens in Past Window Cache into Separator Cache and drop other KV caches.

🔼 이 그래프는 스크래치부터 훈련 과정에서 훈련 단계에 따른 손실 값의 변화를 보여줍니다. Vanilla Transformer, SepLLM(n=64), SepLLM(n=128), SepLLM(n=64, H), SepLLM(n=64, H/T) 등 다양한 모델 아키텍처에 대한 훈련 손실 곡선을 비교하여 SepLLM이 Vanilla Transformer에 비해 훈련 손실이 낮음을 보여줍니다. x축은 훈련 단계를, y축은 손실 값을 나타냅니다.

read the caption

(a) Loss w.r.t steps

🔼 이 그래프는 다양한 모델의 학습 손실 비율을 FLOPs(Floating Point Operations Per Second)에 따라 비교하여 보여줍니다. Vanilla Transformer 모델을 기준으로 SepLLM과 StrmLLM의 손실 비율을 FLOPs 증가에 따라 표시하고 있습니다. SepLLM은 동일한 FLOPs에서 Vanilla Transformer보다 낮은 손실 비율을 보여줍니다. 이는 SepLLM이 연산 효율성 측면에서 Vanilla Transformer보다 우수함을 시사합니다. 또한, SepLLM의 여러 변형(n=64, n=128, n=64, H, n=64, H/T) 또한 Vanilla Transformer보다 낮은 손실 비율을 보입니다.

read the caption

(b) Loss Ratio w.r.t FLOPs

🔼 이 그림은 처음부터 학습할 때의 학습 손실 곡선을 보여줍니다. 5(b)는 FLOPs에 대한 Vanilla 대비 다양한 메서드의 손실 값 비율을 보여줍니다. SepLLM은 Vanilla 모델보다 적은 FLOPs에서 더 낮은 손실을 달성하며, 특히 SepLLM(n=64, H/T)가 두드러집니다. StrmLLM은 Vanilla보다 손실 감소가 느립니다.

read the caption

Figure 5: Training loss curves for training from scratch. 5(b) shows the ratio of the loss values of different methods to that of Vanilla with respect to FLOPs.

🔼 이 그림은 사전 훈련 후 다양한 SepLLM 설정(n=64, n=128, 더 큰 학습률)을 사용한 훈련 손실 곡선을 보여줍니다. n 값을 높이고 학습률을 적절하게 높이면 손실 감소에 도움이 된다는 것을 알 수 있습니다. 또한 전체 어텐션 트랜스포머 체크포인트에서 SepLLM 아키텍처의 요구 사항에 맞는 모델로 사후 훈련을 통해 신속하게 전환할 수 있음을 보여줍니다.

read the caption

Figure 6: Training loss curves for the post-training.

🔼 스트리밍 설정에서 KV 캐시의 진화 과정을 보여줍니다. 그림에서 볼 수 있듯이, 토큰 mo 이후에는 n과 Sizerun이 주기적인 함수 형태를 띠게 되며, 사용되는 평균 KV 캐시 크기는 최대 용량인 c보다 훨씬 작습니다.

read the caption

Figure 7: The evolution of KV caches in the streaming setting.

🔼 이 표는 처음부터 학습된 모델의 성능(다운스트림 작업 성능 및 런타임 KV 캐시 사용량)을 보여줍니다. 다양한 모델 아키텍처(기준, StrmLLM, SepLLM의 변형)에 따른 성능 지표와 런타임 KV 캐시 사용량을 비교하여 SepLLM의 효과를 보여줍니다. Attn(%)는 어텐션 맵의 아랫부분 삼각형에서 ‘1’의 비율을 나타내고, r.KV(%)는 Vanilla 모델에 비해 각 모델이 사용하는 KV 캐시의 비율을 나타냅니다.

read the caption

Table 2: The performance of downstream tasks and the usage of running-time KV cache in the training from scratch setting.

🔼 이 표는 FLOPs(부동 소수점 연산)와 어텐션 맵 비율(Attention Map Ratio)을 비교하여 SepLLM이 기존 모델 대비 얼마나 효율적인지 보여줍니다. 어텐션 맵 비율은 어텐션 마스크의 아래쪽 삼각형에서 ‘1’의 비율을 나타내며, 전체 입력 토큰 중 SepLLM에서 사용하는 토큰의 비율을 의미합니다. FLOPs는 모델의 계산 복잡도를 나타내는 지표입니다. 표에서 SepLLM은 다양한 설정(n=64, n=128, n=64,H, n=64,H/T)에서 FLOPs를 약 30% 줄이는 것을 확인할 수 있으며, 이는 계산 효율성이 크게 향상되었음을 시사합니다.

read the caption

Table 3: The comparison of FLOPs and Attention Map Ratio.

🔼 PG19 테스트 세트에서 StreamingLLM과 SepLLM의 perplexity를 비교한 표입니다. KV 캐시 용량(c)은 324로 고정하고, Sink Cache(a)는 두 모델 모두 4로 설정하여 공정한 비교를 진행했습니다. SepLLM은 StreamingLLM보다 perplexity가 일관되게 낮으며, 이는 SepLLM이 긴 텍스트 생성에서 더 나은 성능을 보인다는 것을 시사합니다.

read the caption

Table 4: The perplexity comparison on the PG19 test set (Rae et al., 2020). For fair evaluation, we keep the KV cache capacity c as 324 and keep Sink Cache a as 4 for both StreamingLLM and SepLLM.

🔼 PG19 테스트 세트에서 다양한 길이의 텍스트(20K 및 64K 토큰)에 대한 Vanilla, StreamingLLM, SepLLM의 perplexity와 실행 시간 비교입니다. r.KV는 생성 과정에서 평균 런타임 KV 캐시 사용량을 나타냅니다. SepLLM은 StreamingLLM과 동일한 최대 KV 캐시 용량 c에서 더 낮은 perplexity, 더 짧은 실행 시간, 더 낮은 평균 런타임 KV 사용량을 달성합니다. 특히 긴 시퀀스에서 그 차이가 두드러집니다.

read the caption

Table 5: The perplexity and runing time comparison on the PG19 test set (Rae et al., 2020). r.KV means the average runtime KV cache usage in the generation process.

🔼 이 표는 SepLLM의 다른 Separator Cache 크기(s)에 따른 WikiText 데이터셋에서의 perplexity와 평균 런타임 KV 캐시 사용량을 보여줍니다. 여기서 a=4, w=224, c=324로 설정되어 있습니다. Separator Cache 크기가 커질수록 perplexity가 감소하는 경향이 있음을 알 수 있습니다.

read the caption

Table 6: The perplexity and average runtime KV cache usage of SepLLM with respect to different Separator Cache sizes (s) on WikiText (Merity et al., 2017), in which a=4, w=224, c=324.

🔼 이 표는 다양한 입력 길이와 다양한 c, w에 따른 WikiText에서의 평균 다운스트림 성능(ppl)과 평균 런타임 KV 사용량을 보여줍니다. 두 방법 모두 a=4이고 SepLLM의 경우 s=64입니다. c와 w는 KV 캐시의 최대 용량과 로컬 윈도우 캐시의 최대 용량을 나타내는 하이퍼파라미터입니다. ppl은 perplexity의 약자로, 언어 모델의 성능을 측정하는 지표입니다. 숫자가 낮을수록 성능이 좋습니다. r.KV는 평균 런타임 KV 캐시 사용량을 나타냅니다.

read the caption

Table 7: Average downstream performance (ppl) over different input lengths and average runtime KV usage with different c,w on WikiText, in which a=4 for both methods and s=64 for SepLLM.

🔼 이 표는 WikiText 데이터셋에서 StreamingLLM과 SepLLM의 성능(퍼플렉시티)과 평균 KV 캐시 사용량을 비교합니다. 두 모델 모두 최대 KV 캐시 용량(c)은 324로 동일하며, 초기 토큰 유지를 위한 캐시 용량(a)은 0 또는 4로 설정했습니다. SepLLM의 경우, 구분자 토큰 캐시 용량(s)은 32, 로컬 윈도우 캐시 용량(w)은 224로 설정했습니다. 이 표를 통해 초기 토큰 유지, 포지션 인코딩 이동 등의 요소가 스트리밍 설정에서 모델 성능에 미치는 영향을 분석할 수 있습니다.

read the caption

Table 8: The perplexity and average runtime KV cache usage of SepLLM and StreamingLLM tested on WikiText (Merity et al., 2017). c=324, a=0/4 for both methods. s=32,w=224 for SepLLM

🔼 이 표는 훈련 가속화에 대한 세부 정보를 제공합니다. Vanilla(전체 어텐션) 모델, SepLLM(n=64) 및 SepLLM(n=128)의 훈련 시간을 비교하여 SepLLM이 훈련 속도를 상당히 향상시키는 것을 보여줍니다.

read the caption

Table 9: The details about training acceleration.

🔼 다양한 디코더 전용 모델, 즉 Llama-3와 Pythia 백본에 SepLLM을 적용한 성능 비교표입니다. PG19 테스트 데이터셋에서 64K 토큰을 생성하는 작업을 기반으로 합니다. SepLLM의 경우, a=4, s=64, w=256, c=800으로 설정했습니다.

read the caption

Table 10: The comparison of SepLLM adapted to different architectures.

🔼 이 표는 다양한 크기의 Pythia 모델에 SepLLM을 적용한 결과를 비교하여 SepLLM의 일반화 성능을 보여줍니다. 구체적으로는 Pythia-6.9B와 Pythia-12B 모델에 대해 PG19 테스트 데이터셋에서 20K 토큰을 생성하는 작업을 수행했습니다.

read the caption

Table 11: The comparison of SepLLM adapted to Pythia (Biderman et al., 2023) with different scales.

🔼 표는 SepLLM을 Falcon-40B 모델에 적용한 결과를 비교하여 보여줍니다. 비교 대상은 StreamingLLM과 Vanilla Transformer 모델이며, 20K 토큰과 64K 토큰 생성 시의 perplexity와 실행 시간을 측정했습니다. SepLLM은 StreamingLLM보다 더 낮은 perplexity를 달성하면서도 실행 시간과 KV 캐시 사용량 측면에서 효율적인 것으로 나타났습니다.

read the caption

Table 12: The comparison of SepLLM adapted to Falcon-40B (Almazrouei et al., 2023).

🔼 이 표는 기본 모델과 지시사항 조정 모델을 사용한 Llama-3-8B 모델에서 SepLLM을 적용한 결과를 비교합니다. GSM8K-CoT 벤치마크에서 추론 능력을 평가하며, SepLLM을 적용해도 기존 어텐션 메커니즘을 사용하는 원본 모델과 비슷한 성능을 보이는 것을 확인할 수 있습니다.

read the caption

Table 13: The comparison of SepLLM adapted to Llama-3-8B (Dubey et al., 2024) with base or instruct versions.