Очистка: Линеаризация, аналогичная свертке, для предварительно обученных трансформеров диффузии Вверх

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

December 20, 2024

Авторы: Songhua Liu, Zhenxiong Tan, Xinchao Wang

cs.AI

Аннотация

Трансформеры диффузии (DiT) стали ведущей архитектурой в области генерации изображений. Однако квадратичная сложность механизмов внимания, отвечающих за моделирование отношений между токенами, приводит к значительной задержке при создании изображений высокого разрешения. Для решения этой проблемы в данной статье мы стремимся к механизму линейного внимания, который снижает сложность предварительно обученных DiT до линейной. Мы начинаем наше исследование с подробного обзора существующих эффективных механизмов внимания и выделяем четыре ключевых фактора, необходимых для успешной линеаризации предварительно обученных DiT: локальность, согласованность формулировок, высокоранговые карты внимания и целостность признаков. Основываясь на этих идеях, мы представляем стратегию локального внимания, аналогичную свертке, названную CLEAR, которая ограничивает взаимодействия признаков локальным окном вокруг каждого токена запроса и, таким образом, достигает линейной сложности. Наши эксперименты показывают, что, донастраивая слой внимания всего лишь на 10 тыс. самопорожденных образцов в течение 10 тыс. итераций, мы можем эффективно передавать знания от предварительно обученного DiT к модели ученика с линейной сложностью, давая результаты, сравнимые с моделью учителя. Одновременно это снижает вычисления внимания на 99,5% и ускоряет генерацию в 6,3 раза для создания изображений разрешением 8K. Кроме того, мы исследуем благоприятные свойства в дистиллированных слоях внимания, такие как обобщение без обучения на различных моделях и плагинах, а также улучшенная поддержка параллельного вывода на нескольких GPU. Модели и коды доступны здесь: 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.