Reformer

Reformer — экономичный по памяти трансформер, представленный Никитой Китаевым, Лукашом Кайзером и Ансельмом Левскаей из UC Berkeley и Google Research в работе Reformer: эффективный трансформер на ICLR 2020. Он решает две проблемы, ограничивающие длину последовательности. Механизм внимания на основе хеширования с учетом локальности (LSH) заменяет полную матрицу n × n внимания и требует около L log L операций для последовательности длины L; обратимые остаточные слои сохраняют активации, необходимые для обратного распространения ошибки, один раз для всей сети, а не один раз на слой. Google сообщила, что Reformer обрабатывает контекстные окна размером до 1 миллиона слов на одном ускорителе с 16 ГБ памяти и обработала весь роман Преступление и наказание как один обучающий пример.

Внимание LSH#

После применения softmax внимание каждого запроса сосредоточено на нескольких наиболее похожих на него ключах, поэтому Reformer ищет только их. Он использует одну общую проекцию для запросов и ключей и хеширует каждый вектор с помощью углового хеширования с учетом локальности: умножает его на случайную матрицу R и выбирает в качестве бакета позицию наибольшего значения в [xR, −xR], поэтому векторы, направленные в похожие стороны, обычно попадают в один бакет. Затем токены сортируются по бакетам, отсортированная последовательность разбивается на равные блоки, и каждый токен обращается только к своему блоку и предыдущему. Несколько раундов хеширования выполняются параллельно, чтобы обнаружить похожие пары, которые один раунд разделил бы.

LSH — классический метод приближённого поиска ближайших соседей с помощью векторного поиска: он находит большинство k ближайших соседей запроса с гораздо меньшими затратами, чем точный поиск. Шаг распределения по бакетам достаточно короткий, чтобы выполнить его отдельно:

import torch

torch.manual_seed(0)
dim, n_buckets, rounds = 64, 16, 1000
query = torch.randn(dim)
vectors = {"query": query, "near": query + 0.3 * torch.randn(dim), "far": torch.randn(dim)}

# Angular LSH as in Reformer: random projection R, then bucket = argmax over [xR, -xR]
R = torch.randn(rounds, dim, n_buckets // 2)
buckets = {}
for name, x in vectors.items():
    proj = torch.einsum("d,rdb->rb", x, R)
    buckets[name] = torch.cat([proj, -proj], dim=-1).argmax(dim=-1)

for name in ("near", "far"):
    cosine = torch.cosine_similarity(query, vectors[name], dim=0)
    shared = (buckets[name] == buckets["query"]).float().mean()
    print(f"{name}: cosine similarity {cosine:.2f}, same bucket as the query in {shared:.0%} of rounds")

Близкий вектор с косинусным сходством 0,96 относительно запроса попадает в тот же бакет в 72% из 1000 раундов, а несвязанный вектор со сходством 0,04 — в 9%. Поэтому внимание, ограниченное одним бакетом, сохраняет большую часть важных пар.

Обратимые слои и разбиение на блоки#

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

Во время обратного прохода каждый слой восстанавливает входные данные по выходным: x2 = y2 − FeedForward(y1), затем x1 = y1 − Attention(x2), поэтому сохраняются только итоговые активации, а каждый слой пересчитывается один раз. Контрольные точки для градиентов используют похожий компромисс: сохраняют активации в нескольких контрольных точках и пересчитывают промежуточные слои. Reformer также выполняет полносвязные слои блоками позиций, поскольку их скрытый слой в несколько раз шире модели. Объём памяти для активаций ограничивает размер батча при обучении Ultralytics YOLO26: с помощью batch=-1 AutoBatch измеряет объём памяти GPU для нескольких размеров батча, строит линейную зависимость и выбирает размер батча, который, согласно прогнозу, займёт 60% свободной памяти.

Результаты и реализации#

Reformer не уступал стандартным трансформерам при обработке символьного текста из enwik8 и генерации изображений ImageNet-64 пиксель за пикселем, потребляя гораздо меньше памяти и быстрее работая с длинными последовательностями. В Hugging Face Transformers он реализован как ReformerModel; доступны опубликованные контрольные точки, обученные на enwik8 и Преступлении и наказании.

Reformer и родственные модели#

• Longformer: заранее задаёт шаблон внимания: скользящее окно вокруг каждого токена и несколько глобальных токенов.
• Разреженное внимание: семейство методов, которые оценивают только некоторые пары токенов. Внимание LSH — один из его методов, учитывающих содержимое.
• Линейное внимание: заменяет softmax отображением признаков ядра, чтобы вычислять внимание за линейное время.
• Flash attention: вычисляет точное внимание в тайлах встроенной памяти GPU, не сохраняя матрицу оценок, поэтому объём памяти растёт линейно, а количество оценок остаётся квадратичным.