PLADIS: 희소성 활용을 통해 추론 시점에서 디퓨전 모델의 어텐션 한계 돌파

PLADIS: Pushing the Limits of Attention in Diffusion Models at Inference Time by Leveraging Sparsity

March 10, 2025

저자: Kwanyoung Kim, Byeongsu Sim

cs.AI

초록

디퓨전 모델은 Classifier-Free Guidance(CFG)와 같은 가이던스 기법을 사용하여 고품질의 조건부 샘플을 생성하는 데 있어 인상적인 결과를 보여주었습니다. 그러나 기존 방법들은 추가적인 학습이나 신경망 함수 평가(NFEs)를 필요로 하기 때문에 가이던스 증류 모델과 호환되지 않는 경우가 많습니다. 또한, 이들은 특정 타겟 레이어를 식별해야 하는 휴리스틱 접근 방식에 의존합니다. 본 연구에서는 PLADIS라는 새로운 효율적인 방법을 제안하며, 이는 희소 주의(sparse attention)를 활용하여 사전 학습된 모델(U-Net/Transformer)의 성능을 향상시킵니다. 구체적으로, 우리는 추론 과정에서 교차 주의(cross-attention) 레이어 내에서 소프트맥스(softmax)와 그 희소 버전을 사용하여 쿼리-키 상관관계를 외삽하며, 추가적인 학습이나 NFEs를 필요로 하지 않습니다. 희소 주의의 노이즈 강건성을 활용함으로써, PLADIS는 텍스트-이미지 디퓨전 모델의 잠재력을 극대화하여, 이전에는 어려움을 겪었던 영역에서도 새로운 효과를 발휘할 수 있게 합니다. 이 방법은 가이던스 증류 모델을 포함한 다양한 가이던스 기법과 원활하게 통합됩니다. 광범위한 실험을 통해 텍스트 정렬 및 인간 선호도 측면에서 뚜렷한 개선을 확인하였으며, 이는 매우 효율적이고 보편적으로 적용 가능한 솔루션을 제공합니다.

English

Diffusion models have shown impressive results in generating high-quality conditional samples using guidance techniques such as Classifier-Free Guidance (CFG). However, existing methods often require additional training or neural function evaluations (NFEs), making them incompatible with guidance-distilled models. Also, they rely on heuristic approaches that need identifying target layers. In this work, we propose a novel and efficient method, termed PLADIS, which boosts pre-trained models (U-Net/Transformer) by leveraging sparse attention. Specifically, we extrapolate query-key correlations using softmax and its sparse counterpart in the cross-attention layer during inference, without requiring extra training or NFEs. By leveraging the noise robustness of sparse attention, our PLADIS unleashes the latent potential of text-to-image diffusion models, enabling them to excel in areas where they once struggled with newfound effectiveness. It integrates seamlessly with guidance techniques, including guidance-distilled models. Extensive experiments show notable improvements in text alignment and human preference, offering a highly efficient and universally applicable solution.