RidgeFill, 재학습 없이 긴 컨텍스트 Prefill 가속

들어가며

긴 프롬프트를 처리하는 대규모 언어 모델에서는 답변을 생성하는 시간보다 첫 토큰이 나오기 전의 Prefill 시간이 더 큰 병목이 될 수 있습니다. RidgeFill은 모델을 재학습하지 않고 이 구간의 계산량을 줄이기 위한 근사 추론 기법입니다. 단순히 레이어나 프롬프트 토큰을 삭제하는 대신, 후반부 레이어가 필요로 하는 입력을 예측해 캐시를 채웁니다.

작동 방식

RidgeFill은 모델을 지정된 분할 레이어까지 정상적으로 실행합니다. 이후 분할 레이어의 은닉 상태를 입력으로 삼아, 각 후속 레이어가 원래 받았을 입력을 별도의 맵으로 예측합니다. 이렇게 얻은 값으로 후반부 레이어의 캐시를 생성합니다.

각 맵은 보정용 텍스트를 한 번 순전파한 결과를 이용해 폐쇄형 릿지 회귀로 계산됩니다. 경사 기반 학습도, 원래 모델 가중치의 수정도 없습니다. 후반부 레이어가 수행하는 자체 프로젝션은 그대로 두고, 그 입력만 예측 결과로 대체합니다. 또한 선택된 행, 프롬프트의 마지막 수천 토큰, 모든 디코딩 단계는 전체 모델을 통과시킵니다. 따라서 모든 토큰을 무차별적으로 근사하는 방식이 아니라, 오류에 민감한 구간을 보호하는 구조입니다.

아이디어의 방향은 투영 기능을 학습 과정에서 모델에 넣는 YOCO와 DeepSeek-V4.1-Flash와 유사합니다. 차이는 RidgeFill이 그러한 훈련을 받지 않은 기존 모델에도 이 기능을 사후 적용하려 한다는 점입니다.

측정 결과

논문은 모든 토큰의 평균 손실만으로는 채워진 영역의 오류를 파악하기 어렵다고 설명합니다. 이를 보완하기 위해 해당 영역에서 needle retrieval을 수행하고, 먼 컨텍스트에 의존하는 토큰의 손실을 측정했으며, 예측된 key나 value를 정확한 값으로 교체하는 splice test도 사용했습니다.

오류의 주요 원인은 모델에 따라 달랐습니다. Llama-3.1-8B에서는 key가 특히 중요했고, 그중에서도 key의 스케일이 상당한 영향을 보였습니다. Qwen3.5-9B에서는 value와 반복적으로 전달되는 상태의 영향이 더 크게 나타났습니다.

llama.cpp 포크와 A100에서 실행한 Qwen3.8-27B는 16K120K 토큰 입력에서 TTFT가 기준선보다 1.251.31배 빨라졌습니다. 120K 토큰 프롬프트에서는 약 28초의 절감이 보고됐습니다. RULER는 약 1.21.4포인트, LongBench v2는 약 1.4포인트 낮아졌습니다. 95% 구간은 0을 포함하지만, 몇 포인트 수준의 성능 저하 가능성을 배제하지는 못합니다. vLLM의 Qwen3.8-Flash-Next에서는 9K32K 토큰 범위에서 1.50~1.60배의 속도 향상이 보고됐습니다.

의미와 한계

RidgeFill의 가장 큰 의미는 공개된 모델을 다시 학습하지 않고도 근사 Prefill을 실험할 수 있다는 데 있습니다. 한 번 맵을 만들고 지원되는 추론 구현을 사용하면 긴 입력의 대기 시간을 조절 가능한 방식으로 줄일 수 있습니다.

다만 보정 데이터와 성격이 크게 다른 프롬프트에서도 맵이 안정적으로 작동하는지는 실제 배포 환경에서 확인해야 합니다. 속도 향상 폭 역시 모델, 컨텍스트 길이, 하드웨어, 프레임워크, 정확한 계산으로 남겨두는 토큰 수에 따라 달라집니다. 따라서 RidgeFill은 무손실 대체재라기보다 속도와 품질 사이의 선택지를 제공하는 근사 추론 기법으로 보는 것이 타당합니다.

출처: Hugging Face Daily Papers