💬 NLP

Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

이 논문은 다양한 작업과 KV 캐시 예산 요구사항에 맞춰 토큰 선택 계층을 적응적으로 결정하여 추론 속도와 정확도 간의 균형을 최적화하는 새로운 방법인 ASL 을 제안합니다.

원저자: Rei Taniguchi, Yuyang Dong, Makoto Onizuka, Chuan Xiao

게시일 2026-04-17

📖 3 분 읽기☕ 가벼운 읽기

원저자: Rei Taniguchi, Yuyang Dong, Makoto Onizuka, Chuan Xiao

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

🧠 거대한 도서관의 비밀: "ASL"이 어떻게 기억을 정리하는가?

이 논문은 최근 화제가 되는 **거대 언어 모델 (LLM)**이 아주 긴 글을 읽거나 대화할 때 겪는 '기억 과부하' 문제를 해결하는 새로운 방법, **ASL(Adaptive Selection Layer)**을 소개합니다.

이해하기 쉽게 거대한 도서관과 비서의 비유로 설명해 드릴게요.

1. 문제: "모든 책을 다 기억할 수는 없어요!" 📚🤯

LLM 이 긴 문서를 읽을 때, 중요한 정보 (키 - 밸류, KV) 를 KV 캐시라는 메모리에 저장해 둡니다. 마치 비서가 대화 내용을 종이에 적어두고 나중에 참고하는 것과 같습니다.

기존 방식의 문제:
- 기존 기술들은 **"무조건 15 번째 페이지에서 중요한 내용만 골라라"**라고 정해져 있었습니다.
- 쉬운 질문 (예: "책 제목이 뭐야?") 이라면 15 페이지에서 바로 답을 찾을 수 있어 좋습니다.
- 하지만 어려운 질문 (예: "책 100 페이지와 200 페이지의 숨겨진 연관성을 찾아줘") 이라면, 15 페이지에서는 답을 찾을 수 없습니다. 중요한 정보가 훨씬 뒤에 숨어있기 때문입니다.
- 결과: 정해진 규칙 (고정된 층) 을 따르다 보니, 어려운 질문을 할 때 비서가 중요한 정보를 버리고 엉뚱한 답을 내놓는 경우가 많았습니다.

2. 해결책: "상황에 따라 언제 골라야 할지 스스로 판단하는 비서" 🧐✨

이 논문에서 제안한 ASL은 정해진 규칙이 아니라, 상황을 보고 스스로 판단합니다.

🕵️‍♂️ 핵심 아이디어: "주목도 (Attention) 가 안정될 때까지 기다려라"

비서가 책을 읽을 때, 처음에는 모든 페이지에 골고루 눈이 가지만 (혼란), 점점 중요한 부분으로 시선이 집중됩니다.

ASL 의 방법:
1. 비서가 책을 읽으며 "어떤 페이지가 가장 중요한지" 순위를 매깁니다.
2. 몇 페이지를 더 읽을 때마다, 이 순위가 얼마나 변하는지 확인합니다.
3. 순위가 더 이상 크게 변하지 않고 안정화되면, "아, 이제 중요한 부분들이 확실해졌구나!"라고 판단합니다.
4. 그 순간을 선택의 순간으로 정하고, 그 시점부터는 중요한 정보만 기억하고 나머지는 버립니다.

🎯 비유: "등산로에서 휴식 지점 정하기"

기존 방식: "등산로 500m 지점에서 무조건 휴식해." (어떤 산이든 똑같음)
- 가파른 산 (어려운 과제) 이라면 500m 지점에서는 숨이 턱까지 차서 더 이상 오를 수 없습니다.
ASL 방식: "숨이 차고 다리가 떨릴 때까지 계속 올라가다가, 숨이 차는 정도가 일정해지면 그곳을 휴식 지점으로 정해."
- 평탄한 길 (쉬운 과제) 이라면 일찍 휴식하고, 가파른 길 (어려운 과제) 이라면 더 올라간 후 휴식합니다.

3. ASL 의 장점: "똑똑하고, 빠르고, 정확해요" 🚀

어떤 과제도 척척:
- 쉬운 질문에는 일찍 정보를 정리해서 빠르게 답을 내고, 어려운 질문에는 정보를 더 오래 보관했다가 정리해서 정확하게 답을 냅니다.
학습 불필요:
- 모델을 다시 가르칠 필요 없이, 기존 모델을 그대로 쓰면서 이 방법만 적용하면 됩니다. (훈련 없는 방법, Training-free)
기존 기술과 함께 작동:
- ASL 은 '언제' 정보를 골라낼지 정하고, 기존 기술 (SnapKV 등) 은 '어떻게' 정보를 압축할지 도와줍니다. 두 기술이 손잡고 일하는 셈입니다.

4. 실험 결과: "어려운 미션에서 압도적 승리" 🏆

연구진은 아주 긴 문서를 다루는 다양한 테스트 (InfiniteBench, RULER 등) 를 진행했습니다.

쉬운 질문: 기존 방법과 비슷하게 빠르고 정확했습니다.
어려운 질문 (예: 책 속의 특정 숫자 찾기, 복잡한 코드 디버깅): 기존 방법들은 실수하거나 답을 못 찾았지만, ASL 은 거의 완벽하게 정답을 찾아냈습니다.
속도: 정확도가 높아진 대신, 아주 미세하게 시간이 더 걸릴 수 있지만, 그 대가로 얻는 정확도 향상은 매우 큽니다.

📝 한 줄 요약

"ASL 은 LLM 이 긴 글을 읽을 때, '언제' 중요한 정보를 골라내야 할지 스스로 판단하게 해주는 똑똑한 비서입니다. 어려운 문제일수록 더 오래 기다렸다가 정확한 답을 찾아내어, 기존 방식보다 훨씬 똑똑하게 작동합니다."

이 기술 덕분에 앞으로 AI 는 더 긴 문서도 더 정확하게, 그리고 더 효율적으로 처리할 수 있게 될 것입니다! 🌟

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 의 긴 문맥 처리 능력은 중요해졌지만, 이를 위한 키-밸류 (KV) 캐시의 메모리 부담이 심각한 병목 현상이 되고 있습니다. 이를 해결하기 위해 최근 제안된 레이어별 토큰 가지치기 (Layer-wise Token Pruning) 기법들은 특정 레이어에서 중요한 토큰만 선택하여 KV 캐시를 축소하는 방식을 사용합니다.

그러나 기존 방법들 (FastKV, GemFilter 등) 은 사전에 정의된 고정된 레이어 (Fixed Selection Layer) 에서 토큰을 선택한다는 치명적인 한계가 있습니다.

유연성 부재: 작업의 난이도나 특성에 따라 최적의 선택 레이어는 달라지지만, 고정된 레이어를 사용하면 성능이 크게 저하됩니다.
- 쉬운 작업 (예: QA): 초기 레이어에서 중요한 정보를 찾을 수 있어 중간 레이어에서 선택해도 잘 작동합니다.
- 어려운 작업 (예: KV Retrieval, NIAH): 문맥과 질문의 의미적 유사성이 높아 초기 레이어에서는 핵심 토큰을 식별하기 어렵습니다. 이 경우 고정된 중간 레이어에서 선택하면 정확도가 급격히 떨어집니다.
결과: 어려운 작업에서 성능을 유지하려면 선택 레이어를 더 깊은 층으로 미루거나 KV 예산을 늘려야 하므로, 메모리 감소의 이점이 사라집니다.

2. 제안 방법: ASL (Adaptive Selection Layer)

저자들은 작업의 난이도에 따라 적응형으로 선택 레이어를 결정하는 훈련 불필요 (Training-free) 방법인 ASL을 제안했습니다.

핵심 아이디어

관찰: Attention 점수에 따라 정렬된 토큰의 순위 (Rank) 분산 (Variance) 을 레이어를 거치며 모니터링합니다.
- 초기 레이어: Attention 점수가 모든 토큰에 고르게 분포되어 있어 토큰 순위의 변동이 큽니다.
- 깊은 레이어: Attention 점수가 소수의 중요한 토큰에 집중되면서 토큰 순위가 안정화됩니다.
작동 원리:
1. 관측 시작: 모델의 특정 레이어 ( $L_{min}$ ) 부터 시작하여 최근 $L_{obs}$ 개의 레이어에서 Attention 점수의 평균화 (Pooling) 및 토큰 순위를 계산합니다.
2. 상대 분산 계산: 현재 레이어에서의 토큰 순위 분산을 초기 레이어의 분산으로 나누어 상대 분산 (Relative Variance) 을 구합니다.
3. 적응형 선택: 사용자가 설정한 임계값 ( $\tau$ ) 보다 상대 분산이 낮아지면, Attention 이 안정적인 토큰 집합에 집중되었다고 판단하여 현재 레이어를 선택 레이어로 결정합니다.
4. 단일 샷 (One-shot) 선택: 결정된 레이어에서 토큰을 한 번만 선택하고, 이후 모든 깊은 레이어에서는 선택된 토큰만 처리합니다.

시스템 통합

Prefilling 단계: ASL 은 Prefilling 중 작동하며, 오버헤드는 매우 적습니다.
KV 예산 준수: 선택된 레이어 이전의 KV 캐시 크기를 사용자의 예산 (예: 2048) 으로 맞추기 위해 SnapKV나 GemFilter와 결합하여 사용할 수 있습니다.
- ASL + SnapKV: Prefilling 시 SnapKV 로 KV 를 압축하고, ASL 로 선택 레이어를 결정하여 디코딩 효율을 극대화합니다.
- ASL + GemFilter: 2-pass 전략으로 정확도를 더 높일 수 있습니다.

3. 주요 기여 (Key Contributions)

고정된 선택 레이어의 비효율성 규명: 기존 레이어별 가지치기 방법들이 작업 난이도에 따라 최적의 레이어가 다름을 분석하고, 고정된 레이어가 어려운 작업 (KV Retrieval 등) 에서 성능 저하를 초래함을 증명했습니다.
ASL 방법론 제안: Attention 점수에 기반한 토큰 순위의 분산 변화를 모니터링하여 작업 인식 (Task-aware) 방식으로 선택 레이어를 동적으로 결정하는 새로운 알고리즘을 제안했습니다.
유연한 통합성: 기존 KV 캐시 축소 기법 (SnapKV, GemFilter 등) 과 원활하게 통합되어 Prefilling 및 Decoding 단계를 모두 최적화할 수 있음을 보였습니다.
성능 입증: InfiniteBench, RULER, NIAH 등 3 가지 벤치마크에서 SOTA(SOTA) 방법들보다 우수한 정확도 - 효율성 트레이드오프를 달성함을 실험적으로 증명했습니다.

4. 실험 결과 (Results)

정확도 (Accuracy):
- InfiniteBench: KV 예산 2048 조건에서 ASL(2-pass) 은 Llama-3.1-8B-UL 에서 평균 점수가 가장 높았으며, 특히 KV Retrieval과 같은 어려운 작업에서 고정 레이어를 사용하는 FastKV 보다 월등히 우수한 성능을 보였습니다.
- RULER: 긴 문맥 (128k) 에서 ASL 은 FastKV 보다 일관되게 높은 정확도를 기록했습니다.
- NIAH (Needle in a Haystack): Qwen2.5-7B 모델에서 ASL 은 모든 문맥 길이에서 풀 KV(Full KV) 와 동등한 성능 (100% 점수) 을 달성했습니다. 반면, GemFilter 는 긴 문맥에서 성능이 저하되었습니다.
효율성 (Efficiency):
- TTFT (Time to First Token): ASL 은 토큰 선택 시기가 FastKV 보다 늦을 수 있어 TTFT 가 약간 느릴 수 있지만, 여전히 풀 KV 보다 빠릅니다.
- TPOT (Time Per Output Token): 디코딩 속도는 다른 가지치기 방법들과 유사하게 매우 빠릅니다.
- 메모리 사용량: ASL 의 추가적인 메모리 오버헤드는 무시할 수준이며, KV 캐시 축소 효과는 기존 방법들과 동일하게 유지됩니다.
트레이드오프: ASL 은 TTFT 를 약간 희생하더라도 어려운 작업에서 정확도를 크게 향상시킵니다.

5. 의의 및 결론 (Significance)

이 논문은 LLM 추론 시 KV 캐시 최적화 분야에서 고정된 규칙에서 적응형 전략으로의 전환을 주도합니다.

작업별 최적화: "어떤 레이어에서 토큰을 잘라야 하는가?"에 대한 정답이 작업마다 다르다는 점을 인식하고, 이를 실시간으로 판단하는 메커니즘을 도입했습니다.
실용성: 추가적인 미세 조정 (Fine-tuning) 이 필요 없으며, 기존 모델과 다양한 KV 압축 기법과 호환되어 즉시 적용 가능합니다.
미래 방향: 특히 긴 문맥 처리가 필수적인 복잡한 추론 작업 (KV Retrieval, 다중 키 검색 등) 에서 기존 방법들의 한계를 극복하고, 높은 정확도와 효율성을 동시에 달성할 수 있는 새로운 기준을 제시했습니다.

요약하자면, ASL은 LLM 이 문맥을 이해하는 과정에서 Attention 패턴이 안정화되는 시점을 감지하여, 작업의 난이도에 맞춰 유연하게 토큰을 선택함으로써 메모리 효율과 정확도 간의 최적 균형을 찾는 혁신적인 방법론입니다.