DUIDELIJK: Convolutie-achtige Linearisatie Herziet Vooraf Getrainde Diffusie Transformers Omhoog

CLEAR: Conv-Like Linearization Revs Pre-Trained Diffusion Transformers Up

December 20, 2024

Auteurs: Songhua Liu, Zhenxiong Tan, Xinchao Wang

cs.AI

Samenvatting

Diffusion Transformers (DiT) zijn een toonaangevende architectuur geworden in beeldgeneratie. Echter, de kwadratische complexiteit van aandachtsmechanismen, die verantwoordelijk zijn voor het modelleren van token-gerelateerde relaties, leidt tot aanzienlijke vertraging bij het genereren van afbeeldingen met een hoge resolutie. Om dit probleem aan te pakken, streven we in dit artikel naar een lineair aandachtsmechanisme dat de complexiteit van vooraf getrainde DiTs verlaagt naar lineair niveau. We beginnen onze verkenning met een uitgebreid overzicht van bestaande efficiënte aandachtsmechanismen en identificeren vier cruciale factoren die essentieel zijn voor een succesvolle linearisatie van vooraf getrainde DiTs: localiteit, formulering consistentie, hoge-rang aandachtskaarten en kenmerkintegriteit. Op basis van deze inzichten introduceren we een convolutie-achtige lokale aandachtsstrategie genaamd CLEAR, die de kenmerkinteracties beperkt tot een lokaal venster rond elk query-token en daardoor lineaire complexiteit bereikt. Onze experimenten tonen aan dat door de aandachtslaag te fine-tunen op slechts 10K zelf gegenereerde voorbeelden gedurende 10K iteraties, we effectief kennis kunnen overdragen van een vooraf getrainde DiT naar een studentmodel met lineaire complexiteit, resulterend in vergelijkbare resultaten als het docentmodel. Tegelijkertijd vermindert het de aandachtsberekeningen met 99.5% en versnelt het de generatie met 6.3 keer voor het genereren van 8K-resolutie afbeeldingen. Bovendien onderzoeken we gunstige eigenschappen in de gedistilleerde aandachtslagen, zoals zero-shot generalisatie over verschillende modellen en plug-ins, en verbeterde ondersteuning voor multi-GPU parallelle inferentie. Modellen en codes zijn beschikbaar op: https://github.com/Huage001/CLEAR.

English

Diffusion Transformers (DiT) have become a leading architecture in image generation. However, the quadratic complexity of attention mechanisms, which are responsible for modeling token-wise relationships, results in significant latency when generating high-resolution images. To address this issue, we aim at a linear attention mechanism in this paper that reduces the complexity of pre-trained DiTs to linear. We begin our exploration with a comprehensive summary of existing efficient attention mechanisms and identify four key factors crucial for successful linearization of pre-trained DiTs: locality, formulation consistency, high-rank attention maps, and feature integrity. Based on these insights, we introduce a convolution-like local attention strategy termed CLEAR, which limits feature interactions to a local window around each query token, and thus achieves linear complexity. Our experiments indicate that, by fine-tuning the attention layer on merely 10K self-generated samples for 10K iterations, we can effectively transfer knowledge from a pre-trained DiT to a student model with linear complexity, yielding results comparable to the teacher model. Simultaneously, it reduces attention computations by 99.5% and accelerates generation by 6.3 times for generating 8K-resolution images. Furthermore, we investigate favorable properties in the distilled attention layers, such as zero-shot generalization cross various models and plugins, and improved support for multi-GPU parallel inference. Models and codes are available here: https://github.com/Huage001/CLEAR.