LoopVQ: 순환 Transformer의 KV 캐시를 1비트 미만으로 압축

순환 Transformer는 동일한 레이어 그룹을 반복 호출하여 훨씬 적은 파라미터로 몇 배 더 큰 일반 모델의 성능을 따라잡을 수 있지만, 각 순환 단계마다 자체 KV를 저장해야 하므로 캐시가 순환 단계 수에 따라 선형으로 증가하여 주요 메모리 병목 현상이 됩니다. 저자들은 KV가 순환 축(axis)에서 높은 중복성을 보이며, 순환 단계마다 양자화 민감도가 크게 다르다는 점을 관찰했습니다. 이에 따라 LoopVQ를 제안했습니다. LoopVQ는 연속된 여러 순환 단계의 채널을 하나의 서브 벡터로 결합하고 하나의 코드북 인덱스를 공유하는 순환 축 벡터 양자화를 적용한 후, Fisher 가중치 k-means 코드북을 미세 조정하여 양자화된 모델의 출력이 전체 정밀도 모델에 가깝도록 만듭니다(학습 시 각 순환 단계에서 양자화된 캐시를 읽음). 이 방법은 일반 LLM을 위한 방법론과 직교하며, TaSQ의 키 변환과 결합하여 사용할 수 있습니다.

원제: LoopVQ: Sub 1-bit KV Cache Quantization for Looped Transformers

순환 Transformer는 동일한 레이어 그룹을 반복 호출하여 훨씬 적은 파라미터로 몇 배 더 큰 일반 모델의 성능을 따라잡을 수 있지만, 각 순환 단계마다 자체 KV를 저장해야 하므로 캐시가 순환 단계 수에 따라 선형으로 증가하여 주요 메모리 병목 현상이 됩니다. 저자들은 KV가 순환 축(axis)에서 높은 중복성을 보이며, 순환 단계마다 양자화 민감도가 크게 다르다는 점을 관찰했습니다. 이에 따라 LoopVQ를 제안했습니다. LoopVQ는 연속된 여러 순환 단계의 채널을 하나의 서브 벡터로 결합하고 하나의 코드북 인덱스를 공유하는 순환 축 벡터 양자화를 적용한 후, Fisher 가중치 k-means 코드북을 미세 조정하여 양자화된 모델의 출력이 전체 정밀도 모델에 가깝도록 만듭니다(학습 시 각 순환 단계에서 양자화된 캐시를 읽음). 이 방법은 일반 LLM을 위한 방법론과 직교하며, TaSQ의 키 변환과 결합하여 사용할 수 있습니다.

좋아요 8 · 코드 스타 수 3

이유

  • 저자들은 요소당 0.5 비트일 때, Ouro 및 Huginn 모델의 일반 작업, 긴 사고 체인 추론 및 긴 컨텍스트 검색 작업에서 기존 일반 LLM용 KV 벡터 양자화 방법보다 우수하다고 주장합니다. TaSQ를 결합하면 0.25 비트에서도 대부분의 BF16 정밀도를 유지합니다.
  • 시스템 이점이 실질적으로 입증되었습니다. 저자들이 SGLang에 구현한 결과, 32K 컨텍스트에서 디코딩 지연이 거의 절반으로 줄었으며, Ouro 1.4B / 2.6B의 최대 서비스 처리량은 BF16 기준선의 각각 4.3배와 6.1배로 향상되었습니다.
  • 코드는 Apache-2.0 라이선스입니다: 캘리브레이션 프로세스, 패킹된 인덱스 저장소 및 융합된 디코딩 커널을 포함한 SGLang 포크, 그리고 논문의 모든 표에 대한 재현 스크립트가 리포지토리에 포함되어 있습니다.

사용법

  • ./setup_envs.sh를 실행하여 세 개의 환경(loopvq-serve 엔진, loopvq-eval 평가, loopvq-build 캘리브레이션)을 생성하고, export MODELS_DIR을 Ouro-1.4B, Ouro-2.6B, Ouro-*-Thinking 및 huginn-0125를 가리키도록 설정합니다. 일반 표 재현: MODEL=ouro_1p4b GPUS="0 1 2 3" bash experiments/loopvq_general/build.sh를 실행한 후, MODEL=ouro_1p4b bash experiments/loopvq_general/run_all.sh를 실행하고, 마지막으로 python3 experiments/loopvq_general/make_table.py를 실행합니다.
  • 서비스 시작: loopvq-serve 환경에서 PREFIX_TOKENS=8 RECENT_TOKENS=64 PREFILL_BACKEND=triton ./scripts/serve_method.sh $MODELS_DIR/Ouro-1.4B vq <코드북 디렉토리> --port 30000 --disable-piecewise-cuda-graph를 실행하면 SGLang의 OpenAI 호환 인터페이스를 얻을 수 있습니다. 처음과 최근 토큰은 bf16로 유지됩니다.

위험 요소

  • 순환 Transformer에만 적용 가능합니다(논문에서는 Ouro 1.4B / 2.6B 및 Huginn-3.5B 사용). 일반 LLM에는 직접 적용할 수 없으며, 캘리브레이션에는 GPQA 시퀀스 16개만 사용되었습니다.
  • 지연 및 처리량 수치는 저자의 환경에서 측정된 것입니다. 좋아요 8, 코드 스타 수 3(관심도만 반영).

열기 · 코드 · 논문 페이지 열기

Hugging Face Daily Papers

최신 AI 소식 보기