기술 요약: CoReDiT

문제 정의

Diffusion Transformers(DiT) 는 어텐션 메커니즘을 통해 장기 의존성을 효과적으로 모델링함으로써 이미지 및 비디오 생성 분야에서 최첨단 성능을 입증했습니다. 그러나 이들의 계산 비용은 토큰 수에 따라 이차적으로 증가하여, 고해상도 합성 및 다프레임 비디오 생성이 리소스가 제한된 모바일 플랫폼에서는 과도하게 비용이 많이 듭니다. 잠재 공간의 상당 부분의 토큰이 시각적으로 중복된 영역 (예: 균일한 배경) 에 해당함에도 불구하고, 기존 토큰 가지치기 방법은 확산 추론에 적용될 때 다음과 같은 세 가지 주요 과제를 직면합니다:

1. 효율적인 국소화: 전역 어텐션 점수를 계산하는 이차적 오버헤드를 발생시키지 않고 저-주목도 토큰을 식별하는 것.
2. 의미 보존: 토큰을 건너뛸 때 시각적 연속성을 유지하고 아티팩트 (예: 텍스처 끊김, 경계 왜곡) 를 방지하는 것. 이는 판별 작업과 달리 생성 작업에 있어 필수적입니다.
3. 적응형 스케줄링: 다양한 트랜스포머 블록과 노이즈 제거 시간 단계에 걸쳐 가변적인 중복 수준에 적응하는 가지치기 전략을 결정하는 것.

방법론: CoReDiT

CoReDiT 는 출력 충실도를 유지하면서 사전 훈련된 DiT 를 가속화하도록 설계된 구조화된 토큰 가지치기 프레임워크입니다. 이는 세 가지 핵심 구성 요소를 통해 작동합니다:

1. 공간 일관성 기반 토큰 선택

전역 어텐션 점수 ($O(N^2)$) 에 의존하는 대신, CoReDiT 는 선형 시간 ($O(N)$) 에 지역적 중복성을 추정하기 위해 공간 일관성 (SC) 점수를 활용합니다.

• 메커니즘: 토큰 격자는 겹치지 않는 그리드로 분할됩니다. 각 토큰에 대해 SC 점수는 해당 토큰의 정규화된 임베딩과 해당 그리드의 평균 정규화된 임베딩 간의 평균 코사인 유사도로 계산됩니다.
• 논리: 높은 SC 점수를 가진 토큰은 공간적 이웃에 의해 잘 설명되므로 중복하다고 간주됩니다.
• 선택: 가장 높은 SC 점수를 가진 상위-K 개 토큰이 건너뛰기 위해 선택됩니다 ($X_S$), 나머지 토큰 ($X_R$) 은 멀티헤드 셀프 어텐션 (MHSA) 레이어로 진행됩니다.

2. 일관성 유도 토큰 재구성

토큰을 건너뛰는 것으로 인한 공간적 불연속성을 방지하기 위해, CoReDiT 는 유지된 이웃의 출력 ($Y_R$) 을 사용하여 건너뛴 토큰에 대한 어텐션 출력 ($Y_S$) 을 재구성합니다.

• 근사화: 재구성을 위한 쌍별 유사성을 계산하는 것 (오버헤드를 추가함) 대신, 이 방법은 이웃의 사전 계산된 SC 점수를 사용하여 유사성 가중치를 근사화합니다.
• 로컬 서브그리드: 과도한 평활화를 피하기 위해 재구성은 더 작은 서브그리드 내에서 수행됩니다.
• 마이크로 설계:
  • 교번 그리드 크기: 그리드 차원은 연속된 블록 사이에서 교번하여 경계 아티팩트를 이동시키고 그리드 경계를 가로질러 정보 전파를 가능하게 합니다.
  • m-스트라이드 유지: 결정론적인 스트라이드 패턴은 모든 서브그리드에서 적어도 하나의 토큰이 유지되도록 보장하여 재구성을 위한 필수 앵커를 제공합니다.

3. 점진적, 블록 적응형 가지치기 스케줄

중복성이 트랜스포머 블록과 노이즈 제거 단계에 따라 변한다는 점을 인식하여, CoReDiT 는 경량 미세 조정 중에 점진적인 가지치기 전략을 사용합니다.

• 블록 적응성: 가지치기 예산은 가장 높은 측정된 중복성 (후보 토큰의 SC 점수 합으로 추정됨) 을 가진 블록에 할당됩니다.
• 시간 단계 감쇠: 가지치기는 초기 노이즈 제거 단계 (노이즈가 지배적이고 구조가 거칠 때) 에 더 공격적이며, 후기 단계 (세부 사항이 정제될 때) 에는 감소합니다. 단계별 계산 그래프 관리의 복잡성과 효율성 사이의 균형을 맞추기 위해 2 단계 감쇠 스케줄이 사용됩니다.
• 훈련: 모델은 원본 모델 출력 손실과 가지치기된 블록 손실을 결합한 증류 손실을 사용하여 미세 조정되어 성능을 회복합니다.

주요 기여

1. 선형 시간 중복성 추정: 전역 어텐션 분석의 계산 병목 현상을 피하고 $O(N)$ 시간 내에 중복 토큰을 식별하는 공간 일관성 점수의 도입.
2. 밀집 격자 보존: 건너뛴 어텐션 출력을 이웃으로부터 합성하여 재구성하는 메커니즘으로, 후속 레이어를 위해 토큰 격자가 밀집된 상태를 유지하고 시각적 아티팩트를 방지합니다.
3. 적응형 가지치기 전략: 깊이와 시간에 걸친 경험적 중복 통계를 기반으로 가지치기 예산을 동적으로 할당하는 점진적이고 블록 적응형 스케줄로, 안정적인 적응을 촉진합니다.

실험 결과

CoReDiT 는 PixArt-α, PixArt-Σ, MagicDrive-V2를 포함한 최첨단 DiT 백본에서 평가되었습니다.

• 효율성 향상:
  • 셀프 어텐션 FLOP 를 최대 55% 감소시켰습니다.
  • 추론 속도 향상: 클라우드 GPU(Nvidia H100) 에서 1.33 배, 모바일 NPU(Qualcomm Snapdragon 8 Elite) 에서 1.72 배 향상되었습니다.
  • 메모리 여유 공간: 모바일 장치에서 1600 해상도에서 Out-of-Memory(OOM) 오류를 겪었던 기준 모델과 달리, 최대 1920 해상도까지 생성을 가능하게 했습니다.
• 품질 보존:
  • PixArt-α-1024에서 40% 가지치기로 CoReDiT 는 FID 28.7(기준 27.3 대비) 을 유지하며 경쟁력 있는 CLIP/IS 점수를 보였습니다. 증류를 통해 FID 격차는 27.4 로 축소되었습니다.
  • PixArt-Σ-2048에서 23% 의 평균 가지치기 비율은 품질 지표의 modest 한 감소만으로 셀프 어텐션 지연 시간을 22% 감소시켰습니다.
  • MagicDrive-V2(비디오) 에서 39% 의 셀프 어텐션 FLOP 감소는 프레임별 품질 (PSNR, LPIPS, SSIM) 과 조건부 정렬 (mAP, mIoU) 을 유지했으나, 시간적 일관성 (FVD) 에서 약간의 감소가 관찰되었습니다.
• 절대 실험 (Ablation Studies):
  • 무작위 토큰 선택은 치명적인 품질 저하 (FID 644.5) 를 초래하여 일관성 기반 선택의 필요성을 입증했습니다.
  • 재구성을 비활성화하면 FID 30.4(28.5 대비) 로 상당한 품질 손실이 발생하여 재구성 단계의 중요성을 확인했습니다.
  • 균일 가지치기 후 미세 조정은 제안된 점진적 적응형 스케줄보다 성능이 낮았습니다.

중요성 및 주장

이 논문은 CoReDiT 가 생성형 토큰 가지치기의 특정 과제를 해결함으로써 엣지 장치에 고충실도 확산 모델을 배포하기 위한 실용적인 솔루션을 제공한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 확장성: 계산 및 메모리 풋프린트를 줄여 모바일 하드웨어에서 더 높은 해상도 합성 및 비디오 생성을 가능하게 합니다.
• 아키텍처 호환성: 구조적 변경이나 복잡한 라우팅 모듈 없이 기존 DiT 아키텍처와 호환성을 유지합니다.
• 효율성 - 품질 트레이드오프: 구조화되고 일관성 유도된 건너뛰기가 전체 모델과 비교 가능한 지각적 품질을 유지하면서 추론 비용을 극적으로 줄일 수 있음을 보여줍니다 (셀프 어텐션 FLOP 최대 55% 감소).

저자들은 현재 접근 방식이 공간적 중복성에 초점을 맞추고 있지만, 향후 작업은 FFN/MLP 구성 요소로 효율성을 확장하고 비디오 시간적 일관성을 더욱 개선하기 위해 프레임 간 일관성 목표를 통합하는 것을 목표로 한다고 언급했습니다.