합성곱 텐서 레이아웃 재배열로 마이크로스케일링 양자화의 이중 연산과 메모리 이동을 줄이는 MicroQonv

무엇인가

마이크로스케일링 양자화는 텐서를 일정 길이의 블록으로 나누고 블록마다 하나의 스케일을 공유해, 가중치·활성값·그래디언트를 4비트 수준까지 낮추면서도 거의 원래 정밀도의 정확도를 유지한다. 문제는 이를 합성곱 계층에 그대로 적용하면 효율이 크게 떨어진다는 점이다. 가중치와 활성값, 그래디언트는 각각 서로 다른 방향의 두 개의 GeMM(일반 행렬 곱)에 참여하는데, 마이크로스케일링 블록은 내적 방향과 정렬되어야 하므로 같은 텐서를 두 번 양자화해야 한다. 활성값 쪽 부담은 더 크다. 양자화 전에 im2col 변환이 적용되면 텐서 크기가 커널 크기 Kh×Kw만큼 부풀어 오르고, 같은 픽셀이 서로 다른 이웃과 묶이면서 블록마다 다른 양자화 값을 갖게 된다. 결국 역전파를 위해 FP32 원본을 저장하거나 부풀려진 양자화 텐서를 저장해야 해서, 양자화로 얻으려던 메모리 이득이 상당 부분 사라진다. 논문은 이 지점을 AI 메모리 월 문제의 일부로 규정한다.

어떻게 동작하나

제안 방법인 MicroQonv의 핵심은 활성값을 im2col 확장 이전에 한 번만 양자화하고, 채널-배치 우선(channel-batch-first)이라는 새 메모리 레이아웃으로 저장하는 것이다. 입력 활성값 X를 X_{Cin,N} ∈ R^{Cin·(Hin·Win·N)} 형태로 재배열하면, 각 공간 위치 (h,w)마다 모든 채널과 모든 배치 원소가 연속된 블록으로 놓인다. 순전파에서는 합성곱 윈도우를 슬라이딩하며 Cin 방향으로 연속 읽기가 가능해 im2col 블록을 온칩에서 즉석 생성할 수 있고, 역전파에서 입력 그래디언트를 계산할 때는 배치 차원 N 방향으로 읽으면 되므로 역시 연속 접근이 된다. 기존 channel-first 레이아웃이 순전파에만 유리하고 역전파에서는 Win 차원을 따라 불연속 접근이 생기는 것과 대비된다. 이 레이아웃은 채널 방향이나 배치 방향의 연속 블록을 오프칩에서 한 번만 읽어 온칩에서 반복 사용하는 데이터 재사용도 가능하게 한다.

무엇과 다른가

정확히 한 번만 양자화하기 위해 논문은 2D 행렬 형태에서 정사각형 블록 양자화를 쓴다. W_f를 W_{fCout,Cin} ∈ R^{Cout·(Kh·Kw·Cin)}, G_{Yf}를 G_{YfCout,N} ∈ R^{Cout·(Hout·Wout·N)}로 재배열하면 세 개의 GeMM, 즉 Y = W·X', G_W = G_Y·X'^T, G_X' = W^T·G_Y가 모두 같은 2D 표현 위에서 성립하고, 정사각형 블록의 한 변이 항상 계산 방향과 정렬된다. 블록 스케일은 정사각형이 덮는 두 차원 각각에서 √B 원소마다 적용된다. 양자화 형식은 NVFP4를 따른다. 텐서 전체 스케일 T_scale = max|X| / (max B_sf × max B_ef)를 먼저 구하고, 블록 스케일 B_scale = CastE4M3(max|X_i| / (T_scale × max B_ef))을 FP8 E4M3로 저장한 뒤, 원소는 X_i^q = Round(X_i / (B_scale × T_scale))로 FP4 E2M1로 반올림한다. 블록 크기는 B=16이며, 활성값과 가중치는 반올림, 그래디언트는 확률적 반올림(stochastic rounding)을 쓰고, 가중치의 마스터 복사본은 FP32로 유지한다.

어떻게 쓰나

실험은 이미지 분류와 객체 검출 두 갈래로 진행된다. 이미지 분류에서는 ResNet-32와 ResNet-18을 FP32 기준선, LUQ, QLR과 비교했고, 4비트에서 MicroQonv가 가장 좋은 정확도를 냈다. LUQ를 같은 유효 비트폭인 3비트로 맞춘 경우 MicroQonv가 ResNet-32에서 2% 이상, ResNet-18에서 0.8% 앞선다. 활성값을 FP4 E2M1로 표현하면 픽셀당 4비트 대신 3비트(블록 스케일 포함 3.125비트)만 쓰게 되어 저장 효율이 올라간다. QLR은 정확도가 약간 더 높지만 학습 비용이 훨씬 크다. 객체 검출에서는 YOLOv8nano(PascalVOC)와 YOLOv26nano(KITTI)를 FP6 E2M3 원소 형식(NVFP6)으로 실험했다. YOLOv8nano에서 MicroQonv-6는 FP32 기준선과 0.3mAP 이내이고 Dacapo-6, Cuyckens et al.-6와는 0.15mAP 이내 차이이며, per-channel 양자화보다 0.4~0.8mAP 높다. YOLOv26nano에서는 미세조정 시 0.3mAP 이내, 처음부터 학습 시 1.3mAP 이내이고 Dacapo-6보다 0.1mAP 높으며 per-channel 양자화보다 1.2~3mAP 높다. DRAM 시뮬레이터 검증에서 활성값 메모리 이동이 YOLOv8nano에서 ×3.5, YOLOv26nano에서 ×2.2 줄었고, 전체 메모리 이동·저장량은 최대 ×7.53 감소했다. 양자화 연산 비용 자체는 가중치·그래디언트에서 ×2, 활성값에서 최대 ×9 줄어든다.

전제와 한계

지속 학습(continual learning) 실험에서는 지연 재생(latent replay) 버퍼에 4비트 MicroQonv를 적용했다. ImageNet으로 사전학습한 ResNet18을 CIFAR-100, CORE50, CUB-200에서 클래스 증분 학습시키고, 8비트 정수 양자화를 쓰는 QLR과 비교했다. 같은 메모리 예산에서 MicroQonv의 정확도가 CIFAR-100에서 +5.7%에서 +11%, CORE50에서 +4.8%에서 +9.4%, CUB-200에서 +9.6%에서 +16.7% 높다. 활성값 픽셀당 8비트 대신 3.125비트를 저장하므로 약 2.5배 더 많은 샘플을 버퍼에 담을 수 있고, 이것이 정확도 차이를 설명한다. CIFAR-100에서 메모리 예산을 2000 샘플로 고정한 비교에서는 FP32 기준선 48.83±0.36%, QLR 48.77±0.23%, MicroQonv 48.16±0.66%로 기준선과 1% 이내 차이다.

실무 관점에서 이 논문은 합성곱 계층이 들어간 모델을 4비트 수준으로 학습하거나 온디바이스에서 계속 학습시킬 때 참고할 만하다. im2col을 명시적으로 materialize하지 않고 채널·배치 방향으로 연속 접근이 가능한 레이아웃을 쓰면, 양자화를 한 번만 수행하면서도 순전파와 역전파 모두에서 메모리 대역폭을 아낄 수 있다는 점이 핵심이다. 다만 이득은 커널 크기와 stride에 크게 좌우된다. 3×3 커널이 많은 ResNet18에서는 활성값 에너지 이득이 약 ×7인 반면, 1×1 커널 비중이 40%인 YOLOv8nano는 ×3.5, 1×1 비중이 53%이고 depthwise 합성곱을 BF16으로 두는 YOLOv26nano는 ×2.2로 떨어진다. 자신의 모델에서 커널 크기 분포와 depthwise 계층 비중을 먼저 확인해야 한다.

저자들이 밝힌 전제와 한계도 분명하다. 설명의 편의를 위해 stride 1, 패딩 없음을 가정했지만 방법 자체는 이 가정에 의존하지 않는다고 밝힌다. 가중치 마스터 복사본을 FP32로 유지하는 것을 전제로 한다. 입력 계층과 출력 계층은 양자화에 민감해 높은 정밀도로 남겨 두는데, YOLOv8은 64개 합성곱 계층 중 6개 출력 계층과 DFL 계층을, YOLOv26은 126개 중 12개 출력 계층을 FP32로 유지하고 depthwise 합성곱은 BF16으로 둔다. 지속 학습 실험에서는 병렬 잔차 연결이 없는 합성곱 계층만 지연 재생 계층으로 쓸 수 있고, 첫 계층과 마지막 계층은 8비트로 양자화한다. DRAM 시뮬레이션은 양자화 연산 비용을 포함하지 않는데, 포함하면 커널이 1보다 클 때 im2col 확장 텐서를 대략 Kh×Kw/stride배 더 양자화해야 하므로 MicroQonv의 이득이 오히려 더 커진다고 저자는 주장한다. 또한 Cuyckens 등의 정사각형 양자화는 완전연결 계층만 고려했고, 합성곱으로 확장하려면 im2col로 부풀린 텐서를 저장해야 해서 메모리 이득이 상쇄된다는 점을 관련 연구의 한계로 지적한다. 향후 과제로 MicroQonv 레이아웃을 활용하는 하드웨어 가속기 설계와 벡터 양자화 같은 다른 손실 압축 기법과의 관계 연구를 남긴다.