Contraction-Gauge Preconditioning for Quantized Matrix Multiplication

기술 요약: 양자화된 행렬 곱셈을 위한 수축-게이지 프리컨디셔닝 (Contraction-Gauge Preconditioning)

문제 정의

본 논문은 두 인자 $A \in \mathbb{R}^{m \times K}$와 $B \in \mathbb{R}^{K \times n}$가 모두 양자화된 저정밀도 행렬 곱셈($C = AB$)의 과제를 다룹니다. 이 영역의 주요 어려움은 양자화 손실이 균일하게 전파되지 않는다는 점입니다. 즉, 특정 엔트리$A_{ik}$의 오차는 대응하는$B$의 행($B_{k,:}$)의 에너지에 비례하여 출력 행$C_{i,:}$에 영향을 미치며, 그 반대도 마찬가지입니다. 채널 그룹화나 직교 변환과 같은 기존 방법들은 비대칭적인 오차 전파와 변환된 피연산자의 여러 복사본을 저장하는 비용을 고려한 통합된 제품 수준의 목적 함수를 결형하는 경우가 많습니다. 저자들은 확률적 가정 하에서 정확한 오차 항등식을 유도하고, 양자화 전 인자의 표현을 최적화하면서 동시에 변환된 복사본의 저장/트래픽 비용을 명시적으로 관리하는 "수축-게이지(contraction-gauge)" 프레임워크를 구축하는 것을 목표로 합니다.

방법론

1. 이론적 기초: 곱-오차 항등식 (Product-Error Identity)
저자들은 알려진 분산 필드를 가진 상호 독립적이고 평균이 0인 엔트리별 오차 하에서의 기대 제곱 곱 오차에 대한 정확한 유한 차원 항등식을 유도합니다. 양자화된 인자를 $\hat{A} = A + E_A$, $\hat{B} = B + E_B$라고 할 때, 오차의 기대 제곱 프로베니우스 노름(Frobenius norm)은 다음과 같습니다:
$$\mathbb{E}\|\hat{A}\hat{B} - AB\|_F^2 = \sum_{i,k} v^A_{ik} \|B_{k,:}\|_2^2 + \sum_{k,j} v^B_{kj} \|A_{:,k}\|_2^2 + \sum_{k} \left(\sum_i v^A_{ik}\right)\left(\sum_j v^B_{kj}\right)$$
이 식은 비중첩 감산 디더링(non-overloading subtractive dither) 및 잔차 의존 분산을 가진 독립적 확률적 반올림(stochastic rounding)에 대해 정확합니다. 이는 반대편 인자의 에너지에 따라 오차에 가중치를 부여하는 주요 목적 함수 역할을 합니다.

2. 수축-게이지 프리컨디셔닝 (Contraction-Gauge Preconditioning)
이 방법은 임의의 가역 행렬 $T \in GL(K)$에 대해 $AB = (AT)(T^{-1}B)$라는 대수적 동등성을 활용합니다. 저자들은 "수축 게이지"를 양자화 전 공유 수축 차원 $K$에 적용되는 $T$의 선택으로 정의합니다.

• 게이지 도메인 (Gauge Domains): 하나의 게이지가 할당된 고정된 출력 블록(예: 행 블록)입니다.
• 재사용 기술자 (Reuse Descriptor, $n_{opp}$): 프리컨디셔닝의 비용은 필요한 반대편 인자의 서로 다른 변환된 양자화 복사본 수로 측정됩니다. 만약 단일 게이지가 모든 도메인에 공유된다면 $n_{opp}=1$이며, 별도의 게이지가 사용되면 $n_{opp}$가 증가합니다.

3. 구조적 가족 및 최적화 규칙
본 논문은 다양한 게이지 가족에 대한 구체적인 최적화 전략을 제안합니다.

• 대각 폴드 (Diagonal Folds, 스케일링): 공유 도메인에 대해, $(AD, D^{-1}B)$가 오차를 최소화하는 최적의 대각 스케일링 $D$를 찾는 문제는 **기하 프로그래밍(Geometric Program, GP)**으로 정식화됩니다. 이를 통해 공유 폴드의 전역 최적 선택이 가능합니다. 또한, 아이덴티티 폴드가 이미 최적인지 테스트하기 위한 선형 계획법(LP)이 유도됩니다.
• 파티셔닝 (Partitioning): 정확도와 복사본 수 사이의 균형을 맞추기 위해, 본 논문은 스칼라 노름 정렬(무작위 타이로 인해 성능 손실이 발생할 수 있음)을 지양하고, 유사한 프로필을 가진 행들을 그룹화하기 위해 로그 크기 좌표에서의 클러스터링을 제안합니다.
• 회전 (Rotations): 직교 게이지의 경우, 본 논문은 **블록 국소 코히런스(block-local coherence)**를 사용하여 잠재적 이득을 제한합니다. 무작위 회전(예: Hadamard)이 좌표 최댓값을 로그 인자 내에서 줄일 수 있음을 보여줍니다.
• 부분 회전 (Householder): 가중 에너지가 저계수(low-rank)인 경우, 저자들은 가중 Gram 행렬의 상위 고유 공간을 평탄한 타겟에 정렬하기 위해 소수의 Householder 반사체($t$)를 사용하는 것을 제안하며, 비용은 $O(tK)$입니다.
• 계층 구조 (Hierarchies): 블록 대각 직교 게이지를 위해, 슬라이스 에너지 역상관 기준이 유도되었습니다. 슬라이스에 따른 에너지 프로필이 서로 역상관 관계에 있으면 계층적 정교화가 오차를 줄이지만, 그렇지 않으면 평탄한 회전이 선호됩니다. 텔레스코핑 항등식을 통해 바텀업 방식의 깊이 선택이 가능합니다.

4. 양자화기 설계 (Quantizer Design)

• 비트 할당 (Bit Allocation): 고정된 비트 폭 합 제약($b_A + b_B = b_{sum}$) 하에서, 최적의 분할은 비대칭적입니다: $b_A - b_B = \frac{1}{2} \log_2(P_A/P_B)$, 여기서 $P_A$와 $P_B$는 곱-가중 오차 계수입니다.
• 클리핑 및 진단 (Clipping and Diagnostics): 본 논문은 결정론적 RTN(Round-to-Nearest) 반올림에서 격자 정렬(lattice alignment)을 감지하여 노이즈 모델의 독립성 가정을 위반할 수 있는지를 탐지하는 절단된 특성 함수 진단 도구($\Xi$)를 도입합니다.

주요 기여

1. 정확한 오차 계측: 디더 및 확률적 반올림 하에서 양자화된 행렬 곱셈에 대한 기대 제곱 곱 오차에 대한 정확한 항등식을 유도하여, 선행 오차 항과 쌍선형 교차 항을 명시적으로 분리했습니다.
2. 수축-게이지 프레임워크: 변환 $T$, 공유 패턴(게이지 도메인), 그리고 결과적인 표현 재사용 비용($n_{opp}$)을 결합하는 프리컨디셔닝을 위한 통합된 정식화를 제공합니다.
3. 전역 최적 폴드 선택: 제한된 대각 가족 내에서 전역 최적 선택이 기하 프로그래밍임을 증명하고, 아이덴티티 폴드 최적성을 위한 정확한 LP 테스트를 유도했습니다.
4. 의사결정 통계량: 모든 가족에 대해 정확한 최적해를 계산하지 않고도 변환 선택을 가이드할 수 있는 계산 가능한 통계량을 개발했습니다:
   • 스케일링을 위한 테일 인덱스(Tail index)
   • 파티셔닝을 위한 프로필 확산(Profile spread)
   • 회전을 위한 블록 코히런스 및 가중 Gram 스펙트럼
   • 계층 구조 깊이를 위한 슬라이스 에너지 공분산
5. 재사용 인식 설계: 구현 비용(저장 및 트래픽)과 오차 감소 사이의 분석을 가능하게 하는 메트릭으로서 $n_{opp}$를 도입하여, 게이지 공유 패턴에 따른 재사용성을 정량화했습니다.

결과

제어된 합성 실험:

• 스케일링: 로그 크기 클러스터링 기반 파티셔닝은 스칼라 노름 정렬(2.116x)에 비해 근사 최적의 확산(1.015x)을 달로 달성했습니다.
• 폴드 (Folds): 도메인 공유 GP 폴드는 아이덴티티 폴드 대비 모델 오차를 약 161배, 별도로 최적화된 일방향 경계 대비 4.86배 감소시켰습니다.
• 비트 할당: 비대칭 비트 분할(예: 8/8 대신 9/7)은 대칭 분할보다 1.7배 낮은 오차를 기록하며 이론적 예측과 일치하는 결과를 보였습니다.
• 회전: 무작위 Hadamard 회전은 아웃라이어가 주입된 헤비 테일 데이터에서 최대 60배의 오차 감소를 제공했으며, $t=3$인 Householder 반사체는 랭크-3 데이터에서 전체 회전 성능과 일치했습니다.
• 계층 구조: 역상관 기준은 계층적 정교화가 평탄한 회전보다 우수한 교차점을 성공적으로 예측했습니다.

학습된 분류기 검증:
손글씨 숫자 이미지에 대한 3-블록 이미지 분류기(ViT 스타일) 사용 시:

• 상관관계: 디더 모델 예측은 결정론적 RTN 오차와 높은 중앙값 Spearman 상관관계를 보였습니다 (8비트에서 0.937, 4비트에서 0.918).
• 성능: 아이덴티티 폴드 대신 GP 폴드를 사용하면 모델링된 제품 오차가 8비트에서 18.0%, 4비트에서 20.5%(기하 평균) 감소했습니다.
• 비교: GP 폴드는 캘리브레이션을 통해 선택된 SmoothQuant 스타일의 그리드 서치보다 12개 제품 중 10개에서 더 우수한 성능을 보였으며, 기하 평균에서 테스트 선택된 오라클보다 낮은 오차를 달성했습니다.
• 엔드 투 엔드 (End-to-End): GP 폴드를 사용하여 12개 제품을 함께 양자화했을 때, 아이덴티티 폴드 모델 대비 로짓 MSE가 8비트에서 15.4%, 4비트에서 26.4% 감소했습니다.

의의 및 주장

본 논문은 휴리스틱한 스케일링이나 임의적인 그룹화 이상의 엄격하고 확률적인 기초를 양자화된 행렬 곱셈에 제공한다고 주장합니다. 정확한 오차 항등식을 확립하고 이를 "수축-게이지" 프레임워크와 연결함으로써, 저자들은 다음을 제공합니다:

• 인증된 선택: 기하 프로그래밍을 통한 제한된 대각 가족 내에서의 전역 최적 선택.
• 통합된 목적 함수: 다양한 변환 가족(폴드, 회전, 계층 구조)과 그 공유 패턴을 평가하기 위한 공통된 목적 함수(곱-가중 오차).
• 실질적인 트레이드오프: 구현 비용($n_{opp}$)에 대한 오차 감소를 명시적으로 균형 있게 조절할 수 있는 메커니즘.

저자들은 디더 및 확률적 반올림 하에서 정확한 항등식이 성립하지만, 유도된 통계량과 선택 규칙이 결정론적 RTN 양자화로 효과적으로 전이됨을 학습된 분류기 실험을 통해 입증하였습니다. 본 연구는 $GL(K)$ 최적화 문제를 완전히 해결하는 것이 아니라, 제한적이고 실용적인 가족 내에서 고성능 후보를 선택하는 구조화되고 실행 가능한 절차를 제공하는 데 중점을 둡니다.