🤖 machine learning

BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization

본 논문은 비가분산적인 비터비 argmax 를 온도 제어 BCJR 합-곱 알고리즘으로 대체하여 엔드투엔드 양자화 인식 훈련을 가능하게 하고, 대규모 언어 모델에서 기존 최첨단 PTQ 최전선을 경험적으로 능가하는 가분산 트레리스 부호화 가중치 양자화의 완화 기법인 BCJR-QAT 를 소개합니다.

원저자: Venugopalan Iyengar

게시일 2026-05-12

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Venugopalan Iyengar

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

거대한 책 도서관 (대형 언어 모델) 이 표준 노트북이나 스마트폰에 들어갈 수 있도록 축소하고 싶다고 상상해 보세요. 이를 위해 도서관 내부의 '단어'(가중치) 를 압축해야 합니다.

이 논문은 이러한 책들을 단어당 단 2 비트까지 더 압축하는 새로운 방법인 BCJR-QAT를 소개합니다 (이는 고화질 사진을 작고 거친 썸네일로 압축하는 것과 같습니다).

다음은 그들이 이를 수행한 방식을 간단한 비유로 설명한 이야기입니다:

1. 문제: '일방통행' 함정

이전까지 이러한 모델을 축소하는 가장 좋은 방법은 QTIP이라는 방법이었습니다. QTIP을 매우 똑똑한 사서라고 생각해보세요. 이 사서는 텍스트 페이지를 보고 "좋아, 이 복잡한 문장을 내 사전에서 가장 가까운 간단한 구로 바꾸겠어"라고 말합니다.

그러나 한계가 있었습니다. 사서가 한 번 그 선택을 하면, 책 전체가 더 잘 들리도록 그 선택을 되돌려 변경할 수 없었습니다. 이는 마치 '일방통행'과 같았습니다. 사서가 초기에 약간 잘못된 구를 선택하면 나머지 책이 피해를 보게 되었고, 결정이 '단단하고' 최종적이었기 때문에 이를 수정할 수 없었습니다.

이를 해결하기 위해 연구자들은 보통 **학습 (QAT)**을 사용했습니다. 즉, 사서가 선택을 연습하고 조정하도록 하는 것입니다. 하지만 여기에는 함정이 있었습니다. 사서의 결정 과정은 복잡한 미로 (트레리스) 를 포함합니다. 미로에서 최상의 경로를 찾기 위해 그들은 **비터비 (Viterbi)**라는 규칙을 사용하는데, 이는 '절대적으로 최상의 옵션을 선택'하는 버튼과 같습니다.

문제점: '최고의 것을 선택'하는 버튼을 통해 컴퓨터가 학습하도록 가르칠 수 없습니다. 그 버튼은 수학적 막다른 길이기 때문입니다. 버튼을 누르면 컴퓨터는 더 나아지기 위해 선택을 어떻게 미세 조정해야 하는지 알지 못합니다. 그저 가장 가까운 옵션으로 딱 맞춰질 뿐입니다. 마치 차를 운전할 때 다음 차선으로 즉시 점프하는 것만 허용하여 부드러운 회전은 불가능하게 만드는 것과 같습니다.

2. 해결책: '부드러운' 결정 (BCJR)

저자 Venugopalan Iyengar 는 사서가 연습할 수 있는 새로운 방법을 고안했습니다. 즉시 '최고의 것을 선택'하는 단단한 결정을 강요하는 대신, **온도 ( $T$ )**라는 개념을 도입했습니다.

높은 온도 (뜨거움): 사서가 매우 편안하고 결정을 못 내리는 상태라고 상상해 보세요. 하나의 구만 선택하는 대신, 여러 구를 동시에 고려하여 각각에 '확률'을 부여합니다. 안개 낀 날처럼 하나의 경로가 아니라 여러 경로를 볼 수 있는 것과 같습니다. 이 '부드러운' 관점은 수학적으로 매끄러워, 컴퓨터가 책 전체를 개선하기 위해 선택을 어떻게 조정해야 하는지 쉽게 계산할 수 있습니다.
낮은 온도 (차가움): 학습이 마무리됨에 따라 사서는 '차갑게' 되어 더 결단력 있게 변합니다. 안개가 걷히고 그들은 단일 최상의 구 (단단한 결정) 로 딱 맞춰집니다.

BCJR이라고 불리는 이 방법은 '일방통행'을 주행 가능한 매끄러운 도로로 바꿉니다. 이제 컴퓨터는 개별 문장뿐만 아니라 책 전체가 더 잘 들리도록 가중치를 어떻게 조정해야 하는지 정확히 학습할 수 있습니다.

3. '오버슈트' 실수

저자들은 이 '온도'가 작동하는 방식에서 놀라운 특이점을 발견했습니다.

전통적인 물리학에서는 최상의 해결책을 찾으려 할 때 보통 매우 뜨겁게 (매우 결정을 못 내리는 상태) 시작하여 모든 가능성을 탐색한 다음 서서히 식히는 방식입니다.

논문의 발견: 이 특정 유형의 압축의 경우, '너무 뜨겁게' 시작하는 것은 재앙입니다. 사서가 초기에 너무 결정을 못 내리면, 나쁜 이웃 (더 나쁜 보로노이 분지) 으로 헤매게 되어 그곳에 갇히게 됩니다. 그들이 식어 최상의 경로를 찾으려 할 때는 이미 나쁜 지역에 빠져 있어 돌아올 수 없습니다.
해결책: 그들은 적당한 온도 (너무 뜨겁지도, 너무 차갑지도 않은) 에서 시작하는 것이 가장 효과적임을 발견했습니다. 안개 속에서 맹목적으로 헤매는 것이 아니라 명확한 지도를 가지고 하이킹을 시작하는 것과 같습니다. '매우 뜨겁게' 시작하는 단계를 건너뛰어 길을 잃는 것을 피하고 더 나은 해결책을 찾았습니다.

4. 결과: 소형 컴퓨터를 위한 승리

팀은 두 가지 다른 유형의 모델에서 이를 테스트했습니다.

'프록시' 테스트 (OLMoE): 모델의 각 계층이 자신의 데이터를 얼마나 잘 재구성하는지 (예: 복사본이 원본과 비슷한지 확인) 만 보고 압축을 최적화하려고 시도했습니다. 결과: 실패했습니다. 모델은 오히려 이전 방법보다 성능이 떨어졌습니다. 이는 단순히 '더 나은 복사본'을 만드는 것이 책이 더 잘 읽힌다는 것을 의미하지 않는다는 것을 가르쳐 주었습니다.
'실제' 테스트 (Llama-3.2): 모델이 실제로 텍스트를 읽고 이해하는 능력을 향상시키도록 최적화했습니다 (지식 증류라고 불리는, 똑똑한 교사 모델이 학생을 안내하는 방법 사용). 결과: 성공했습니다!
- 모델의 특정 계층에서 그들의 새로운 방법은 작지만 중요한 차이로 이전 최첨단 방법을 능가했습니다.
- 이를 여러 계층에 적용했을 때, 개선 사항이 '초가산적'인 방식으로 누적되었습니다 (전체가 부분의 합보다 커짐).

5. 엔진의 내부

이 수학을 수행하는 것은 보통 매우 느리고 값비싼 슈퍼컴퓨터가 필요합니다. 저자들은 또한 단일 소비자용 그래픽 카드 (RTX 4080 과 같은) 에서 이 복잡한 수학을 incredibly 빠르게 실행하는 전용 '엔진' (Triton 커널) 을 구축했습니다. 이를 표준 방법보다 6.5 배 빠르게 만들어, 이 고급 기술이 데이터 센터 없이도 실행될 수 있음을 증명했습니다.

요약

이 논문은 소비자용 기기에 맞도록 AI 모델을 축소하는 새로운 방법을 제시합니다.

트릭: 학습할 수 없는 경직된 결정 과정을 점진적으로 단단해지는 '부드러운' 학습 가능한 것으로 대체했습니다.
통찰: 학습 과정을 너무 '뜨겁게' (너무 무작위하게) 시작해서는 안 됩니다. 적절한 시작은 모델이 길을 잃는 것을 방지합니다.
결과: 데이터를 단순히 복사하는 것이 아니라 텍스트를 이해하도록 올바르게 학습되었을 때, 이 방법은 이전 최첨단 기술보다 더 나은 결과를 산출하며 일반 게이밍 PC 에서 충분히 빠르게 실행됩니다.

저자들은 다른 사람들이 시도해 볼 수 있도록 코드와 학습된 모델을 공개했습니다.

기술 요약: BCJR-QAT

문제 정의

소비자 하드웨어에서 대규모 언어 모델 (LLM) 의 가중치당 2 비트 (bpw) 이하의 양자화를 달성하는 것은 중요한 병목 현상입니다. 트렐리스 부호화 양자화 (TCQ) 를 사용하여 2 비트 양자화를 위한 강력한 최전선을 확립한 QTIP 와 같은 훈련 후 양자화 (PTQ) 방법들은 "PTQ 한계"에 의해 제한됩니다. 이를 넘어서기 위해서는 양자화 인식 훈련 (QAT) 이 필요합니다. 그러나 표준 QAT 기법은 트렐리스 코드에서 실패합니다. 비터비 (Viterbi) 디코딩 알고리즘이 지수적으로 많은 경로에 대한 비미분 가능한 argmax 연산을 포함하기 때문입니다. 트렐리스에 직통 추정기 (STE) 를 적용하면 인접한 가중치 결정 간의 결합이 파괴되는데, 이는 스칼라 양자화에 비해 트렐리스 부호화의 주요 장점입니다.

방법론: BCJR-QAT

본 논문은 하드 비터비 argmax 를 BCJR (Bahl-Cocke-Jelinek-Raviv) 순방향 - 역방향 알고리즘에서 유도된 소프트 기대값으로 대체하는, 미분 가능한 트렐리스 부호화 양자화의 완화 기법인 BCJR-QAT를 소개합니다.

핵심 메커니즘

통계 역학 비유: 트렐리스는 1 차원 이징 (Ising) 유사 스핀 사슬과 동등한 인자 그래프로 간주됩니다. 비터비 argmax 는 경로에 대한 볼츠만 분포의 영온도 ( $T \to 0$ ) 한계로 간주됩니다.
소프트 코드워드: 유한 온도 $T > 0$ 에서, 이 방법은 BCJR 알고리즘을 사용하여 각 상태의 정확한 주변 사후 확률을 계산합니다. 양자화된 출력은 방출된 코드워드의 볼츠만 기대값입니다:
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
이 소프트 코드워드는 가중치 전상 (pre-image) 및 트렐리스 매개변수에 대해 정확히 미분 가능합니다.
어닐링 (Annealing): 훈련 동안 온도 $T$ 는 초기 값 $T_0$ 에서 거의 0 인 $T_{end}$ 로 어닐링됩니다. $T \to 0$ 로 갈 때, 소프트 코드워드는 하드 비터비 출력으로 수렴하여 연속 최적화를 효과적으로 이산 해로 "결정화"합니다.

주요 기술적 혁신

미분 가능성: BCJR 순방향 - 역방향 재귀 (복잡도 $O(L \cdot S^2)$ ) 는 트렐리스를 통해 정확한 기울기를 제공하여, 스칼라 STE 방법이 잃어버리는 상태 간 결합을 보존합니다.
Triton 커널: 저자는 gather, log-sum-exp, 덧셈 연산을 단일 런칭으로 결합하는 퓨즈드 Triton 커널을 구현했습니다. 이는 순방향/역방향 패스를 수백 개의 커널 런칭에서 하나로 줄여, FP32 수치적 동등성을 유지하면서 참조 PyTorch autograd 구현 대비 6.57 배의 속도 향상을 달성했습니다.
스케줄 수정 (Skip-High-T): 본 논문은 표준 시뮬레이션 어닐링 스케줄 ( $T_0=1.0$ 에서 시작) 의 치명적인 실패 모드를 진단합니다. 고온에서 소프트 코드워드 기울기는 "흐려져" 정보성이 없어져, 냉각 단계가 회복하기 전에 최적화가 더 나쁜 보로노이 (Voronoi) 분지로 drifting 하게 만듭니다. 제안된 해결책은 고온 단계를 건너뛰는 것으로, 기울기가 이미 정보성을 갖는 더 낮은 온도 (예: $T_0=0.3$ ) 에서 초기화하는 것입니다.

실험 결과

본 논문은 OLMoE-1B-7B(MoE) 및 Llama-3.2-1B(Dense) 두 가지 아키텍처에서 BCJR-QAT 를 평가했습니다.

1. OLMoE-1B-7B (부정적 결과: 프록시 간극)

설정: 재구성 MSE(프록시 목적 함수) 를 최적화하는 레이어별 탐욕적 QAT.
결과: BCJR-QAT 는 QTIP-PTQ 대비 레이어별 재구성 MSE 를 17~20% 성공적으로 감소시켰습니다. 그러나 이 개선은 최종 작업 성능으로 이어지지 않았습니다. WikiText-2 에서 모델은 QTIP-PTQ 기준 대비 약 1.3 perplexity (PPL) 만큼 성능이 저하되었습니다.
결론: 이는 2 bpw 에서 국소 재구성 오차를 최소화하는 것이 더 나은 최종 작업 성능을 보장하지 않는 "프록시 간극"을 확인시켜 줍니다.

2. Llama-3.2-1B (긍정적 결과: 엔드투엔드 KL)

설정: FP16 교사 모델에 대한 엔드투엔드 순방향-KL 증류를 최적화하는 단일 레이어 QAT.
결과:
- 스케줄 오버슈트: 순진한 고온 스케줄 ( $T_0=1.0$ ) 을 사용하면 PTQ 보다 약간 나쁜 성능 (+0.005~+0.022 PPL) 이 나타났습니다.
- Skip-High-T 수정: 수정된 스케줄 ( $T_0=0.3$ ) 을 사용하면 오버슈트가 제거되었습니다. 레이어 4에서 모델은 10.1347 PPL을 달성하여 QTIP-PTQ 기준 (10.2189) 대비 −0.084 PPL 개선을 보였습니다.
- 초가산성 (Super-Additivity): 훈련된 레이어 (레이어 4 및 레이어 8) 를 결합했을 때, 결합된 이득 (−0.077 PPL) 은 개별 이득의 합 (−0.062 PPL) 을 초과하여 BCJR 훈련 코드워드들이 서로를 강화하는 "협력 잉여"를 입증했습니다.

3. 이론적 분석

드리프트 - 예산 이론: 저자들은 PTQ 보로노이 분지를 탈출하기 위한 조건을 유도했습니다: 누적 가중치 드리프트 ( $\eta \cdot N \cdot g_{max}$ ) 는 분지 반지름 ( $r_{Voronoi} \approx 10^{-3}$ ) 을 초과해야 합니다. 실험은 이 조건을 충족하지 못하는 스케줄은 PTQ 해를 탈출할 수 없음을, 반면 이를 초과하는 스케줄은 목적 함수와 스케줄이 정확하다면 더 나은 해를 찾을 수 있음을 확인시켜 줍니다.

중요성 및 주장

본 논문은 세 가지 주요 기여를 주장합니다:

알고리즘적: BCJR-QAT 는 경로 결합의 이점을 잃지 않고 TCQ 에서 QAT 를 가능하게 하는, 트렐리스 부호화 양자화를 네이티브로 미분 가능하게 만든 최초의 방법입니다.
공학적: 퓨즈드 Triton 커널은 이 방법을 소비자 하드웨어 (단일 12GB GPU) 에서 계산적으로 실현 가능하게 만들어, 연구의 추론 시간 배포 장벽을 제거했습니다.
경험적 통찰:
- 레이어별 재구성 MSE 는 2 비트 QAT 를 위한 나쁜 프록시이며, 엔드투엔드 KL 증류가 효과적임을 확인했습니다.
- 이산 완화를 위한 시뮬레이션 어닐링에 대한 통념에 도전하여, 고온 ( $T=1.0$ ) 에서 시작하는 것이 BCJR-QAT 에 해롭다는 것을 보였습니다. 고온 기울기는 잡음에 지배적이기 때문입니다. 최적 전략은 기울기가 이미 정보성을 갖는 온도에서 시작하는 것입니다.

본 논문은 단일 레이어 결과가 유망하지만, BCJR-QAT 의 전체 잠재력은 모든 레이어에 걸친 공동 엔드투엔드 훈련에 있을 것이라고 결론지었습니다. 이 작업은 공개된 커널로 인해 비용 효율적이 되었지만, 10 억 개 이상의 파라미터를 가진 모델의 경우 H100/H200 과 같은 클라우드 규모 하드웨어가 필요합니다. 저자들은 이 작업에서 이러한 전체 규모의 공동 훈련을 수행하지 않았다고 명시적으로 밝혔습니다.