2025-11-13T04:07:09.837900

Optimal Quantization for Matrix Multiplication

Ordentlich, Polyanskiy
Recent work in machine learning community proposed multiple methods for performing lossy compression (quantization) of large matrices. This quantization is important for accelerating matrix multiplication (main component of large language models), which is often bottlenecked by the speed of loading these matrices from memory. Unlike classical vector quantization and rate-distortion theory, the goal of these new compression algorithms is to be able to approximate not the matrices themselves, but their matrix product. Specifically, given a pair of real matrices $A,B$ an encoder (compressor) is applied to each of them independently producing descriptions with $R$ bits per entry. These representations subsequently are used by the decoder to estimate matrix product $A^\top B$. In this work, we provide a non-asymptotic lower bound on the mean squared error of this approximation (as a function of rate $R$) for the case of matrices $A,B$ with iid Gaussian entries. Algorithmically, we construct a universal quantizer based on nested lattices with an explicit guarantee of approximation error for any (non-random) pair of matrices $A$, $B$ in terms of only Frobenius norms $\|\bar{A}\|_F, \|\bar{B}\|_F$ and $\|\bar{A}^\top \bar{B}\|_F$, where $\bar{A},\bar{B}$ are versions of $A,B$ with zero-centered columns, respectively. For iid Gaussian matrices our quantizer achieves the lower bound and is, thus, asymptotically optimal. A practical low-complexity version of our quantizer achieves performance quite close to optimal. In addition, we derive rate-distortion function for matrix multiplication of iid Gaussian matrices, which exhibits an interesting phase-transition at $R\approx 0.906$ bit/entry, showing necessity of Johnson-Lindestrauss dimensionality reduction (sketching) in the low-rate regime.
academic

행렬 곱셈을 위한 최적 양자화

기본 정보

  • 논문 ID: 2410.13780
  • 제목: Optimal Quantization for Matrix Multiplication
  • 저자: Or Ordentlich (Hebrew University of Jerusalem), Yury Polyanskiy (MIT)
  • 분류: cs.IT cs.AI cs.CL cs.LG math.IT
  • 발표 시간: 2024년 10월 (arXiv 사전인쇄본)
  • 논문 링크: https://arxiv.org/abs/2410.13780

초록

본 논문은 대규모 행렬 곱셈의 양자화 문제를 심층적으로 연구합니다. 전통적인 벡터 양자화와 달리, 본 연구의 목표는 행렬 자체를 근사하는 것이 아니라 행렬의 곱을 근사하는 것입니다. 두 개의 실수 행렬 A, B가 주어질 때, 인코더는 각 행렬을 독립적으로 압축하며 각 항목마다 R 비트를 사용하고, 디코더는 이러한 압축된 표현을 활용하여 행렬 곱 A⊤B를 추정합니다. 본 논문은 독립동일분포 가우시안 항목을 가진 행렬의 경우에 대해 근사 평균제곱오차의 비점근적 하한을 제공하고, 중첩 격자 기반의 범용 양자화기를 구성하며, R≈0.906 비트/항목에서의 흥미로운 상전이 현상을 발견합니다. 이는 저비율 경우에 Johnson-Lindenstrauss 차원 축소 기법의 필요성을 시사합니다.

연구 배경 및 동기

문제 정의

심층 신경망과 대규모 언어 모델의 등장으로 행렬 곱셈이 계산의 주요 병목이 되었습니다. 현대 컴퓨팅 하드웨어는 종종 계산 능력이 아닌 메모리 대역폭에 의해 제한됩니다. 따라서 메모리 전송을 줄이기 위해 행렬을 손실 압축하는 것이 중요한 문제가 되었습니다.

실제 필요성

대규모 언어 모델의 경우, 저자들은 필요한 양자화율을 다음과 같이 추정했습니다:

  • 생성 단계에서 CPU는 계산 자원을 충분히 활용하기 위해 1-3 비트/항목의 양자화율이 필요합니다
  • 사전 채우기 단계에서 빠른 GPU에서 실행되는 소형 LLM의 경우 약 11.7 비트/항목의 양자화율이 필요합니다

기존 방법의 한계

  1. 고전적 벡터 양자화: 행렬 A와 B를 독립적으로 직접 양자화한 후 양자화된 행렬의 곱을 계산하면 O(n²)의 오차가 발생합니다
  2. 스케치 방법: 불편향 추정을 제공하지만 분산은 여전히 O(n²)입니다
  3. 결정론적 양자화기: 구면 위의 벡터에 대해 Ω(n²)의 하한이 존재합니다

핵심 기여

  1. 이론적 하한: iid 가우시안 항목을 가진 행렬에 대한 행렬 곱셈 양자화의 비점근적 하한 제공
  2. 범용 양자화기: 중첩 격자 기반의 범용 양자화기 구성으로 임의의 행렬에 대한 명확한 오차 보장 제공
  3. 점근적 최적성: 제안된 양자화기가 iid 가우시안 행렬에 대해 하한을 달성하므로 점근적으로 최적임을 증명
  4. 상전이 현상: R≈0.906 비트/항목에서의 상전이를 발견하여 저비율에서 차원 축소의 필요성을 드러냄
  5. 실용적 알고리즘: 최적 성능에 근접한 저복잡도 구현 제공

방법론 상세 설명

작업 정의

행렬 A ∈ R^(n×a)와 B ∈ R^(n×b)가 주어질 때, 목표는 인코더 f₁: R^(n×a) → 2^(naR)과 f₂: R^(n×b) → 2^(nbR) 및 디코더 g를 설계하여 다음을 최소화하는 것입니다:

EABg(f1(A),f2(B))F2E\|A^⊤B - g(f_1(A), f_2(B))\|_F^2

여기서 각 행렬 항목은 R 비트로 설명됩니다.

핵심 함수 Γ(R)

논문은 다음과 같은 핵심 율-왜곡 함수를 정의합니다:

Γ(R)={1(1(222R24R))RRR<R222R24RRR\Gamma(R) = \begin{cases} 1 - \left(1 - (2 \cdot 2^{-2R^*} - 2^{-4R^*})\right) \frac{R}{R^*} & R < R^* \\ 2 \cdot 2^{-2R} - 2^{-4R} & R \geq R^* \end{cases}

여기서 R* ≈ 0.906은 고정점 방정식 R = ½log₂(1 + 4R ln 2)의 해입니다.

중첩 격자 양자화 방식

1. 내적 양자화(기초 구성 요소)

구면 위의 ρ-상관 벡터 U, V에 대해 중첩 격자 Λc ⊂ Λf를 사용하여 양자화합니다:

인코딩 과정:

  • U와 V에 각각 독립적인 디더 벡터 Z₁, Z₂를 추가합니다
  • 세밀한 격자 Λf로 양자화합니다
  • 거친 격자 Λc의 코셋 표현을 출력합니다

디코딩 과정:

  • 코셋에서 양자화 포인트를 복원합니다
  • 디더를 제거합니다
  • 내적 추정을 계산합니다

2. 행렬 곱셈 양자화

전처리 단계:

  1. 영점 중심화: Ā = A - (1/n)1·1^⊤A, B̄ = B - (1/n)1·1^⊤B를 계산합니다
  2. 범수 양자화: 각 열의 평균과 범수를 고정밀도로 설명합니다
  3. 무작위 회전: 동일한 직교 행렬 S를 Ā와 B̄에 적용합니다

양자화 단계:

  • 회전된 각 열에 내적 양자화기를 적용합니다
  • 시간 공유 매개변수 κ와 MMSE 스케일링 매개변수 α를 사용합니다

기술적 혁신점

  1. 디더 기법: 양자화 오차를 입력과 독립적으로 만들어 결정론적 양자화기의 O(n²) 오차를 회피합니다
  2. 중첩 격자 구조: 유한 비율을 제공하면서 우수한 양자화 성능을 유지합니다
  3. 시간 공유: 저비율에서 차원 축소를 통해 최적 성능을 달성합니다
  4. 무작위 회전: 임의의 벡터를 구면 균등 분포로 변환하여 분석을 용이하게 합니다

실험 설정

이론적 검증 실험

데이터 생성:

  • 행렬 A, B ∈ R^(n×n), 항목은 iid N(0,1)
  • n = 3 × 2¹¹

구현 세부사항:

  • 기초 격자: D₃ 격자(3차원)
  • 중첩 비율: q = 6
  • 조회 테이블 크기: < 64KB(L1 캐시에 저장 가능)
  • 유효 비율: ≈ 3.015 비트/기호

비교 방법

  • 3비트 스칼라 양자화기(ℓ∞ 정규화)
  • 이론적 하한 Γ(R)

실험 결과

주요 결과

성능 비교:

  • 제안 방법: 1/n³ ∥Â⊤B - A⊤B∥²F ≈ 0.0593
  • 3비트 스칼라 양자화: ≈ 0.1668(약 3배 차이)
  • 이론적 하한: Γ(3.015) = 0.0304

주요 발견:

  1. D₃ 격자 기반 방식은 스칼라 양자화보다 현저히 우수합니다
  2. 성능은 이론적 최적값에 근접합니다(약 2배 차이)
  3. n의 증가에 따라 성능 차이는 더욱 감소합니다

복잡도 분석

인코딩 복잡도: O(n log n)(빠른 Hadamard 변환 사용) 디코딩 복잡도: O(1)(조회 테이블 사용) 저장 오버헤드: 각 양자화기마다 스케일링 인수를 설명하기 위해 O(log n) 추가 비트 필요

관련 연구

무작위 선형 대수

  • Monte Carlo 행렬 곱셈(MCMM): 행 무작위 샘플링을 통한 근사
  • 국소 민감 해싱(LSH): 코사인 유사도의 저차원 스케치
  • 한계: 상대 오차는 ∥A∥²F∥B∥²F/∥A⊤B∥²F에 따라 증가합니다

신경망 양자화

  • 훈련 후 양자화: OPTQ, GPTQ 등의 방법
  • 회전 기법: QuIP, QuaRot은 Hadamard 변환을 사용합니다
  • 격 양자화: QUIP#은 가중치 양자화에 E₈ 격자를 사용합니다

정보 이론

  • 분산 압축: 선형 함수 계산을 위한 압축
  • 코드북 설계: Voronoi 코드와 중첩 격자 코드

결론 및 토론

주요 결론

  1. 최적성: iid 가우시안 행렬의 경우 제안된 방식이 정보 이론적 하한을 달성합니다
  2. 범용성: 임의의 행렬에 대해 명확한 성능 보장을 제공합니다
  3. 상전이 현상: R* ≈ 0.906 비트/항목이 핵심 임계값입니다
  4. 실용성: 저복잡도 구현이 이론적 성능에 근접합니다

한계

  1. 공유 무작위성: 인코더와 디코더가 무작위 시드를 공유해야 합니다
  2. 행렬 조건: 행렬 항목이 유계여야 합니다(M = n^(10^22000))
  3. 고차원 격자: 이론적 최적성을 위해 고차원 "좋은" 격자가 필요하지만, 실제로는 저차원 격자의 곱을 사용합니다
  4. 결정론적 방식: 무작위성이 필요 없는 최적 결정론적 방식의 존재 여부가 불명확합니다

향후 방향

  1. 다중 행렬 곱: k>2개 행렬의 곱으로 확장
  2. 다른 거리 척도: KL 발산 등 MSE가 아닌 척도
  3. 결정론적 양자화기: 공유 무작위성이 필요 없는 방식 탐색
  4. 심층 신경망 응용: 실제 LLM에서의 배포 및 최적화

심층 평가

장점

  1. 이론적 엄밀성: 상한과 하한을 포함한 완전한 정보 이론 분석 제공
  2. 실용적 가치: LLM 추론의 실제 병목 문제 해결
  3. 기술적 혁신: 격 양자화, 무작위 회전, 시간 공유를 교묘하게 결합
  4. 범용성: 특정 행렬 분포 가정에 의존하지 않음

부족한 점

  1. 복잡성: 이론 분석이 상당히 복잡하며 실제 구현에 여러 기술 요소 필요
  2. 상수 인수: 점근적으로 최적이지만 유한 표본에서의 상수가 클 수 있음
  3. 하드웨어 적응: 다양한 하드웨어 플랫폼에 맞춘 구현 최적화 필요
  4. 확장성: 2개 행렬에서 다중 행렬 곱으로의 확장이 비자명함

영향력

이론적 기여:

  • 행렬 곱셈 양자화의 정보 이론적 기초 확립
  • 상전이 현상과 차원 축소의 필요성 규명
  • 해당 분야에 중요한 이론적 기준 제공

실제 응용:

  • LLM 양자화에 새로운 이론적 지침 제공
  • 후속 연구인 NestQuant가 실제 LLM에서 SOTA 성능 달성
  • 하드웨어 가속기 설계에 이론적 근거 제공

적용 가능 분야

  1. 대규모 언어 모델 추론: 가중치와 활성화의 결합 양자화
  2. 엣지 컴퓨팅: 메모리 제약 환경의 행렬 연산
  3. 분산 컴퓨팅: 통신 제약 행렬 곱셈
  4. 과학 계산: 대규모 수치 선형 대수 문제

참고문헌

논문은 정보 이론, 격 이론, 무작위 선형 대수, 신경망 양자화 등 여러 분야의 중요한 연구를 포함한 44개의 관련 문헌을 인용합니다. 특히 주목할 만한 것은:

  • Zamir의 격 부호화 전문서가 이론적 기초 제공
  • Erez와 Zamir의 중첩 격자에 관한 개척적 연구
  • OPTQ, QuIP 등 최근 LLM 양자화 방법
  • 무작위 행렬 이론과 구면 기하학의 관련 결과

전체 평가: 이는 이론과 실제 모두에서 중요한 기여를 하는 우수한 논문으로, 행렬 곱셈 양자화 문제에 견고한 정보 이론적 기초를 제공하며 거의 최적에 가까운 실용적 알고리즘을 제시합니다. 본 연구는 대규모 기계학습 시스템의 양자화 기술을 이해하고 개선하는 데 중요한 의미를 가집니다.