Token Fusion: Объединение отсечения и слияния токенов для эффективных Vision Transformers

1. Введение и обзор

Vision Transformers (ViTs) произвели революцию в компьютерном зрении, но страдают от высоких вычислительных затрат из-за квадратичной сложности механизма self-attention относительно количества входных токенов. В данной статье, Token Fusion: Bridging the Gap between Token Pruning and Token Merging, представлен метод Token Fusion (ToFu) — гибридный подход, который динамически выбирает между отсечением и слиянием токенов на основе поведения модели для оптимизации компромисса между эффективностью и точностью.

Ключевая идея заключается в том, что ни отсечение (удаление токенов), ни слияние (усреднение токенов) не являются универсально оптимальными. В статье предлагается принципиальный способ выбора подходящей операции для каждого слоя, а также новая техника слияния под названием MLERP (Multi-token Linear intERPolation), предназначенная для решения проблем сдвига распределения при стандартном усреднении.

2. Основная методология: Token Fusion (ToFu)

ToFu основан на анализе реакции модели на интерполированные входные данные, определяя её пригодность для слияния или отсечения.

2.1. Дилемма: отсечение против слияния

Авторы определяют ключевой критерий: линейность модели. Если слой модели реагирует почти линейно на интерполированные входные данные (например, $f(\alpha x_1 + (1-\alpha)x_2) \approx \alpha f(x_1) + (1-\alpha)f(x_2)$), то слияние похожих токенов путём усреднения эффективно и сохраняет информацию. Однако в ранних/глубоких нелинейных слоях (как визуализировано на их Рисунке 1) линейная интерполяция во входном пространстве приводит к сильно нелинейным выходам, что делает усреднённое слияние проблематичным и потенциально вызывает сдвиги распределения. В таких случаях отсечение менее важных токенов является более безопасной, хотя и более «потерянной» альтернативой.

2.2. Фреймворк ToFu

ToFu работает на уровне каждого трансформер-блока:

Оценка важности токенов: Назначает оценку важности каждому токену (например, на основе нормы внимания или градиента).
Оценка линейности: Оценивает приблизительную линейность слоя, часто определяемую эмпирически или с помощью лёгкого зонда.
Адаптивная операция: Для целевого коэффициента сокращения токенов:
- В областях с высокой линейностью: Объединить наименее важные токены с их наиболее похожими важными соседями.
- В областях с низкой линейностью: Непосредственно отсечь наименее важные токены.

Это создаёт динамичный, контекстно-зависимый конвейер сжатия.

2.3. MLERP: Слияние с сохранением нормы

Для улучшения простого усреднения авторы предлагают MLERP, адаптацию Сферической Линейной Интерполяции (SLERP) для слияния $K$ токенов. Для токенов $t_1, t_2, ..., t_K$ с нормами $n_i = ||t_i||$, MLERP сначала интерполирует направления на единичной сфере, а затем масштабирует взвешенным средним исходных норм:

$t_{\text{merged}} = \left( \frac{\sum_{i=1}^K w_i n_i}{\| \sum_{i=1}^K w_i \frac{t_i}{n_i} \|} \right) \left( \sum_{i=1}^K w_i \frac{t_i}{n_i} \right)$

где $w_i$ — веса, основанные на важности. Это сохраняет статистическое распределение норм признаков, смягчая сдвиг распределения, вызванный наивным усреднением, и приводит к более стабильной производительности, особенно в нелинейных режимах.

3. Технические детали и математическая формулировка

В статье формализуется проблема сокращения токенов. Пусть слой имеет $N$ входных токенов $T = \{t_1, ..., t_N\}$. Цель — создать сокращённое множество $T'$ с $M < N$ токенами.

Ключевые уравнения:

Оценка важности: $I(t_i) = ||\text{Attn}(t_i)||_1$ или мера на основе градиента.
Метрика схожести: Обычно косинусное сходство $S(t_i, t_j) = \frac{t_i \cdot t_j}{||t_i|| \, ||t_j||}$.
Метрика линейности ($\mathcal{L}$): Измеряется отклонением выходов слоя от линейной интерполяции входов. Низкое значение $\mathcal{L}$ благоприятствует слиянию; высокое — отсечению.

Алгоритм ToFu может быть применён к предварительно обученным моделям без дообучения (zero-shot) или улучшен с помощью лёгкого обучения.

4. Экспериментальные результаты и производительность

Авторы оценивают ToFu на задачах классификации изображений (ImageNet с ViT-B/16, DeiT) и генерации изображений (латентные диффузионные модели).

Ключевые показатели производительности

Классификация: ToFu достигает лучшего компромисса между точностью и FLOPs, чем отдельные методы отсечения (например, DynamicViT) или слияния (ToMe). Например, при сокращении FLOPs на 40% ToFu теряет <0.5% точности top-1 на ImageNet, превосходя ToMe примерно на ~0.3%.
Генерация изображений: В Stable Diffusion ToFu сохраняет более высокую визуальную точность (измеряемую по FID) при сниженных вычислительных затратах по сравнению с ToMe, особенно при сокращении большого количества токенов. Слияние MLERP показывает более явное преимущество в задачах генерации, где критически важно распределение выходных данных.
Абляционное исследование: Показано, что адаптивная стратегия (выбор слияния/отсечения) превосходит использование только одной операции во всех слоях. MLERP последовательно превосходит усреднённое слияние.

Описание графика (на основе Рисунка 1 статьи): На рисунке иллюстрируется нелинейность слоёв ViT. Две точки входных признаков (x1, x2) линейно интерполируются (цветная линия). Выходы (f1-f4) из четырёх различных MLP-слоев внутри ViT отображены на графике. Выходы ранних и поздних MLP-слоев (f1, f4) показывают значительное отклонение от прямой линии, что указывает на сильную нелинейность. Среднее двух входов (фиолетовая звезда) отображается в выходную точку, далёкую от среднего выходов, наглядно демонстрируя, почему усреднённое слияние может давать сбой в нелинейных слоях.

5. Фреймворк анализа и пример использования

Пример: Применение ToFu к предварительно обученному ViT для развёртывания на периферийных устройствах

Сценарий: Разработчику необходимо запустить модель ViT-B на мобильном устройстве для классификации изображений в реальном времени. Полная модель слишком медленная.

Применение фреймворка:

Профилирование: Прогнать небольшой калибровочный набор данных через модель. Для каждого трансформер-блока вычислить метрику линейности $\mathcal{L}$, выбирая пары токенов и проверяя ошибку интерполяции выходов.
Карта стратегии: Создать профиль: Блоки 1-3 (низкая линейность) → предпочтительно отсечение. Блоки 4-8 (высокая линейность) → предпочтительно слияние MLERP. Финальные блоки (низкая линейность) → предпочтительно отсечение.
Конфигурация: Установить глобальный бюджет сокращения токенов (например, 35%). Применить отсечение в блоках с низкой линейностью и слияние MLERP в блоках с высокой линейностью, соблюдая бюджеты для каждого блока, выведенные из оценок важности.
Оценка: Развернуть сжатую модель. Адаптивный подход обеспечивает минимальное падение точности по сравнению с универсальным методом, так как избегает агрессивного слияния в чувствительных нелинейных слоях.

Этот пример демонстрирует практическую полезность ToFu как структурированного фреймворка сжатия, а не просто монолитного алгоритма.

6. Будущие применения и направления исследований

Мультимодальные трансформеры: Расширение ToFu на видео, аудио или мультимодальные трансформеры (например, CLIP, Flamingo), где динамика токенов более сложна.
Аппаратно-зависимое совместное проектирование: Оптимизация алгоритма принятия решений ToFu (отсечение/слияние) и реализации MLERP для конкретных AI-ускорителей (NPU, GPU) для максимизации реального ускорения.
Интеграция с другими техниками: Комбинирование ToFu с квантованием, дистилляцией знаний или эффективными механизмами внимания (такими как Linformer) для совокупного повышения эффективности.
Автоматический поиск гиперпараметров: Использование нейроархитектурного поиска (NAS) или обучения с подкреплением для автоматического определения оптимального соотношения отсечения/слияния для каждого слоя и порога линейности.
За пределами зрения: Исследование эффективности метода в больших языковых моделях (LLM) для сжатия последовательностей, хотя семантика токенов там значительно отличается.

7. Ссылки

Dosovitskiy, A., et al. "An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale." ICLR 2021.
Bolya, D., et al. "Token Merging: Your ViT But Faster." ICLR 2023 (ToMe).
Wang, Y., et al. "DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification." NeurIPS 2021.
Rombach, R., et al. "High-Resolution Image Synthesis with Latent Diffusion Models." CVPR 2022.
Zhu, J.Y., et al. "Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks." ICCV 2017 (CycleGAN).
Vaswani, A., et al. "Attention Is All You Need." NeurIPS 2017.

8. Экспертный анализ и критические замечания

Ключевая идея: ToFu — это не просто ещё один инструмент сжатия; это формальное признание того, что слои трансформера гетерогенны. Обработка всех слоев одним и тем же примитивом сжатия наивна. Блеск статьи заключается в её диагностическом подходе — измерении линейности слоя для назначения правильной «операции» (отсечение или слияние). Это напоминает то, как современные компиляторы профилируют код для применения оптимизаций — уровень изощрённости, часто отсутствующий в исследованиях эффективности ML.

Логическая последовательность: Аргументация убедительна: 1) Показать, что усреднённое слияние даёт сбой в нелинейных слоях (Рис. 1). 2) Предложить метрику для обнаружения этого режима сбоя (линейность). 3) Использовать метрику для маршрутизации токенов. 4) Исправить проблемную операцию (усреднённое слияние) с помощью MLERP. Последовательность от идентификации проблемы до многокомпонентного решения чиста и логична.

Сильные стороны и недостатки:
Сильные стороны: Гибридный подход теоретически обоснован и эмпирически проверен на различных задачах. MLERP — это простое, но умное решение реальной проблемы (коллапс нормы). Возможность применения без дообучения (zero-shot) является серьёзным практическим преимуществом для развёртывания существующих моделей.
Недостатки: В статье несколько недооценивается накладные расходы на «оценку линейности». Это предварительно вычисленный профиль (статический) или вычисление на лету (динамические накладные расходы)? Преимущества MLERP, хотя и очевидны, кажутся скромными в классификации; его истинная ценность, по-видимому, более выражена в генеративных задачах, что согласуется с выводами из литературы по диффузионным моделям, где распределение выходных данных имеет первостепенное значение. Сравнение, хотя и справедливое, могло бы быть более агрессивным по отношению к современным методам посттренировочного квантования, которые предлагают ортогональные преимущества.

Практические выводы: Для практиков: Немедленно принять ToFu/MLERP в качестве метода сокращения токенов первой линии для ViT, особенно для генеративных задач. Он заменяет ToMe в качестве стратегии слияния по умолчанию. Для исследователей: Ключевой вывод — парадигма «сжатия с учётом слоя». Будущая работа должна быть сосредоточена на автоматизации обнаружения областей модели, дружественных к сжатию и чувствительных к нему, возможно, черпая вдохновение из работ по отсечению сетей в CNN или анализу коллапса мод в GAN, таких как CycleGAN. Следующий рубеж — создание моделей, изначально эффективных по дизайну, с использованием идей из таких диагностических исследований для информирования поиска архитектур, выходя за рамки простого постфактумного сжатия.