OrbitQuant: 이미지 및 비디오 확산 트랜스포머를 위한 데이터-무관 양자화

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

Diffusion Transformer(DiT)는 다단계 샘플링과 커진 파라미터 수 때문에 추론 비용이 매우 높아 PTQ(포스트 트레이닝 양자화)의 필요성이 컸다. 기존 PTQ 방법은 활성화 분포가 타임스텝·프롬프트·CFG 브랜치마다 이동하므로 각 체크포인트나 모달리티별로 보정 데이터가 필요했다. OrbitQuant은 회전 후 고정된 차원 주변분포에 기반한 단일 코드북으로 이 문제를 회피해 보정 데이터 없이도 저비트에서 안정적인 생성 품질을 확보했다.

왜 중요한가

핵심 기여

차원별 고정 분포에 기반한 데이터-무관 활성화 코드북 설계

활성화를 정규화하고 회전한 뒤 각 좌표가 차원 d에 의해 결정되는 고정 주변분포 f_d를 따르도록 하여 그 분포에 대해 오프라인으로 Lloyd–Max 코드북을 계산했다. 이렇게 만든 코드북은 모든 타임스텝, 프롬프트, 레이어에서 동일하게 재사용되므로 입력별 범위 보정(calibration)이 불필요해졌다. 이 접근은 DiT 활성화의 시간적·프롬프트 의존적 변동을 제거하여 저비트에서 안정적인 양자화를 가능하게 했다.

가중치 행에 회전을 흡수하는 공동 기준(shared-basis) 양자화로 런타임 역회전 불필요화

차원별 회전 행렬 Π_d를 가중치에 곱해 가중치 행을 회전된 기준으로 옮긴 뒤 행 단위로 방향을 코드북으로 양자화하고 크기(ℓ2 노름)는 저정밀형(BF16)으로 저장했다. 이러한 처리로 선형층 내부에서 회전이 상쇄되어 런타임에 가중치에 대한 역회전 복원이 필요 없고 활성화에 한 번의 순방향 회전만 남았다. 결과적으로 온라인 오버헤드를 낮추면서도 가중치와 활성화가 동일한 기준으로 양자화되도록 했다.

효율적인 RPBH(Randomized Permuted Block-Hadamard) 회전 도입

무작위 치환과 블록 Hadamard 변환, 랜덤 부호(diagonal Rademacher)를 결합한 RPBH는 저장 비용과 계산 복잡도를 O(d log h)로 줄이며 거의 모든 입력에 대해 회전 후 좌표의 분포가 f_d에 근접하도록 보장했다. 균일 무작위 치환은 블록 간 질량을 고르게 분산시켜 특정 블록에 집중된 아웃라이어가 전체 분포를 깨뜨리지 않도록 한다. 이 구조는 치환을 데이터 종속적으로 학습·보정할 필요를 제거하며, dense Haar 회전에 비해 속도와 저장 면에서 유리했다.

이미지와 비디오 DiT에서 보정 데이터 없이 저비트 PTQ 성능 우위 입증

FLUX.1, Z-Image-Turbo, Wan 2.1, CogVideoX 등 이미지·비디오 DiT 벤치에서 OrbitQuant은 여러 저비트 설정에서 기존 PTQ 기법을 앞섰다. W4A4에서 다수 모델에 대해 FP16 대비 거의 손실 없는 성능을 보였고, W2A4에서는 기존 방법들이 무너지는 상황에서도 의미 있는 생성 품질을 유지했다. 또한 Latency·메모리 비교에서 활성화 코드북 조회 방식의 고정 오버헤드가 동적 보정 기반 방법들보다 낮아 실용적 이점을 제공했다.

핵심 아이디어 이해하기

DiT의 활성화는 각 타임스텝과 프롬프트에 따라 채널별 아웃라이어와 전체 분포가 크게 이동하므로, 전통적 범위 보정(range calibration)을 한 번만 수행하는 방식이 성능을 보장하지 못한다. 이러한 이동은 각 토큰이 들어오는 매번 서로 다른 스케일과 분포를 만들며, 그래서 타임스텝 또는 모달리티가 바뀔 때마다 재캘리브레이션이 필요해진다. 문제의 핵심은 활성화 좌표의 분포가 입력에 따라 달라지는 것이며, 이 때문에 데이터-무관적 양자화가 어렵다.

방법론

OrbitQuant은 활성화를 먼저 ℓ2 노름으로 정규화한 다음, 차원별로 미리 샘플된 RPBH 회전을 적용해 각 좌표가 차원 d에 의해 결정되는 고정 주변분포에 가깝게 만든다. 오프라인 단계에서는 각 입력 차원 d에 대해 RPBH 회전 Π_d를 샘플링하고 해당 f_d에 대해 MSE-optimal한 Lloyd–Max 코드북을 계산해 저장한다. 가중치는 오프라인에서 W' = W Π_d^T로 회전시킨 뒤 각 행을 크기(r_i')와 단위 방향(w~_i')로 분리하여 방향을 코드북으로 양자화하고 크기는 BF16으로 저장한다. 온라인 단계에서는 토큰 활성화를 Π_d로 회전하고 노름 s를 분리한 뒤 정규화된 방향만 코드북으로 양자화하고 s를 곱해 복원한다. 회전이 가중치에 흡수되어 있으므로 양자화된 가중치와 양자화된 활성화의 행렬곱은 원래 Wx를 근사하며 런타임에는 활성화에 대한 순방향 회전만 수행되어 역회전 비용이 발생하지 않는다.

주요 결과

이미지 생성 평가(GenEval)에서 OrbitQuant은 W4A4 설정에서 FLUX.1-schnell과 Z-Image-Turbo에서 FP16 대비 동등하거나 더 나은 Overall 점수를 보였고 FLUX.1-dev에서는 FP16에 0.034 포인트만 뒤졌다. 구체적으로 FLUX.1-schnell의 W4A4 Overall은 0.703, W2A4에서 0.604를 기록해 낮은 비트에서도 의미 있는 성능을 유지했다. Z-Image-Turbo는 W4A4에서 0.767 Overall을 달성했고 W2A4에서도 0.319의 점수를 내어 기존 PTQ 기법들이 거의 붕괴하는 구간에서 유의미한 결과를 보였다. 비디오 평가(VBench)에서 Wan 2.1-1.3B의 경우 OrbitQuant은 W4A6에서 Overall Consistency 24.35를 기록해 동일한 비트 설정의 다른 PTQ 방법들을 앞섰고 W4A4에서는 23.86을 기록해 전체 일관성 측면에서 우수했다. 성능 외에도 Latency·메모리 분석에서 이미지 실험(FLUX.1-dev) 기준 OrbitQuant은 기존 가중치·활성화 양자화 기법들보다 낮은 오버헤드를 보였고, 비디오에서는 활성화 주도의 메모리 증가로 인해 QuaRot와 SmoothQuant보다 약간 높은 피크 메모리를 기록했으나 ViDiT-Q보다는 낮은 값을 유지했다.

기술 상세

RPBH 회전은 blkdiag(H_h D_1, ..., H_h D_{d/h}) · P_π 형태로 구성되며 각 구성요소는 블록 Hadamard 행렬 H_h, 랜덤 부호 대각행렬 D_i, 그리고 균일 무작위 치환 행렬 P_π로 이루어진다. 여기서 h는 d의 가장 큰 2의 거듭제곱 약수로 선택되어 Fast Walsh–Hadamard Transform를 블록 단위로 적용해 계산 복잡도를 O(d log h)로 유지한다. 회전 이후 정규화된 단위 벡터의 좌표는 이론적으로 차원 종속 주변분포 f_d를 따르며 d ≥ 64이면 f_d는 거의 N(0,1/d)로 근사된다. 논문은 행렬곱 항등성 W' x' = W Π_d^T Π_d x = W x를 활용해 가중치에 Π_d^T를 흡수하고 활성화에 Π_d만 적용하도록 설계하여 온라인에서 역회전을 할 필요를 제거했다. Proposition 1은 단일 고정 단위벡터 x~에 대해 RPBH의 치환에 대한 확률적 분산 집중을 보장하며, 이로써 각 좌표 z_i의 분산이 [ (1-ρ)/d, (1+ρ)/d ] 범위에 들어가도록 하여 f_d 근사를 유지한다. 여기서 ρ = d μ_∞ sqrt( (1/(2h)) log(4k/δ) )이며 μ_∞는 단위벡터의 최대 좌표 제곱이다, 따라서 한 좌표가 과도하게 노름을 차지하지 않는 한 ρ는 작아져 분산이 안정화된다.

한계점

논문에서 명시된 제약으로는 AdaLN처럼 타임스텝 의존적이며 크기-시프트를 수행하는 modulation 투영을 메인 코드북으로 흡수할 수 없으므로 이들 투영은 INT4 같은 별도 저비트 처리를 유지해야 한다는 점이 있다. RPBH의 이론적 보장은 입력 단위 방향의 최대 좌표 제곱 μ_∞가 작다는 가정에 의존하므로 극단적으로 한 좌표에 질량이 집중된 입력에서는 분산 집중 보장이 약해질 수 있다. 또한 회전 적용은 런타임에 추가 비용을 발생시키며 특히 비디오에서 활성화 중심의 인덱스·gather 텐서를 생성해 피크 메모리를 소폭 증가시킬 수 있다.

실무 활용

OrbitQuant은 보정 데이터 수집이 어려운 프로덕션 환경에서 DiT 기반 이미지·비디오 생성 모델을 저비트로 배포할 때 실무적으로 유용하다. 오프라인으로 차원별 코드북과 회전만 준비하면 런타임에는 활성화에 하나의 순방향 회전과 최근접 중심 조회만 수행하므로 운영 오버헤드가 상대적으로 작다. 다만 AdaLN와 같이 타임스텝 의존적 스케일·시프트를 갖는 소규모 프로젝션은 별도 처리(INT4 권장)가 필요해 전체 압축률-품질 트레이드오프를 고려해야 한다.

대규모 DiT 모델을 GPU 메모리 제약 하에서 배포해 추론 비용과 저장 공간을 절감하는 용도에 적합하다.
이미지·비디오 생성 파이프라인에서 체크포인트나 해상도 변경 시 보정 데이터 재수집 없이 빠르게 양자화된 모델을 적용해야 하는 상황에 적합하다.
온프레미스 또는 엣지에 모델을 옮겨야 하며 가중치·활성화 모두의 저비트 압축이 필요한 워크로드에서 사용 가능하다.

코드 공개 여부: 비공개

키워드

RPBH(Randomized Permuted Block-Hadamard) 회전Lloyd–Max 코드북(로이드-맥스 코드북)post-training quantization(포스트 트레이닝 양자화)diffusion transformers(확산 트랜스포머)activation-weight shared-basis(가중치-활성화 공유 기준)

용어 해설

RPBH rotation: — 무작위 치환 블록-Hadamard(RPBH) 회전은 입력 벡터에 블록 단위 Hadamard 변환과 랜덤 부호 행렬을 적용한 뒤 균일한 무작위 치환을 수행하는 선형 변환이다. 이 변환은 계산 복잡도를 O(d log h)로 낮추면서도 회전 후 각 좌표의 분포가 차원에만 의존하도록 집중시키며, 따라서 오프라인으로 설계한 차원별 코드북으로 안정적인 양자화가 가능하다. 논문 맥락에서는 이 회전이 활성화와 가중치의 분포를 표준화하여 단일 Lloyd–Max 코드북을 모든 타임스텝과 레이어에 재사용하도록 만든다.
Lloyd–Max codebook: — Lloyd–Max 코드북은 주어진 연속 분포에 대해 평균제곱오차(MSE)를 최소화하는 유한 개의 대표값(centroid) 집합을 계산한 것이다. OrbitQuant에서는 회전 후의 차원별 이론적 주변분포 f_d를 기준으로 코드북을 미리 계산해 모든 레이어와 타임스텝에 재사용한다. 이렇게 하면 입력별 범위 보정(calibration)이 불필요해지고 온라인 단계에서는 단순한 최근접 중심(nearest-centroid) 조회로 양자화가 이루어진다.
TurboQuant: — TurboQuant은 정규화(norm)와 무작위 직교 회전을 결합해 벡터의 좌표 분포를 차원에만 의존하게 만든 뒤, 해당 분포로 설계한 Lloyd–Max 코드북으로 검증-불필요한 벡터 양자화를 수행하는 이전 작업이다. OrbitQuant은 TurboQuant의 분포 기반 코드북 아이디어를 계승하되, DiT의 선형층 내부에서 회전을 흡수(가중치에 적용)하고 효율적인 RPBH 구조를 사용해 실시간 회전 비용을 줄였다. 논문에서는 TurboQuant이 KV-cache 압축용인 반면 OrbitQuant은 DiT 투입용으로 회전을 흡수하여 역회전 불필요성을 확보했다고 보고된다.