MC-Sparse, 확산 트랜스포머의 희소·밀집 어텐션 간격을 좁히다

긴 시퀀스를 처리하는 확산 트랜스포머에서 어텐션 계산은 중요한 지연 요인이다. 특히 비디오 생성과 고해상도 3D 에셋 생성에서는 토큰 수가 빠르게 늘어나기 때문에, 모든 질의가 모든 키·밸류 쌍을 계산하는 밀집 어텐션의 비용이 커진다. 희소 어텐션은 일부 상호작용만 남겨 이 비용을 줄이는 방법이지만, 희소성을 높일수록 생성 품질과 세부 묘사의 충실도가 떨어질 수 있다.

MC-Sparse는 이 문제를 더 세분화해 살펴본다. 논문은 통제된 oracle 비교를 통해 밀집 어텐션과 희소 어텐션의 차이가 세 가지 요인에서 발생한다고 분석했다. 첫째, 질의를 그룹으로 묶으면 실제로는 다른 질의가 같은 희소 패턴을 공유하게 된다. 둘째, 중요한 상호작용을 고르는 선택기가 잘못된 KV 토큰을 선택할 수 있다. 셋째, 삭제된 토큰이 어텐션 출력에 기여한 부분이 별도 보정 없이 사라진다.

제안된 프레임워크는 다음과 같은 방식으로 이 문제를 다룬다.

• 개별 KV 토큰 선택: 고정된 큰 토큰 블록에만 의존하지 않고, 정확한 어텐션 확률을 이용해 개별 키·밸류 토큰을 선택한다. 이를 통해 유지할 상호작용을 더 정밀하게 지정할 수 있다.
• GPU 친화적인 질의 그룹화: 세밀한 선택은 불규칙한 메모리 접근을 유발해 GPU 실행 효율을 떨어뜨릴 수 있다. MC-Sparse는 유사한 질의를 타일에 맞는 그룹으로 정렬해, 선택의 정밀도와 하드웨어 병렬성을 함께 고려한다.
• 메타데이터와 잔차 재사용: 질의 그룹, 선택된 KV 인덱스, 밀집 어텐션 출력과 희소 어텐션 출력 사이의 잔차를 캐시한다. 이후 디노이징 단계에서 이를 재사용해 반복적인 선택 비용을 줄이고, 제거된 토큰의 정보 손실을 보완한다.

MC-Sparse는 모델을 다시 학습하지 않는 방식으로 설계됐다. 단순히 더 많은 상호작용을 생략하는 것이 목표가 아니라, 실제 GPU 실행 비용을 관리하면서 희소 어텐션의 결과를 밀집 어텐션에 가깝게 만드는 것이 핵심이다. 논문이 제시한 비디오 및 3D 생성 실험에서는 기존 희소 어텐션 기준선보다 밀집 어텐션 출력에 대한 충실도와 디노이징 가속 모두에서 개선을 보였다. 밀집 어텐션과 비교하면 Minimax-H3-Base에서 1.80배, 3D 에셋 생성에서 2.32배의 디노이징 가속을 보고했으며, 품질 손실은 무시할 수 있는 수준이라고 설명한다.

이 연구의 의미는 희소 어텐션 최적화를 단순한 마스크 설계 문제로 보지 않았다는 데 있다. 효율적인 희소화에는 중요한 상호작용을 정확히 찾는 일, GPU가 처리하기 좋은 형태로 정리하는 일, 버려진 정보의 기여를 보정하는 일이 함께 필요하다. 다만 제공된 자료만으로 모든 모델과 희소성 설정에서 동일한 이득을 보장할 수는 없다. 시퀀스 길이, 캐시 오버헤드, GPU 커널 구현은 실제 성능에 영향을 줄 수 있다. 그럼에도 MC-Sparse는 장기 시퀀스 확산 모델의 학습 후 추론 가속을 위한 실용적인 방향을 제시한다.

출처: Hugging Face Daily Papers