KronQ: 그래디언트 공분산과 Kronecker-팩터화 헤시안 근사를 통합한 사후 학습 양자화

사후 학습 양자화는 대형 언어 모델을 재학습 없이 압축하는 현실적인 방법으로 널리 채택되어 있다. 기존의 두 번째 차수 기반 PTQ 방법들은 활성화 통계만을 기준으로 계층별 재구성 손실을 구성해 출력 채널별 민감도 차이를 무시하는 경향이 있었다. KronQ는 그래디언트 공분산을 양자화 파이프라인에 통합해 출력 차원 쪽 민감도를 반영함으로써 특히 극저비트 설정에서 기존 방법들이 발산하거나 열등한 결과를 보이는 상황에서도 성능을 보존했다.

대형 언어 모델의 계층별 양자화는 입력 활성화 통계만으로 계층 손실을 구성하는 관행이 일반적이다. 이 경우 출력 채널마다 손실에 기여하는 정도가 동일하다고 가정하게 되어 실제로는 출력 쪽 민감도가 큰 채널에 과도한 양자화 오류가 발생할 수 있다. 특히 2비트와 같은 극저비트에서는 이러한 가정이 모델 성능의 급격한 열화를 초래한다. KronQ의 핵심은 손실 구성을 활성화 통계뿐 아니라 그래디언트 공분산까지 확장해 계층 손실이 입력과 출력의 공분산 정보를 동시에 반영하도록 만드는 것이다. Kronecker-팩터화 헤시안 근사는 헤시안의 거대한 행렬을 입력과 출력 쌍의 크로네커 곱 형태로 근사해 계산 비용을 억제한다. 이 근사를 통해 활성화와 그래디언트의 통계가 결합된 손실 항을 실용적인 비용으로 얻을 수 있으며, 출력 차원에서의 민감도 차이를 정량적으로 반영할 수 있다. 이 접근은 두 가지 실질적 효과를 낳는다. 첫째, 입력과 출력 양쪽에서 가중치의 크기 분산이 줄어들어 균일한 양자화 스텝으로도 오류가 감소한다. 둘째, 그래디언트 기반 민감도 지표를 통해 혼합 정밀도 할당이 데이터 기반으로 이루어져 리소스를 더 효율적으로 사용하게 된다. 이러한 변화는 극저비트 설정에서 기존 방법들이 발산하거나 매우 높은 perplexity를 보일 때도 모델 성능을 유지하도록 만든다.

전체 접근은 Kronecker-팩터화 헤시안 근사 하에서 양자화 손실을 재정의하는 것이다. 기존 두 번째 차수 PTQ는 입력 활성화 공분산만을 사용해 계층별 재구성 항을 구성했으나 KronQ는 그래디언트 공분산을 추가해 손실이 활성화와 그래디언트 통계의 곱으로 표현되도록 했다. 이 손실식은 각 가중치 블록이 출력과 입력 쌍에 대해 가지는 민감도를 반영한다. 핵심 메커니즘으로는 양방향 비동조 처리가 있다. 입력 측의 랜덤 회전과 대응되게 출력 차원도 그래디언트 공분산을 이용해 회전 및 정규화를 수행함으로써 가중치 분포의 편차를 양쪽 차원에서 동시에 줄였다. 이 과정은 양자화 대상 행렬의 스펙트럼과 크기 분포를 안정화해 낮은 비트폭에서도 양자화 오류를 억제한다. 혼합 정밀도 할당을 위해서는 그래디언트 및 활성화 헤시안의 트레이스를 결합한 민감도 지표를 도출했다. 이 지표는 각 계층의 양자화로 인한 손실 증가량을 근사해 주어진 비트 예산 하에서 비트폭을 효율적으로 배분하도록 설계되었다. 계층별 비트 할당은 이 지표를 기준으로 우선순위를 매겨 진행되며 실험적으로 전반적인 성능을 개선하는 것으로 확인됐다.

주요 결과로서 2비트 가중치 전용 양자화 설정에서 LLaMA-3-70B를 대상에 두었을 때 GPTQ와 GPTAQ는 WikiText-2에서 발산하거나 2000을 초과하는 perplexity를 보였으나 KronQ는 7.93 perplexity를 달성했다. 이 수치는 동일 조건에서 KronQ가 극저비트 환경에서도 의미 있는 언어 모델 성능을 유지함을 직접적으로 보여준다. 논문은 이 결과를 통해 그래디언트 정보의 통합이 안정성과 성능 유지에 결정적임을 입증했다. 추가적인 비교실험과 절제 연구에서 양방향 비동조 처리와 민감도 기반 혼합 정밀도 할당이 각각 성능 개선에 유의미한 기여를 하는 것으로 확인됐다. 절제 실험에서는 출력 측 회전 처리를 제거하면 저비트 환경에서 성능 저하가 커지는 경향이 관찰되었다. 이는 출력 쪽 통계가 양자화 민감도에 실질적으로 영향을 준다는 근거가 된다. 효율성 관점에서 KronQ의 계산 오버헤드는 Kronecker 근사를 활용해 실용적인 범위로 억제되었다. 논문은 대규모 모델에 적용 가능한 수준의 추가 연산과 메모리 요구로 설계되었음을 보고했다. 실제 배포 가능성은 실험 설정과 하드웨어 제약에 따라 달라질 수 있으나 극저비트에서의 성능 보존이 주요 성공 지표로 제시되었다.

전체 아키텍처는 각 계층의 가중치 블록을 대상으로 활성화 공분산과 그래디언트 공분산을 결합한 손실을 최소화하도록 설계되었다. Kronecker-팩터화 헤시안 근사는 계층별 헤시안을 입력 공분산과 출력 공분산의 크로네커 곱으로 분해해 큰 행렬을 효율적으로 처리할 수 있게 만든다. 이 근사로 인해 두 통계의 결합 항을 계산하는 비용이 실용적으로 유지된다. 출력 차원에 대한 랜덤 회전과 정규화는 그래디언트 공분산의 고유구조를 이용해 수행된다. 회전은 가중치 행렬의 스펙트럼을 평탄화하고 분포의 치우침을 줄여 균일 양자화 스텝으로도 오류를 줄일 수 있게 설계되었다. 이 과정은 입력 측 회전과 대칭적으로 동작해 양방향에서 분산을 감소시킨다. 혼합 정밀도 민감도는 그래디언트 헤시안 트레이스와 활성화 헤시안 트레이스를 결합한 형태로 도출되었다. 수식적으로는 계층별 재구성 손실의 2차 근사를 통해 파라미터 블록의 기여도를 근사하고 이 값의 크기를 기준으로 비트 예산을 분배한다. 계층별 할당은 제한된 총 비트 예산을 만족하도록 정수 비트폭으로 반올림하는 단계가 포함되며, 논문은 이 절차의 실험적 유효성을 보고했다. 구현 및 학습 세부사항으로는 Kronecker 근사 추정을 위한 미니배치 통계 계산과 회전 변환의 수치적 안정성 확보가 포함된다. 실험은 다양한 비트폭 조합과 데이터셋에서 수행되었으며 극저비트 설정에서는 특히 그래디언트 정보 통합의 효과가 컸다. 코드 저장소에는 핵심 변환과 민감도 계산이 포함되어 있어 재현성이 확보되어 있다.