FNET: СМЕШИВАНИЕ ТОКЕНОВ С ПОМОЩЬЮ ПРЕОБРАЗОВАНИЙ ФУРЬЕ

1. Введение

Архитектуры Transformer стали базовым инструментом для обработки естественного языка и ряда смежных задач, обеспечив значительный прогресс в моделировании последовательностей. Однако практическое применение трансформеров ограничивается вычислительной стоимостью само-внимания: для последовательности длины \(n\) вычислительная сложность классического внимания масштабируется как \(O(n^2)\), что делает обучение и вывод дорогостоящими при больших \(n\).

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

Авторы показывают, что такая замена может дать существенное ускорение при умеренной (а иногда и небольшой) потере качества. В частности, предлагается модель FNet, где вместо внимания используется дискретное преобразование Фурье (и на практике — FFT), выполняющее глобальное смешивание по оси токенов и по оси скрытых признаков.

Основные вклады работы:

• предложена простая архитектура FNet, где слой внимания заменён на преобразование Фурье как механизм глобального смешивания токенов;
• проведена оценка на задачах переноса (GLUE) и на Long-Range Arena (LRA), демонстрирующая практический компромисс «точность–скорость»;
• выполнен детальный анализ скорости и потребления памяти на разных аппаратных платформах (GPU/TPU) и при разных длинах последовательности;
• рассмотрены гибридные варианты (часть слоёв внимания + слои Фурье), а также альтернативные линейные преобразования.

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

2.1. Преобразования Фурье в нейронных сетях

Преобразования Фурье и близкие ортогональные преобразования (например, дискретное косинусное преобразование) применялись в машинном обучении в разных контекстах — от обработки сигналов и изображений до построения признаковых представлений. В современных нейросетях Фурье-операторы встречаются как инструменты для эффективного глобального смешивания, а также как компоненты архитектур, которые заменяют или дополняют стандартные свёртки/внимание.

Важное преимущество FFT состоит в более благоприятной асимптотике по сравнению с прямым умножением на матрицу DFT: \(O(n \log n)\) вместо \(O(n^2)\). Это делает Фурье-подходы особенно интересными для длинных последовательностей.

2.2. Моделирование семантических отношений с помощью внимания

Само-внимание стало ключевым механизмом трансформера (Vaswani et al., 2017), позволяя модели строить контекстно-зависимые представления и обучать богатые зависимости между токенами. Тем не менее, матрица внимания вычисляется для каждой пары токенов, что и создаёт квадратичную стоимость по длине последовательности.

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

2.3. Эффективные модели для длинных последовательностей

Существует большая линия работ по «эффективным трансформерам», ориентированная на обработку длинных последовательностей: Longformer, Reformer, Performer, Linformer, BigBird и др. Эти модели сокращают асимптотику внимания либо за счёт разреженных паттернов внимания, либо за счёт низкоранговых/ядерных аппроксимаций.

На их фоне FNet интересен тем, что полностью исключает вычисление внимания и заменяет его крайне простой операцией, хорошо оптимизируемой на современных ускорителях.

3. Модель

3.1. Дискретное преобразование Фурье

Для сигнала (или вектора) длины \(N\) дискретное преобразование Фурье (DFT) определяется как:

DFT можно записать в матричной форме через матрицу Фурье \(W \in \mathbb{C}^{N \times N}\), где элементы имеют вид:

Тогда преобразование можно представить как умножение матриц. Прямое вычисление DFT требует \(O(N^2)\), но быстрый алгоритм FFT позволяет вычислять DFT за \(O(N \log N)\).

3.2. Архитектура FNet

FNet строится по мотивам BERT-подобного энкодера: входная последовательность токенов преобразуется в эмбеддинги (включая позиционные), далее проходит через стек одинаковых блоков. Отличие состоит в том, что подслой внимания заменён на слой Фурье-смешивания (Fourier mixing).

Пусть вход в блок имеет вид \(x \in \mathbb{R}^{n \times d_h}\), где \(n\) — длина последовательности, \(d_h\) — размерность скрытого пространства. В FNet смешивание определяется как применение двумерного преобразования Фурье (по оси токенов и по оси скрытых признаков) с последующим взятием действительной части:

Здесь \(\mathcal{F}_{\text{seq}}\) обозначает DFT/FFT по оси последовательности, а \(\mathcal{F}_{h}\) — DFT/FFT по оси скрытых признаков. В реализации применяются FFT-операторы.

Как и в стандартном Transformer/BERT-блоке, используются остаточные связи и слой нормализации, а также отдельный полносвязный подслой FFN (feed-forward network). В результате блок сохраняет общую структуру трансформера, но имеет значительно более дешёвый «mixing» подслой.

3.3. Реализация

На практике для вычисления преобразования Фурье используется FFT, что даёт выигрыш по времени на больших \(n\). Отдельно подчёркивается, что результат FFT комплекснозначный, но в FNet используется только действительная часть.

Важно также, что слой смешивания в базовой версии FNet не содержит обучаемых параметров (в отличие от внимания). Это влияет на размер модели и на стабильность обучения.

4. Результаты

4.1. Обучение с переносом (Transfer learning)

Авторы оценивают FNet в стандартной парадигме, аналогичной BERT: сначала предобучение (masked language modeling и next sentence prediction), затем дообучение на задачах GLUE. Рассматриваются модели размеров «Base» и «Large» и несколько базовых линий сравнения: (1) BERT как эталон, (2) Linear-смешивание, (3) Random-смешивание, (4) FF-only (без смешивания), а также (5) гибридные варианты FNet-Hybrid.

В целом FNet демонстрирует заметно более низкое качество по сравнению с BERT при чистом замещении внимания, но при этом выигрывает по скорости. Гибридные модели (FNet-Hybrid) частично компенсируют просадку качества, сохраняя при этом значительную долю ускорения.

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

4.2. Long-Range Arena (LRA)

Чтобы оценить пригодность FNet для длинных последовательностей, авторы используют набор задач Long-Range Arena (Tay et al., 2020), который включает различные типы входов (текст, последовательности токенов «ListOps», задачи поиска/извлечения, изображения и др.) и требует моделировать зависимости на больших расстояниях.

В таблице ниже приведено сравнение FNet с «vanilla Transformer» и рядом эффективных трансформеров. По точности FNet обычно сопоставим с базовым трансформером (а иногда и превосходит отдельные варианты), при этом демонстрирует лучшие масштабируемость и/или скорость при больших длинах.

Примечание: в оригинале авторы отмечают, что по задаче Path-X многие модели либо не обучаются, либо показывают FAIL, а в ряде конфигураций возникают ограничения по памяти (OOM). Также «Среднее» не включает Path-X.

5. Выводы

В статье предложена модель FNet, которая заменяет подслой само-внимания в Transformer-энкодере на слой смешивания токенов на основе преобразования Фурье. Такой подход радикально упрощает вычисления и снижает стоимость по времени и памяти, особенно заметно при длинных последовательностях.

Эксперименты на GLUE показывают, что FNet может достигать конкурентного качества при существенно меньших вычислительных затратах, а гибридные варианты позволяют приблизить качество к BERT, сохраняя ускорение. На LRA FNet демонстрирует сильные свойства масштабируемости: при росте длины последовательности выигрыши по скорости и памяти возрастают.

В качестве дальнейших направлений авторы рассматривают расширение гибридных схем, поиск более «информативных» безпараметрических mixing-преобразований и применение к другим модальностям.

A. Приложения

A.1. Подробности предобучения

Авторы используют те же «фиксированные» конфигурации (Base и Large), что и в оригинальном BERT: архитектурные параметры и базовые настройки оптимизации берутся из Devlin et al. (2019). Важно учитывать, что такая конфигурация исторически «настроена» под механизм внимания, поэтому результаты альтернативных mixing-слоёв могут быть частично смещены в пользу BERT.

В таблице 5 приведены метрики предобучения: значения функции потерь (Total/MLM/NSP) и точности по задачам MLM и NSP (на TPU). Несмотря на более слабые метрики точности, Linear и FNet обучаются существенно быстрее, чем BERT. Авторы также отмечают, что модели, у которых слой смешивания не содержит обучаемых параметров (FNet, Random, FF-only), демонстрируют более стабильную динамику обучения.

В таблице 6 приведены размеры моделей, использованные авторами для построения графиков (рисунки 2 и 3): скрытая размерность \(d_h\), число слоёв и количество параметров (в миллионах) для BERT, Linear, FNet и FNet-Hybrid (без учёта выходных проекционных слоёв).

A.2. Результаты на TPU

Авторы отдельно анализируют характеристики FNet на TPU. В частности, отмечается, что для некоторых диапазонов длины последовательности реализация DFT через матричные операции (mat) может быть выгодной, тогда как FFT становится особенно полезной на больших длинах.

A.3. Альтернативные линейные преобразования

В качестве альтернативы DFT/FFT можно рассматривать другие ортогональные линейные преобразования (например, DCT, Hadamard, Hartley). Они также обеспечивают глобальное смешивание токенов, но по экспериментам авторов не дают заметного превосходства над Фурье-подходом в целевых настройках.

В частности, преобразование Хартли (Hartley transform) является вещественной альтернативой Фурье, что потенциально упрощает реализацию, однако практические результаты не показывают устойчивого выигрыша.

A.4. Профилирование скорости преобразования Фурье

Авторы отдельно профилируют скорость «mixing»-подслоя, чтобы понять, где именно возникает ускорение. В таблице 8 приведены времена выполнения только слоя смешивания (без FFN и прочих компонентов). Видно, что замена внимания на линейный слой и особенно на Фурье-смешивание резко уменьшает стоимость соответствующего подслоя.

A.5. Гибридные модели

Гибридные модели (FNet-Hybrid) комбинируют небольшое число слоёв внимания со слоями Фурье-смешивания. Идея состоит в том, что внимание может давать критически важные «семантические» взаимодействия, а слои Фурье обеспечивают дешёвое глобальное смешивание на остальных уровнях.

Авторы исследуют, как расположение и число attention-слоёв влияет на метрики предобучения и скорость.

A.6. Сравнение с «vanilla Transformer» в LRA

В приложении авторы обсуждают различия между их собственной реализацией базового трансформера и реализациями/результатами, представленными в исходной работе по LRA и в последующих публикациях. Показатели скорости и памяти могут заметно отличаться из-за аппаратной платформы, деталей оптимизации, реализации слоёв внимания и особенностей батчинга.

Также подчёркивается, что «эффективные трансформеры» часто оптимизируются под конкретные диапазоны длины, а прямое сравнение требует аккуратного согласования конфигураций.

A.7. Код базовой модели (как в оригинале)

Ниже приведён листинг «ядра» модели (упрощённая иллюстрация основных классов и прямого прохода), как он представлен в приложении статьи.

# Listing 1: Core model code (from the paper)
import flax.linen as nn
import jax.numpy as jnp
from jax.numpy.fft import fft, fft2

class FourierTransformLayer(nn.Module):
    def __call__(self, x):
        # Apply 2D FFT along the last two dimensions.
        return jnp.real(fft2(x))

class FeedForwardLayer(nn.Module):
    hidden_dim: int
    intermediate_dim: int
    dropout_rate: float

    def setup(self):
        self.dense1 = nn.Dense(self.intermediate_dim)
        self.dense2 = nn.Dense(self.hidden_dim)
        self.dropout = nn.Dropout(self.dropout_rate)

    def __call__(self, x, deterministic: bool):
        x = self.dense1(x)
        x = nn.gelu(x)
        x = self.dropout(x, deterministic=deterministic)
        x = self.dense2(x)
        x = self.dropout(x, deterministic=deterministic)
        return x

class FNetLayer(nn.Module):
    hidden_dim: int
    intermediate_dim: int
    dropout_rate: float

    def setup(self):
        self.fourier = FourierTransformLayer()
        self.ffn = FeedForwardLayer(self.hidden_dim, self.intermediate_dim, self.dropout_rate)
        self.layer_norm1 = nn.LayerNorm()
        self.layer_norm2 = nn.LayerNorm()
        self.dropout = nn.Dropout(self.dropout_rate)

    def __call__(self, x, deterministic: bool):
        y = self.fourier(x)
        x = x + self.dropout(y, deterministic=deterministic)
        x = self.layer_norm1(x)

        y = self.ffn(x, deterministic=deterministic)
        x = x + self.dropout(y, deterministic=deterministic)
        x = self.layer_norm2(x)
        return x

class FNetEncoder(nn.Module):
    num_layers: int
    hidden_dim: int
    intermediate_dim: int
    dropout_rate: float

    def setup(self):
        self.layers = [
            FNetLayer(self.hidden_dim, self.intermediate_dim, self.dropout_rate)
            for _ in range(self.num_layers)
        ]

    def __call__(self, x, deterministic: bool):
        for layer in self.layers:
            x = layer(x, deterministic=deterministic)
        return x

Примечание: листинг воспроизведён по содержанию приложения статьи; для полной актуальной реализации см. репозиторий авторов (ссылка указана в начале документа).