LoopVQ가 루프 트랜스포머 KV 캐시를 0.25비트로 공동 양자화한다

무엇인가

이 논문은 루프 트랜스포머의 KV 캐시 메모리 병목을 다룬다. 루프 트랜스포머는 같은 레이어 묶음을 여러 번 재적용해 파라미터를 아끼지만, 각 루프 단계가 자기 키와 값을 저장하므로 KV 캐시가 루프 수에 비례해 커진다. 논문은 Ouro-1.4B가 Llama-3.1-8B보다 토큰당 KV 캐시를 6배 많이 요구하고, 중간 정도 컨텍스트에서도 KV 캐시가 가중치보다 커진다고 설명한다. 이는 배치 크기와 컨텍스트 길이를 제한하므로 KV 캐시 압축이 필요하다. 기존 스칼라·벡터 양자화는 루프 축의 중복을 활용하지 못하고, 특히 1비트 미만에서 정확도 손실이 크다.

어떻게 동작하나

저자들은 먼저 루프 축의 중복을 관찰한다. 같은 레이어의 키와 값은 동일한 투영 가중치로 계산되어 루프 단계 간 상관이 높다. Ouro-1.4B 전체 레이어·헤드 평균으로 인접 루프 단계 상관은 첫 두 단계에서 0.69, 마지막 두 단계에서 0.94까지 올라간다. 한 헤드에서 같은 채널의 인접 루프 상관은 평균 |r|=0.85인데, 같은 하위 벡터 안의 다른 채널과는 0.16에 그친다. 같은 비트폭에서 한 루프 단계의 G개 채널을 묶는 대신 네 루프 단계의 G/4개 채널씩 묶으면, 0.5비트에서 Ouro-1.4B의 재구성 오차가 대략 절반으로 줄어든다. 또한 루프 단계별 양자화 민감도가 크게 달라 첫 루프 단계 오류가 가장 해롭고 세 번째가 가장 덜 해로우며, KL 발산 기준 약 5배 차이가 난다. 기존 보정 목적함수는 결합된 하위 벡터 안에서 이 차이를 표현하지 못한다.

무엇과 다른가

제안 방법의 첫 구성 요소는 루프 축 벡터 양자화다. T개 루프 단계를 Tc개씩 연속 블록으로 나누고, Tc가 T와 G를 모두 나누도록 한다. 하나의 코드가 블록 안 Tc개 루프 단계 각각에서 c=G/Tc개 채널을 덮는다. 즉 여러 루프 단계의 채널을 하나의 결합 벡터로 만들어 가장 가까운 코드북 엔트리 하나로 인코딩하고, 각 루프 단계는 그 엔트리의 자기 몫만 복원한다. 하나의 인덱스를 Tc개 루프 단계가 공유하므로 비트폭은 기존과 같은 b/G 비트/원소로 유지된다. 코드북은 Fisher 가중 k-means로 초기화하며, 키는 RoPE 적용 전에 양자화하고 어텐션에서 복원된 키에 RoPE를 적용한다.

어떻게 쓰나

두 번째 구성 요소는 루프 인식 코드북 미세조정이다. 모델 가중치는 고정한 채, 보정 시퀀스에서 전체 정밀도 모델의 출력 분포 p_j와 양자화 캐시로 복원한 모델의 출력 q_j 사이 KL 발산을 최소화한다. 이때 bf16 싱크 토큰 이후 위치들을 대상으로 한다. 이 목적함수는 루프 단계 간 오류 전파와 단계별 민감도 차이를 함께 반영한다. 매 반복마다 현재 코드북으로 할당을 다시 계산해 추론 때 쓰는 인코더와 같은 조건에서 코드북을 학습한다. 논문은 Ouro-1.4B에서 미세조정이 전체 정밀도 출력 대비 KL 발산을 대략 절반으로 줄였다고 보고한다.

전제와 한계

시스템 구현은 SGLang에 들어갔다. 캐시는 8토큰 페이지로 구성하고, 처음 싱크 토큰과 최근 토큰은 bf16으로 유지하며 나머지는 코드 인덱스로 저장한다. Tc개 루프 단계가 하나의 인덱스를 공유하므로 양자화 페이지에는 Tc 블록당 한 개의 인덱스 행이 저장되고, 같은 G에서 루프 단계를 따로 코딩할 때와 토큰당 메모리는 같다. 인코딩은 블록의 모든 Tc 단계가 계산된 뒤에야 가능하다. 프리필에서는 각 단계의 키와 값을 마지막 단계까지 보관했다가 행렬곱으로 코드북 거리를 계산해 인코딩하고, 디코딩에서는 최근 윈도에 쓴 뒤 8토큰마다 가장 오래된 페이지를 모든 레이어·루프 단계에 걸쳐 양자화한다. 어텐션 커널은 루프 단계 t에 해당하는 코드북 슬라이스만 읽고, 양자화 토큰을 Flash-Decoding처럼 분할해 logsumexp로 결합한다. 논문은 32K 토큰 컨텍스트에서 디코딩 지연이 거의 절반으로 줄고, BF16 대비 최대 처리량이 Ouro-1.4B에서 4.3배, Ouro-2.6B에서 6.1배 높아졌다고 보고한다.

실험은 Ouro-1.4B, Ouro-2.6B, Huginn-3.5B에서 진행됐다. 모든 방법은 8비트 코드 K=256과 G=16 또는 32를 사용해 각각 0.5비트와 0.25비트/원소를 만든다. 베이스라인은 CQ와 TaSQ다. 보정에는 GPQA main에서 Diamond를 제외한 2,048토큰 시퀀스 16개를 쓰고, LoopVQ는 그중 12개로 Adam 학습률 10^-3, 4에포크 미세조정 후 나머지 4개에서 KL이 가장 낮은 에포크를 고른다. RTX 6000 Ada 한 장에서 k-means 보정과 미세조정은 Ouro-1.4B에 6.1분과 2.3분, Ouro-2.6B에 14.3분과 5.9분 걸린다. 평가는 GSM8K, MATH500, HumanEval, MBPP의 일반 벤치마크, AIME’24·AIME’25·GPQA-Diamond·HumanEval+·MBPP+의 긴 사고 연쇄 추론, RULER의 8개 NIAH 과제 4K~64K 컨텍스트로 이뤄졌다. 논문은 0.5비트에서 LoopVQ가 기존 벡터 양자화 베이스라인을 앞서고 TaSQ+LoopVQ가 BF16에 거의 근접하며, 0.25비트에서는 TaSQ+LoopVQ가 모든 추론 벤치마크에서 최고 결과를 내고 긴 컨텍스트 검색에서도 모든 길이·비트폭에서 가장 높은 정확도를 보인다고 서술한다. 다만 제공된 원문 본문에는 표 2와 표 3의 개별 정확도 셀 값이 실려 있지 않아, 벤치마크별 구체적 정확도 수치는 제시되지 않았다. 절제 실험에서는 루프 축 결합이 이득의 대부분을 차지하고, 미세조정은 HumanEval에서 개선에 기여했다. 보정 코퍼스를 바꿔도 TaSQ+LoopVQ는 GSM8K에서 최대 4포인트, HumanEval에서 1포인트 변해 CQ의 최대 11포인트 변동보다 둔감했다. 또한 CQ와 TaSQ는 GPQA-Diamond에서 32K 토큰 예산을 대부분 소진하고 BF16보다 약 2배 많은 토큰을 생성한 반면, LoopVQ와 TaSQ+LoopVQ는 BF16의 길이 분포를 따랐다.

개발자 관점에서 이 방법은 KV 캐시가 메모리를 지배하는 루프 구조 언어모델을 긴 컨텍스트와 많은 동시 요청으로 서빙할 때 유용하다. 특히 1비트 미만으로 내려가야 하는 상황에서 루프 축을 추가 압축 축으로 쓸 수 있고, 채널 축을 쓰는 TaSQ와 직교적으로 결합할 수 있다. 도입 전에는 대상 모델의 루프 수 T와 채널 그룹 G가 Tc로 나누어지는지, 보정 데이터와 모델별 코드북 준비가 가능한지, bf16 싱크·최근 토큰 정책이 워크로드에 맞는지, 커스텀 어텐션 커널이 배포 스택에서 지원되는지 확인해야 한다. 정확도는 목표 비트폭과 컨텍스트 길이, 배치 크기에서 다시 측정해야 하며, 프리필은 블록 인코딩 때문에 약간 느려질 수 있다.

원문은 별도의 한계 절을 두지 않았지만 전제는 분명하다. 이득을 위해 모델별 보정과 코드북 미세조정이 필요하고, Tc가 T와 G를 나누어야 하며, Huginn의 프렐류드와 코다는 루프 축이 없어 CQ처럼 레이어별로 코딩한다. 0.25비트는 많은 정확도를 유지하지만 BF16과 완전히 같지는 않고, 처리량·지연 이득은 긴 생성과 메모리 병목 배치에서 주로 발생한다. 평가는 특정 Ouro·Huginn 모델과 RTX 6000 Ada 환경을 중심으로 이뤄졌고 A100 결과는 부록에 있다.