TinyFusion: Ускорение диффузионных трансформеров с помощью обрезки слоев

В последние годы диффузионные трансформеры стали основой для генеративных задач, достигнув значительных успехов в таких областях, как генерация изображений и видео. Однако, несмотря на свои выдающиеся возможности, эти модели часто имеют чрезмерную параметризацию, что приводит к значительным затратам на вывод в реальных приложениях. В данной статье мы представим метод обрезки слоев, называемый TinyFusion, который разработан для уменьшения избыточности слоев в диффузионных трансформерах с использованием методов обучения.

Проблема избыточности

Диффузионные трансформеры, такие как DiT (Diffusion Transformer), MAR (Masked Autoregressive) и SiT (Scalable Interpolant Transformer), имеют сложные архитектуры, которые обеспечивают высокое качество генерации, но требуют значительных вычислительных ресурсов. Основные факторы, влияющие на эффективность диффузионных моделей, включают количество шагов выборки, дизайн операторов, вычислительную точность, ширину и глубину сети. В этом контексте обрезка слоев представляет собой мощный инструмент для уменьшения задержки и повышения производительности.

Метод обрезки слоев

Традиционный подход к обрезке слоев заключается в том, чтобы сначала удалить незначимые слои, а затем провести дообучение модели для восстановления производительности. Однако существующие методы в основном основываются на эвристических подходах, таких как оценка важности слоев, что часто не дает гарантии хорошей производительности после дообучения. Мы предлагаем новый подход, который вместо минимизации потерь сразу после обрезки, фокусируется на восстановлении производительности после дообучения.

Основная идея TinyFusion

TinyFusion представляет собой обучаемый метод, который объединяет процессы обрезки и дообучения. Мы моделируем обрезку и дообучение как дифференцируемый процесс выборки масок слоев, что позволяет одновременно оптимизировать параметры модели. Основная цель заключается в том, чтобы итеративно уточнять распределение, так чтобы сети с высокой восстановляемостью имели более высокую вероятность быть выбраны в процессе выборки.

Дифференцируемая выборка

Метод включает в себя дифференцируемую выборку, которая позволяет обновлять распределение масок слоев с использованием градиентного спуска. Мы предполагаем, что маска, полученная "идеальными" методами обрезки, должна следовать определенному распределению. Начальная оценка распределения масок начинается с равномерного распределения, однако прямой выбор из этого распределения неэффективен из-за огромного пространства поиска.

Для решения этой проблемы мы вводим продвинутый алгоритм, который использует результаты оценки для итеративного уточнения распределения масок. Если определенные маски показывают положительные результаты, другие маски с похожими паттернами также должны иметь более высокую вероятность выборки в будущих оценках.

Эффективность метода

Мы провели обширные эксперименты на различных моделях диффузионных трансформеров, включая DiT, MAR и SiT. Результаты показывают, что наш обучаемый подход позволяет эффективно выявлять избыточные слои в диффузионных трансформерах с использованием всего одной эпохи обучения на наборе данных. Например, модели, обрезанные с помощью TinyFusion, быстро восстанавливаются через дообучение, достигая значительно более конкурентного FID (Fréchet Inception Distance) по сравнению с методами, минимизирующими немедленные потери.

Влияние на производительность

Мы также исследовали роль дистилляции знаний для улучшения восстановляемости. Введенный вариант MaskedKD позволяет смягчить негативное влияние больших или выбросных активаций в скрытых состояниях, что может существенно повлиять на производительность и надежность дообучения. Использование MaskedKD улучшает FID от 5.73 до 3.73 при использовании всего 1% от стоимости предобучения.

Связанные работы

Существующие методы обрезки сети, такие как ширинная обрезка и разреженность, помогают снизить память, но часто предлагают ограниченные улучшения скорости, особенно на параллельных устройствах, таких как GPU. В последние годы обрезка глубины привлекла значительное внимание, так как удаление целых слоев обеспечивает лучшую скорость, пропорциональную коэффициенту обрезки.

Методы, такие как MoD и глубинно-осведомленные трансформеры, также были предложены. Тем не менее, большинство существующих методов по-прежнему основаны на эмпирических или эвристических стратегиях, которые часто не обеспечивают сильной гарантии производительности после дообучения.

Экспериментальные результаты

Мы провели эксперименты на различных архитектурах, включая DiT-XL, MAR и SiT, и результаты показывают, что TinyFusion может создать более мелкий диффузионный трансформер с менее чем 7% от стоимости предобучения, достигая 2x ускорения с FID 2.86, превосходя конкурентов с сопоставимой эффективностью.

Сравнение с другими методами

В наших экспериментах мы сравнили TinyFusion с другими методами, такими как Diff-Pruning и ShortGPT. Результаты показывают, что TinyFusion не только превосходит эти методы по метрикам производительности, но и делает это с гораздо меньшими затратами на обучение.

Заключение

Работа представляет TinyFusion, обучаемый метод для ускорения диффузионных трансформеров путем удаления избыточных слоев. Метод моделирует восстановляемость обрезанных моделей как оптимизируемую задачу и включает дифференцируемую выборку для обучения от начала до конца. TinyFusion демонстрирует сильную обобщаемость на различных архитектурах, таких как DiTs, MARs и SiTs, и предлагает значительные преимущества в производительности и эффективности по сравнению с существующими подходами.