RIFT는 추측 디코딩에서 초안 모델이 쓰는 KV 캐시를 압축하는 방법이다. 연구진은 압축성 학습(compressibility training)을 결합한 RIFT에서 Qwen3.5-4B가 Qwen3.5-122B의 초안을 작성할 때, 입력 32K 토큰과 출력 4K 토큰, KV 캐시 예산 8K 토큰 조건으로 전체 캐시 수락 길이의 98% 넘게 유지했다고 밝혔다.
추측 디코딩은 디코드 가속에 널리 쓰이는 방법이며, 점점 더 큰 KV 캐시가 필요한 긴 출력 생성과 특히 관련이 깊다. 긴 문맥에서 초안을 만들 때는 KV 캐시 압축으로 비용을 제한할 수 있다. 그런데 프리필 정보를 중심으로 고르는 압축은 출력이 길어질수록 쓸모가 떨어진다. 연구진은 이를 디코드 노후화(decode staleness)라고 부르며, 선택 기반 압축에서 초안 모델의 정확도를 낮춘다고 설명했다.
RIFT의 구성
RIFT는 프리필 정보와 최근 생성한 토큰의 균형을 맞추는 선택 기반 압축 전략의 묶음이다. 토큰을 생성하는 도중 캐시 선택을 새로 고치면 추가 디코드 지연 시간을 대가로 초안 품질이 크게 개선된다. 이 비용을 줄이려고 연구진은 점수 계산 횟수를 줄이면서도 품질을 유지하는 적응형 KV 갱신 전략을 제안했다.
여기에 희소 어텐션 마스크를 이용한 압축성 중심의 초안 모델 학습 단계를 더했다. 압축된 캐시에서도 초안 품질이 나아진다는 것이 연구진의 설명이다. 98% 넘는 수치는 이 압축성 학습을 결합했을 때의 결과다.
출처
이 기사는 아래 자료를 바탕으로 정리했습니다. 제목을 누르면 원문이 새 창으로 열립니다.