MC-Sparse가 확산 트랜스포머 희소 어텐션의 품질 격차를 좁힌다

무엇인가

확산 트랜스포머(DiT)는 고해상도 영상과 3D 자산 생성의 주력 백본이 됐지만, 시퀀스 길이가 수만에서 수십만 토큰으로 늘면서 어텐션의 이차 비용이 추론 지연의 주 원인이 된다. 어텐션이 소수 토큰 상호작용에 집중된다는 관찰 덕에 희소 어텐션이 자연스러운 해법이지만, 희소도를 높이면 생성 품질이 떨어진다. 기존 방법 대부분은 블록 단위로 희소화하며, 블록 중요도를 평균 풀링된 쿼리·키 표현으로 추정한다. 이 설계는 근사 점수에서 오는 오차와 블록 레이아웃이 선택 자체를 제한해서 생기는 오차를 뒤섞어, 품질 격차가 선택기 탓인지 블록 구조 탓인지 구분할 수 없게 만든다.

어떻게 동작하나

저자들은 동일한 어텐션 밀도에서 통제된 오라클 비교로 이 문제를 분해한다. 블록 오라클은 쿼리 블록마다 KV 블록을, 토큰 오라클은 쿼리 블록은 공유하되 개별 KV 토큰을, per-query 오라클은 쿼리마다 독립적으로 KV 토큰을 정확한 dense 어텐션 확률로 고른다. 이들을 평균 풀링 기반 vanilla BSA와 비교해 세 가지 오차원을 분리한다. KV 토큰이 블록으로 묶이고 같은 그룹의 쿼리가 하나의 선택을 공유하면서 생기는 구조적 결속 오차, 평균 풀링 같은 근사 스코어링에서 생기는 선택 오차, 그리고 버려진 토큰이 여전히 0이 아닌 어텐션 질량을 갖기 때문에 생기는 꼬리 손실 오차다. 여기서 세 가지 설계 원칙이 나온다. 더 세밀한 선택, 정확한 어텐션 스코어링, 버려진 기여에 대한 보상이다.

무엇과 다른가

MC-Sparse는 이 원칙을 학습 없이 구현하는 프레임워크다. 핵심은 앵커 스텝과 재사용 스텝의 교대다. 앵커 스텝(밀집 워밍업 또는 리프레시 구간의 마지막 스텝)에서는 dense 어텐션을 실행하고, 쿼리 그룹을 구성하고, 정확한 확률로 KV 토큰을 선택하고, dense 출력과 sparse 출력의 차이를 잔차로 기록해 캐시한다. 이후 재사용 스텝에서는 현재 스텝의 Q·K·V를 다시 계산하되 어텐션은 캐시된 쿼리 그룹과 KV 인덱스로 수행하고, 캐시된 잔차를 더한다. 즉 Ô_t = SparseAttn(Q_t, K_t, V_t; G_τ, S_τ) + R_τ, 여기서 R_τ = O_full − O_sparse다. 이렇게 하면 그룹화와 정확 선택 비용이 여러 디노이징 스텝에 걸쳐 분할 상환된다.

어떻게 쓰나

구성 요소는 세 가지다. 첫째, 토큰 단위 KV 선택은 블록 경계를 넘어 개별 KV 토큰을 고르고, 커널이 인덱스로부터 키와 값을 온라인으로 gather한다. 둘째, 타일 정렬 쿼리 그룹화는 커널의 쿼리 타일 크기 C에 정확히 맞는 동일 크기 그룹을 만든다. k-means는 균등 크기를 보장하지 않고 강제 균형 제약은 할당 오버헤드를 키우므로, PDDP의 중앙값 분할 변형을 써서 각 그룹을 주방향으로 투영한 뒤 중앙값에서 재귀 분할한다. GPU에서 순진한 PDDP는 비싸기 때문에 배치 power iteration으로 같은 트리 깊이의 노드를 헤드 전체에 걸쳐 병렬 처리하고 쿼리 인덱스 배열만 레벨 간에 전파하는 Fast PDDP를 구현했다. 셋째, 잔차 보상은 FlashAttention이 어텐션 로짓을 노출하지 않는다는 제약 때문에 two-pass 커널로 정확 선택을 수행하고, 앵커에서 기록한 dense−sparse 잔차를 재사용 스텝에 더해 버려진 꼬리 기여를 근사한다. 실행 커널은 CuTeDSL로 구현했고, GPU 메모리가 부족하면 비활성 캐시 항목을 CPU로 오프로딩하고 한 레이어 앞서 프리페치해 전송 지연을 숨긴다.

전제와 한계

실험은 영상 생성에 Minimax-H3-Base, HunyuanVideo-13B, Wan2.1-14B-T2V/I2V, 3D 자산 생성에 HY3D-Internal을 쓴다. 베이스라인은 FA3 기반 dense 어텐션, SVG2, SVG-EAR, Sol-Attn, PISA, SpargeAttn이다. 지표는 dense 출력에 대한 PSNR·SSIM·LPIPS, VBench의 이미지 품질과 배경 일관성, 3D는 Chamfer distance·volumetric IoU·F1과 Uni3D-I·ULIP3D-I다. Minimax-H3-Base에서 MC-Sparse는 25% 밀도에 28.44dB PSNR과 1.61배 가속을 냈고, 같은 조건에서 Sol-Attn은 32.4% 밀도에 23.66dB, 1.59배였다. 같은 앵커 스케줄의 MC-Sparse-Flash는 15% 밀도에서 1.80배 가속과 27.30dB PSNR로 Sol-Attn과 PISA보다 높은 충실도를 유지한다. Hunyuan과 Wan에서는 평가된 희소 방법 중 최고 PSNR·SSIM과 최저 LPIPS를 기록하며 1.52배에서 1.82배 가속을 얻었다. 3D에서는 15% 밀도에서 2.32배 디노이징 가속과 함께 최고 기하 충실도를 보였고, Chamfer distance를 PISA의 0.976에서 0.177로 낮췄다.

시스템 효율과 어블레이션도 수치로 제시된다. 토큰 단위 KV gather를 쓰는 커널은 이상적인 1/s 가속 대비 96~99%를 달성해 SVG2의 80~81%, PISA의 82~87%를 앞선다. Fast PDDP는 Flash-KMeans 대비 분할 상환된 그룹화 비용을 480p에서 3.96배, 720p에서 6.14배 줄인다. 앵커 스텝 오버헤드는 전체 디노이징 스텝 기준 밀도 증가 Δs≈5%에 해당하고, 그룹화 등 부수 연산을 포함하면 5~6%로 SVG2의 클러스터링·가변 길이 커널 오버헤드보다 낮다. 어블레이션에서는 vanilla BSA에 정확 KV 선택과 스텝 간 재사용, 토큰 단위 KV 입도, 쿼리 그룹화, 잔차 보상을 차례로 더할 때마다 모든 밀도에서 PSNR·SSIM·LPIPS가 개선됐다. 잔차 보상은 누적 어블레이션에서 가장 큰 이득을 내지만, vanilla BSA보다 MC-Sparse의 선택·그룹화 설계와 결합했을 때 효과가 훨씬 크다.

개발자 관점에서 이 논문은 기존 DiT 체크포인트를 재학습 없이 추론 단계에서 가속하는 방법을 다룬다. 영상·3D 생성 파이프라인에서 어텐션이 병목일 때 후보가 된다. 도입 시 확인할 것은 앵커 스케줄과 목표 밀도, 캐시 메모리 사용량과 오프로딩·프리페치 필요 여부, CuTeDSL 커널을 쓰는 백엔드 통합 비용, 그리고 배치 크기·해상도별 실측 가속과 품질이다. 특히 앵커 스텝에서 dense 어텐션을 한 번 더 돌리는 비용이 전체 스텝에 분할 상환된다는 전제를 자기 워크로드의 스텝 수와 워밍업 구성에 맞춰 검증해야 한다.

저자들이 밝힌 전제와 한계는 다음과 같다. 이 방법은 어텐션 패턴이 인접한 디노이징 스텝 사이에서 안정적이라는 관찰에 의존하며, 잔차 보상은 단순 가산 보정에 그친다. 앵커 스텝에서는 dense 어텐션과 추가 선택 패스가 필요해 Δs≈5~6%의 오버헤드가 발생하고, GPU 메모리가 부족한 상황에서는 CPU 오프로딩과 프리페치가 전제된다. 평가는 Minimax-H3-Base, HunyuanVideo, Wan2.1, HY3D-Internal 등 특정 영상·3D 모델에 한정되며, 학습 기반 효율적 어텐션 방법과의 직접 비교나 다른 도메인으로의 일반화 수치는 원문에 제시되지 않았다.