River-LLM: KV 공유 기반의 대형 언어 모델 심리스 조기 종료 기법

대형 언어 모델의 추론 속도를 높이기 위해 불필요한 연산 층을 건너뛰는 Early Exit 기법은 이전 토큰의 데이터가 유실되는 KV Cache Absence 문제로 인해 실질적인 속도 향상이 어려웠다. 이 논문은 별도의 추가 학습 없이도 유실된 데이터를 자연스럽게 생성하고 공유하는 구조를 제안하여, 모델의 정확도를 유지하면서도 실제 추론 속도를 획기적으로 개선했다.

Transformer 기반의 LLM은 이전 토큰들의 연산 결과인 Key와 Value를 저장하는 KV Cache를 사용하여 다음 토큰을 생성한다. 하지만 특정 토큰이 중간 레이어에서 연산을 멈추고 나가는 Early Exit을 수행하면, 그 이후 레이어들에 대한 KV Cache가 생성되지 않아 다음 토큰이 연산을 이어갈 수 없는 KV Cache Absence 문제가 발생한다. 기존에는 이를 해결하기 위해 다시 계산하거나 마스킹을 사용했지만, 이는 연산 비용을 높이거나 정확도를 크게 떨어뜨리는 한계가 있었다.

River-LLM은 'Exit River'라는 별도의 가속 경로를 구축하여 이 문제를 해결한다. 이 경로는 원래 모델의 레이어들과 데이터를 공유하도록 설계되어, 토큰이 조기에 종료되더라도 해당 경로를 통과하며 필요한 KV Cache를 자동으로 채워넣는다. 이는 마치 고속도로에서 나들목으로 빠져나가더라도 나란히 달리는 보조 도로를 통해 필요한 정보를 계속 전달받는 것과 유사한 원리이다.

결과적으로 모델은 모든 레이어를 거치지 않고도 데이터의 연속성을 유지할 수 있게 된다. 특히 초기 레이어에서의 데이터 변화량이 적을수록 조기 종료가 안전하다는 점에 착안하여, 입력과 출력의 유사도를 측정해 오차를 관리함으로써 높은 생성 품질을 유지하면서도 연산량을 대폭 줄였다.

River-LLM은 기존 LLM 백본에 경량화된 플러그인 형태의 Exit Layer들을 직렬로 연결한 구조를 가진다. 각 Exit Layer는 대응하는 백본 디코더 블록의 가중치를 복사한 후 W4A16 양자화를 적용하여 연산 효율을 극대화한다. [백본 레이어의 가중치 입력] → [4비트 양자화 수행] → [경량화된 가중치 출력] → [백본 대비 2.4배 빠른 연산 수행]

종합적인 종료 결정은 상태 전이 유사도(State Transition Similarity)를 통해 이루어진다. 현재 레이어 l의 입력 h(l-1)과 출력 h(l) 사이의 코사인 유사도를 계산하여 임계값 τ보다 클 경우 조기 종료를 결정한다. [h(l-1)과 h(l) 벡터 입력] → [내적 및 크기 정규화 연산] → [유사도 점수 출력] → [점수가 높으면 이후 레이어들이 유사한 연산을 반복할 것으로 판단하여 종료]

추론 단계에서는 Prefill 단계와 Generation 단계를 구분하여 처리한다. Prefill 단계에서는 병렬 연산 효율을 위해 시퀀스 단위 종료를 수행하고, Generation 단계에서는 토큰별로 최적의 깊이에서 종료하도록 유연하게 전환하여 최대의 가속 효과를 얻는다.

Llama 3.2 1B 모델을 사용한 GSM8K 벤치마크 실험에서 River-LLM은 2.16배의 실질적인 속도 향상을 기록했다. 이는 이론적인 레이어 감소 수치와 실제 속도 향상 사이의 간극을 좁힌 결과로, 기존의 KV Recompute 방식이 속도 향상이 거의 없었던 것과 대조적이다.

정확도 측면에서도 우수한 성능을 보였다. Llama 3.1 8B 모델의 경우 HumanEval 데이터셋에서 백본 모델(57.3)과 유사한 55.5의 점수를 유지하면서도 1.77배의 속도 향상을 달성했다. 특히 일부 작업에서는 불필요한 심층 연산을 건너뜀으로써 발생하는 노이즈 억제 효과 덕분에 원본 모델보다 소폭 높은 정확도를 기록하기도 했다.

메모리 효율성 분석 결과, River-LLM은 공유 KV 메커니즘을 통해 추가적인 KV Cache 메모리 점유 없이 동작한다. 64K 이상의 긴 문맥 처리 시에도 기존 Early Exit 프레임워크들이 메모리 사용량이 폭증하는 것과 달리, 원본 백본 모델보다 낮은 메모리 점유율을 유지하며 안정적으로 동작함이 확인됐다.

River-LLM은 디코더 전용 아키텍처에서 Early Exit의 고질적 문제인 KV Cache 무결성을 보장하기 위해 'KV Share' 메커니즘을 도입했다. 이는 백본 레이어와 Exit 레이어가 동일한 KV Cache 주소 지정 체계를 공유하도록 설계되어, Exit 경로에서 생성된 KV 값이 즉시 백본의 캐시에 반영되도록 한다.

수학적으로는 레이어 l에서의 상태 전이 유사도 s(l)을 정의하고, 이를 누적 양자화 오차의 대리 지표(Proxy)로 활용한다. 실험적으로 첫 번째 레이어의 유사도와 최종 레이어의 출력 일치도 사이에 유의미한 상관관계(r=0.5536)가 있음을 입증하여, 연산 초기에 종료 여부를 결정할 수 있는 이론적 근거를 마련했다.

또한 'Backbone Offloading' 전략을 통해 사용되지 않는 백본 레이어들을 VRAM에서 제거함으로써, 전체 모델 파라미터 메모리를 14.96GB에서 6.73GB로 줄이면서도 4비트 정적 양자화 모델보다 높은 정확도를 유지하는 유연한 배포 구조를 지원한다.

현재 평가는 8B 파라미터 이하의 모델들에 집중되어 있어, 70B 이상의 초대형 모델에서의 확장성 및 성능 유지는 추가 검증이 필요하다. 또한 토큰별 조기 종료가 빈번한 Generation 단계에 최적화되어 있어, Prefill 단계가 지배적인 작업에서는 가속 효과가 상대적으로 적게 나타날 수 있다.