FLARE: Fast Low-rank Attention Routing Engine

Основная информация

• ID статьи: 2508.12594
• Название: FLARE: Fast Low-rank Attention Routing Engine
• Авторы: Vedant Puri, Aditya Joglekar, Kevin Ferguson, Yu-hsuan Chen, Yongjie Jessica Zhang, Levent Burak Kara (Carnegie Mellon University)
• Классификация: cs.LG (Машинное обучение)
• Дата публикации: 15 октября 2025 г. (arXiv v2)
• Ссылка на статью: https://arxiv.org/abs/2508.12594

Аннотация

Квадратичная сложность традиционного механизма самовнимания ограничивает его применимость и масштабируемость на больших неструктурированных сетках. В данной работе предлагается быстрый механизм маршрутизации внимания низкого ранга (FLARE) — механизм самовнимания с линейной сложностью, который маршрутизирует внимание через последовательность скрытых переменных фиксированной длины. Каждая голова внимания проецирует входную последовательность на последовательность скрытых переменных длины M≪N, используя обучаемые токены запроса, обеспечивая глобальную коммуникацию между N токенами. Благодаря маршрутизации внимания через узкое место последовательности FLARE изучает внимание низкого ранга, которое может быть применено с затратами O(NM). FLARE не только масштабируется до беспрецедентных размеров задач, но и обеспечивает лучшую точность по сравнению с современными моделями нейронных агентов для УЧП на нескольких эталонах.

Исследовательский контекст и мотивация

Предпосылки проблемы

1. Основная проблема: Механизм самовнимания традиционного Transformer имеет временную и пространственную сложность O(N²), что серьёзно ограничивает его применение на больших неструктурированных сетках (таких как облака точек и сетки в физическом моделировании).
2. Важность приложений: При моделировании агентов уравнений в частных производных (УЧП) каждая точка в облаке точек 3D рассматривается как токен, содержащий геометрические и физические величины (такие как координаты, нормальные векторы, свойства материалов). Моделирование высокоточных физических систем требует больших вычислительных затрат, поэтому модели машинного обучения предоставляют быструю альтернативу приближения.
3. Ограничения существующих методов:
   • PerceiverIO: Выполняет только одно кодирование и декодирование, узкое место может ограничить точность
   • Transolver: Совместное использование весов проекции между головами не позволяет использовать существующие GPU ядра для масштабирования скалярного произведения внимания
   • LNO: Применяет только одну проекцию, не хватает возможностей глубокой модели
4. Исследовательская мотивация: Разработать механизм внимания, который сохраняет способность глобальной коммуникации, но имеет линейную сложность, позволяя Transformer обрабатывать геометрии с миллионами точек.

Основные вклады

1. Смешивание токенов с линейной сложностью: Предложен механизм самовнимания FLARE, который достигает линейной сложности путём замены полного самовнимания на проекцию низкого ранга и реконструкцию.
2. Превосходная точность: На нескольких эталонах УЧП FLARE достигает точности предсказания, превосходящей ведущие модели нейронных агентов, с меньшим количеством параметров и более низкой вычислительной сложностью.
3. Беспрецедентная масштабируемость: FLARE полностью построен на стандартных примитивах объединённого внимания, обеспечивая высокое использование GPU и поддерживая сквозное обучение неструктурированных сеток с миллионами точек.
4. Новый эталонный набор данных: Выпущен крупномасштабный высокоразрешающий набор данных аддитивного производства металла для исследования предсказания остаточного смещения.

Подробное описание метода

Определение задачи

Дана входная последовательность X ∈ R^(N×C), где N — количество токенов, C — размерность признаков. FLARE направлен на изучение механизма внимания с линейной сложностью, обеспечивающего эффективную глобальную коммуникацию между токенами.

Архитектура модели

Основной механизм FLARE

FLARE вводит M≪N обучаемых скрытых токенов в качестве узкого места для обмена информацией, состоящее из двух этапов:

1. Этап кодирования: Входная последовательность проецируется на скрытые токены через перекрёстное внимание

   Z_h = SDPA(Q_h, K_h, V_h, s=1)
   

   
   где Q_h ∈ R^(M×D) — матрица обучаемых запросов, K_h, V_h ∈ R^(N×D)
2. Этап декодирования: Скрытые токены проецируются обратно на входную последовательность

   Y_h = SDPA(K_h, Q_h, Z_h, s=1)
   

Матрица коммуникации низкого ранга

Весь процесс эквивалентен:

Y_h = (W_decode,h · W_encode,h) · V_h

где:

• W_encode,h = softmax(Q_h · K_h^T) ∈ R^(M×N)
• W_decode,h = softmax(K_h · Q_h^T) ∈ R^(N×M)
• W_h = W_decode,h · W_encode,h ∈ R^(N×N) — матрица глобальной коммуникации ранга не более M

Структура блока FLARE

X = X + FLARE(LayerNorm(X))
X = X + ResMLP(LayerNorm(X))

Технические инновации

1. Независимая проекция между головами: В отличие от Transolver, который совместно использует веса проекции, FLARE назначает каждой голове различные срезы скрытых токенов, позволяя каждой голове изучать независимые отношения внимания.
2. Глубокая остаточная MLP: Использование глубокой остаточной сети для проекции ключей/значений позволяет изучать взаимодействия признаков более высокого порядка по сравнению с простыми линейными слоями.
3. Симметричный дизайн кодирования-декодирования: Симметричность операций кодирования и декодирования способствует стабильному потоку информации.
4. Совместимость с объединёнными ядрами: Полностью построен на стандартных операциях SDPA, может использовать оптимизированные алгоритмы, такие как Flash Attention.

Экспериментальная установка

Наборы данных

В статье оценены 6 эталонных наборов данных и 1 новый предложенный набор:

Метрики оценки

Основной используемой метрикой является относительная ошибка L2:

Relative L2 = ||û - u||₂ / ||u||₂

Методы сравнения

• Модели универсального внимания: Vanilla Transformer, PerceiverIO
• Агенты УЧП на основе внимания: Transolver, LNO
• Нейронные операторы: GNOT

Детали реализации

• Оптимизатор: AdamW (β₁=0.9, β₂=0.999)
• Расписание скорости обучения: OneCycleLR, пиковая скорость обучения 10⁻³
• Количество эпох: 500 для 2D задач, 250 для LPBF
• Размер пакета: 2 для 2D задач, 1 для 3D задач

Результаты экспериментов

Основные результаты

FLARE достигает оптимальных или близких к оптимальным результатов на всех эталонах:

Примечание: значения — относительная ошибка L2 (×10⁻³)

Эксперименты с геометриями с миллионом точек

FLARE успешно обучается на наборе данных DrivAerML с миллионом точек на одном GPU H100, что является первым методом на основе внимания для нейронных агентов, обрабатывающим миллион точек без выгрузки памяти или распределённых вычислений.

Абляционные исследования

1. Влияние количества блоков (B) и скрытых токенов (M):
   • Увеличение количества блоков постоянно снижает относительную ошибку
   • Увеличение M обычно улучшает производительность, но тенденция не строго монотонна
   • Различные задачи требуют разных рангов
2. Временная и пространственная сложность:
   • FLARE более чем в 200 раз быстрее, чем vanilla attention
   • Использование памяти немного выше, чем vanilla attention, но значительно ниже, чем Physics Attention

Спектральный анализ

Анализ изученных матриц коммуникации с использованием алгоритма собственного разложения с временной сложностью O(M³+M²N):

• На ранних блоках собственные значения быстро убывают, что указывает на эффективное сжатие
• Глубокие блоки используют большую часть скрытой ёмкости
• Различные головы имеют различные спектральные профили, что подтверждает дизайн независимой проекции между головами

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

Нейронные агенты УЧП

• Нейронные операторы: FNO, DeepONet и другие изучают отображения между бесконечномерными функциональными пространствами
• Графовые сети: Используют локальные взаимодействия соседей на сетке
• Архитектуры Transformer: Позволяют агрегировать глобальный контекст, но ограничены квадратичной сложностью

Эффективные механизмы внимания

• Linformer: Проецирует последовательности ключей/значений через обучаемые линейные отображения
• Reformer: Использует локально-чувствительное хеширование
• Nyströmformer: Использует метод Nyström для приближения самовнимания
• LoRA: Низкоранговая адаптация, в основном используется для эффективной настройки

Заключение и обсуждение

Основные выводы

1. FLARE успешно преодолевает квадратичное узкое место сложности самовнимания через механизм внимания низкого ранга
2. Достигает современной точности на нескольких эталонах УЧП с меньшим количеством параметров и более низкой вычислительной сложностью
3. Впервые реализует обучение модели нейронного агента на основе внимания на геометриях с миллионом точек

Ограничения

1. Зависимость от глубокой остаточной MLP: Может вводить последовательные узкие места и увеличивать задержку
2. Ограничение фиксированных скрытых токенов: Выбор M требует настройки для конкретной задачи
3. Применимость к некоторым задачам высокого ранга: Например, в задаче Darcy vanilla transformer всё ещё имеет преимущества

Направления будущих исследований

1. Постепенное увеличение количества скрытых токенов во время обучения
2. Разработка условных по времени скрытых токенов для моделирования диффузии
3. Разработка вариантов только декодера для авторегрессивного моделирования
4. Решение проблемы последовательного узкого места глубокой остаточной MLP

Глубокая оценка

Преимущества

1. Сильная техническая инновативность:
   • Умело преобразует проблему маршрутизации внимания в разложение матрицы низкого ранга
   • Дизайн независимой проекции между головами позволяет специализированные паттерны маршрутизации
   • Полностью совместим с существующими GPU ядрами
2. Полнота экспериментов:
   • Охватывает 6 различных эталонов УЧП
   • Подробные абляционные исследования и спектральный анализ
   • Впервые реализует эксперименты в масштабе миллиона точек
3. Глубокий теоретический анализ:
   • Предоставляет алгоритм собственного разложения с временной сложностью O(M³+M²N)
   • Объясняет эффективность коммуникации низкого ранга с математической точки зрения
   • Проверяет предположения дизайна через спектральный анализ
4. Высокая практическая ценность:
   • Выпущен новый набор данных аддитивного производства
   • Код открыт для воспроизведения
   • Может быть напрямую интегрирован в существующие архитектуры Transformer

Недостатки

1. Ограничения применимости метода:
   • Ограниченная эффективность на задачах высокого ранга (например, Darcy)
   • Выбор M требует настройки для конкретной задачи
   • Глубокая MLP может стать новым вычислительным узким местом
2. Ограничения экспериментальной установки:
   • Отсутствие сравнения с большим количеством более новых методов
   • Относительно небольшой масштаб некоторых эталонов
   • Необходима дополнительная проверка универсальности для различных типов задач УЧП
3. Недостаточный теоретический анализ:
   • Отсутствует анализ сходимости
   • Ограниченное теоретическое руководство по выбору оптимального M
   • Необходимо дальнейшее обоснование обоснованности предположения низкого ранга для всех задач УЧП

Влияние

1. Академический вклад: Предоставляет новую парадигму проектирования для эффективных механизмов внимания, особенно в области научных вычислений
2. Практическая ценность: Позволяет Transformer обрабатывать крупномасштабные геометрические задачи, способствуя развитию AI4Science
3. Воспроизводимость: Открытый код, подробные параметры экспериментов, удобство для последующих исследований

Применимые сценарии

• Решение УЧП на больших неструктурированных сетках
• Обработка облаков точек и геометрическое глубокое обучение
• Задачи последовательного моделирования, требующие глобальной коммуникации с ограниченными вычислительными ресурсами
• Приложения моделирования агентов в научных вычислениях

Библиография

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

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