Layer-Condensed KV Cache 논문 리뷰: 층의 기억을 공유할 때 얻는 처리량과 학습 비용

언어모델을 작게 저장해도 긴 대화를 여러 개 동시에 처리하면 KV cache가 메모리를 크게 차지합니다. Layer-Condensed KV Cache는 과거 토큰을 적게 남기는 대신, 그 기억을 보관하는 층의 수를 줄입니다. 일부 층의 KV를 다른 층도 함께 참조하도록 모델 구조를 바꾸어 배치를 키우려는 접근입니다. 다만 최대 26배라는 처리량 이득은 특정 하드웨어와 offload 조건의 결과이며 학습 비용과 프롬프트 처리의 제약을 함께 보아야 합니다.

토큰을 줄이는 압축과 층을 줄이는 압축

이 리뷰는 ACL 2024에 실린 논문을 대상으로 합니다. 일반 decoder는 생성 중 각 층의 key와 value를 저장해 다음 토큰 계산에 재사용합니다. 이 공간은 문맥 길이·동시 처리량·층 수가 커질수록 늘어납니다. 기존의 토큰 선택 방식이 어떤 과거 토큰을 남길지 묻는다면, 이 연구는 그 토큰의 KV를 모든 층에서 따로 만들 필요가 있는지 묻습니다.

모든 층의 query를 최상위 층의 KV와 연결하면 대부분의 층에서 KV 계산과 저장을 없앨 수 있습니다. 해당 층의 key·value 투영 가중치도 줄일 수 있습니다. 다만 이미 학습된 모델에서 cache 몇 개를 지우는 설정 변경만으로 같은 결과를 얻었다는 연구는 아닙니다. 새로운 연결 구조와 그에 맞는 학습 절차를 함께 제안합니다.

원문 Figure 1, PDF 1쪽 / 인쇄 11175쪽. 표준 Transformer와 상위층 KV 공유 구조의 연결 비교. 노드는 토큰의 층별 계산을, 가로 연결은 질의와 KV의 짝을 나타낸다. 출처: Wu, H., & Tu, K. (2024). Layer-condensed KV cache for efficient inference of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) (pp. 11175–11188). Association for Computational Linguistics. https://aclanthology.org/2024.acl-long.602/ · CC BY 4.0, https://creativecommons.org/licenses/by/4.0/ · 원문 영역 캡처(여백만 제외), 내용 변경 없음.

Figure 1의 연결이 만드는 순환 문제

Figure 1에서 노드는 한 토큰의 한 층 계산이고, 가로 연결은 query와 참조할 KV의 관계입니다. 아래층에서 현재 토큰의 최상위 KV를 요구하면, 그 최상위 층은 아래층 계산이 끝나야 만들어지는 순환이 생깁니다. 저자들은 현재 토큰이 자기 자신을 attention으로 참조하는 대각 연결을 가려 이 문제를 풉니다. 첫 토큰에는 0으로 된 임시 KV를 사용합니다.

이때 현재 토큰 정보가 전부 사라지는 것은 아닙니다. 잔차 연결을 통한 아래에서 위로의 계산에는 여전히 들어갑니다. 그러나 모든 층이 최상위 KV만 쓰는 구성은 언어모델링 성능을 떨어뜨렸습니다. 이를 보완하기 위해 일부 층은 원래 attention을 유지하는 warmup layer로 남깁니다. w는 그 층의 수이며 학습률을 올리는 warmup step과는 다른 개념입니다.

아래와 위에 남기는 warmup layer

기본 배치는 아래쪽 w/2층과 위쪽 w/2층을 남기는 sandwich 형태입니다. 낮은 층과 높은 층이 서로 다른 정보를 처리한다는 직관을 반영합니다. w를 늘리면 원래 Transformer에 가까워져 품질을 회복할 여지가 있지만, KV를 별도로 저장하는 층도 늘어 메모리 절감은 작아집니다. 정확도와 처리량을 조절하는 구조적 선택인 셈입니다.

학습에는 또 다른 문제가 있습니다. 각 토큰이 앞 토큰의 최상위 KV에 의존하면 일반 Transformer처럼 모든 토큰을 한 번에 병렬 계산하기 어렵습니다. 논문은 이전 반복에서 얻은 KV로 다음 반복을 계산하는 방식으로 바꾸고, KV가 비교적 빨리 수렴하는 현상을 이용해 반복 횟수를 줄입니다. 이는 추론 cache를 줄이는 대신 학습 쪽에 복잡성을 더하는 교환입니다.

m=7과 b=2가 뜻하는 것

기본 학습은 앞쪽 m=7회 반복으로 KV를 근사하고, 마지막 b=2회에 대해서만 역전파합니다. 모든 반복의 계산 그래프를 저장하는 메모리 부담을 피하기 위한 gradient stopping입니다. 여기서 b는 역전파 반복 수이며 처리량 공식에서 사용하는 batch 변수와 역할이 다릅니다. 논문은 적은 반복으로 충분한지 별도 실험하지만 모든 모델·입력에서 같은 수렴 속도를 증명한 것은 아닙니다.

프롬프트를 넣을 때도 반복 계산이 필요합니다. 따라서 생성할 토큰이 많으면 그 준비 비용을 나눌 수 있지만, 매우 긴 문서를 넣고 짧게 요약하는 작업에서는 부담이 커질 수 있습니다. 이 구조의 메모리 이득을 이해하려면 입력 길이와 생성 길이를 하나의 ‘문맥 길이’로 합쳐 보지 말고 둘을 따로 남겨야 합니다.

원문 Table 1, PDF 5쪽 / 인쇄 11179쪽. GPU·모델 크기·입력+생성 길이별 최대 배치와 처리량(tokens/s). CPU-offload 조건과 w=2·10 구분을 보존했다. 출처: Wu, H., & Tu, K. (2024). Layer-condensed KV cache for efficient inference of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) (pp. 11175–11188). Association for Computational Linguistics. https://aclanthology.org/2024.acl-long.602/ · CC BY 4.0, https://creativecommons.org/licenses/by/4.0/ · 원문 영역 캡처(여백만 제외), 내용 변경 없음.

최대 26배가 나온 행과 비교 범위

Table 1은 RTX 3090 24GB와 A100 80GB에서 1.1B·7B·30B 구성을 비교합니다. x+y 표기는 입력 x토큰과 생성 y토큰입니다. 처리량은 프롬프트 처리와 전체 생성을 합친 시간 t에 대해, 배치 수와 생성 토큰 수의 곱을 t로 나눈 tokens/s입니다. 사용자가 한 요청의 첫 토큰을 얼마나 빨리 받는지와 같은 지표는 아닙니다.

최대 26배는 RTX 3090에서 CPU-offload를 사용하는 30B, 입력 512·생성 1,024 조건입니다. 표준 Llama는 최대 배치 4, 처리량 0.23 tokens/s이고 w=2 모델은 배치 83, 5.99 tokens/s입니다. 같은 행에서 w=10은 배치 23, 1.63 tokens/s로 약 7.1배입니다. 최대 배수가 일반적인 단일 요청 가속을 뜻하지 않으며, 원래 조건의 낮은 처리량과 배치 확대가 함께 반영된 결과입니다.

offload가 없는 A100의 30B, 입력·생성 각 2,048 조건에서는 Llama가 배치 1·14.10 tokens/s, w=2가 배치 32·108.29 tokens/s로 약 7.7배입니다. 논문은 일부 속도 측정 길이가 학습 시 최대 입력 길이를 넘는다고도 밝힙니다. 해당 길이의 메모리·처리량을 측정했다는 사실을 그 길이에서 답변 품질까지 검증했다는 주장으로 바꾸면 안 됩니다.

원문 Table 2, PDF 6쪽 / 인쇄 11180쪽. 1.1B 모델의 일곱 상식 추론 zero-shot 정확도와 평균. warmup 층 수 w=2·10 비교이며 높은 점수가 좋다. 출처: Wu, H., & Tu, K. (2024). Layer-condensed KV cache for efficient inference of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) (pp. 11175–11188). Association for Computational Linguistics. https://aclanthology.org/2024.acl-long.602/ · CC BY 4.0, https://creativecommons.org/licenses/by/4.0/ · 원문 영역 캡처(여백만 제외), 내용 변경 없음.

품질 비교는 1.1B의 별도 학습 실험

품질 평가는 두 1.1B 모델을 SlimPajama 1,000억 토큰 부분집합에서 처음부터 학습하고, 같은 크기의 TinyLlama와 비교합니다. w=2와 w=10 모두 m=7·b=2이며 학습에는 A800 80GB 128장을 사용합니다. 속도 표에 7B·30B가 있다고 그 크기의 모든 품질 결과도 같은 표에서 입증된 것으로 보아서는 안 됩니다.

Table 2의 일곱 상식 추론 zero-shot 평균은 TinyLlama 46.65, w=2는 45.45, w=10은 46.84입니다. w=10의 평균이 0.19%포인트 높아도 모든 과제에서 우세하지는 않습니다. PIQA는 TinyLlama 68.93, w=10 67.90입니다. 평균 점수가 가깝다는 결론과 서비스의 모든 입력에서 동일하게 답한다는 결론은 다릅니다.

언어모델링에서는 SlimPajama 개발 자료 1,000만 토큰 부분집합의 perplexity도 평가합니다. TinyLlama는 9.219, w=2는 9.746, w=10은 9.265입니다. 낮을수록 좋은 지표이므로 w=10이 기준에 가깝고 w=2는 더 손해를 보는 경향입니다. 이 결과는 KV 절감의 최대값보다 목표 품질에 맞춰 w를 선택해야 한다는 판단을 뒷받침합니다.

원문 Figure 4, PDF 5쪽. A100의 7B 모델, 입력·생성 길이 각각 2,048 조건. 가로축 배치 크기는 2배 간격, 세로축 처리량은 tokens/s다. 배치 32 부근 이후의 포화와 큰 배치에서의 감소도 보존했다. 출처: Wu, H., & Tu, K. (2024). Layer-condensed KV cache for efficient inference of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) (pp. 11175–11188). Association for Computational Linguistics. https://aclanthology.org/2024.acl-long.602/ · CC BY 4.0, https://creativecommons.org/licenses/by/4.0/ · 원문 영역 캡처(여백만 제외), 내용 변경 없음.

배치를 더 키우면 언제까지 빨라지는가

Figure 4는 A100의 7B, 입력과 생성 각 2,048토큰에서 배치 수를 바꾼 처리량입니다. 가로축은 batch size, 세로축은 tokens/s이며 로그 간격의 배치 눈금을 같은 선형 간격으로 읽지 않아야 합니다. 제안 모델은 같은 배치에서도 표준 Llama보다 높지만, 배치 32 이후에는 처리량이 더 늘지 않고 128에서는 오히려 줄어듭니다.

이는 메모리 부족을 완화한 뒤 계산 자체가 병목이 될 수 있음을 시사합니다. 저자들은 더 큰 배치보다 효율적인 커널이 필요할 수 있다고 해석합니다. 따라서 최대 배치가 32배 커졌다는 숫자를 처리량 32배와 같다고 쓸 수 없습니다. 두 측정값은 서로 다른 열이며 목표 서비스에서 함께 확인해야 합니다.

학습 비용과 적용 대상의 경계

저자들은 같은 데이터 양의 TinyLlama보다 사전학습 시간이 약 3배 든다고 명시합니다. KV를 줄이는 추론 이득을 얻는 대신 반복 학습 비용을 치르는 것입니다. 또한 입력이 생성보다 훨씬 긴 문서 요약에서는 처리량이 나빠질 수 있다고 지적합니다. 기존 모델에 즉시 붙여 쓰는 무비용 cache 옵션으로 소개하기 어려운 이유입니다.

실무 검토에서는 입력·출력 길이의 분포와 동시 요청 수를 먼저 확인하고, 허용할 품질 차이에 맞춰 w를 비교할 수 있습니다. 이어 같은 장비·offload 여부·엔진에서 전체 처리량과 개별 요청 지연을 따로 측정해야 합니다. 이는 이 리뷰의 적용 제안입니다. 이 논문의 의미는 모든 토큰의 모든 층에 별도 기억을 둘 필요가 있는지 질문하고, 절감한 메모리가 실제 처리량으로 바뀌는 조건과 비용을 함께 제시한 데 있습니다.

논문 출처: https://aclanthology.org/2024.acl-long.602/ · 도표 라이선스: CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). 원문 도표의 내용은 변경하지 않고 주변 여백을 제외해 인용했습니다.