Transformer-XL

Transformer-XL — это трансформерная языковая модель, которая сохраняет контекст между текстовыми сегментами фиксированной длины: она кэширует скрытые состояния предыдущего сегмента, и следующий сегмент может обращаться к ним. Зихан Дай, Чжилин Ян и их коллеги из Университета Карнеги — Меллона и Google Brain представили эту модель на конференции ACL 2019 в статье «Transformer-XL: языковые модели с вниманием за пределами контекста фиксированной длины». Обычный трансформер, обученный на сегментах фиксированной длины, каждый из которых является отдельным контекстным окном, не может увидеть текст до начала текущего сегмента. Авторы назвали эту проблему фрагментацией контекста. Transformer-XL усваивал зависимости, которые были на 80% длиннее, чем у рекуррентных нейронных сетей (RNN), и на 450% длиннее, чем у обычных трансформеров.

Рекуррентность на уровне сегментов#

При обработке сегмента каждый слой объединяет кэшированные скрытые состояния предыдущего сегмента с состояниями текущего. Запросы формируются только из текущего сегмента, тогда как ключи и значения охватывают оба сегмента; распространение градиентов останавливается на кэше, поэтому вычислительная стоимость обработки каждого сегмента при обучении остаётся постоянной. Один слой внимания в PyTorch показывает этот подход:

import torch
from torch import nn

torch.manual_seed(0)
seg_len, dim = 4, 16
attn = nn.MultiheadAttention(dim, num_heads=2, batch_first=True)
hidden = torch.randn(1, 3 * seg_len, dim)  # hidden states for three consecutive segments

memory = torch.empty(1, 0, dim)  # nothing is cached before the first segment
for i, segment in enumerate(hidden.split(seg_len, dim=1), start=1):
    # Queries come from the current segment; keys and values also cover the cached one
    context = torch.cat([memory, segment], dim=1)
    out, weights = attn(segment, context, context)
    print(f"Segment {i}: {weights.shape[1]} queries attend over {weights.shape[2]} positions")
    # Cache this segment for the next one; detach() stops gradients as in Transformer-XL
    memory = segment.detach()

Четыре запроса первого сегмента учитывают только четыре токена этого сегмента, а запросы каждого следующего сегмента — восемь: собственные четыре и четыре кэшированных. Поскольку слой n сегмента читает кэшированные состояния слоя n − 1 предыдущего сегмента, с каждым слоем охват увеличивается ещё на один сегмент — примерно до N × L токенов при длине сегмента N и L слоях:

Кэш также ускоряет оценку. Обычный трансформер, предсказывающий токены один за другим, вынужден пересчитывать полное окно для каждого нового токена, тогда как Transformer-XL повторно использует кэшированные состояния, благодаря чему во время оценки он работал до 1 800 с лишним раз быстрее. Функция отслеживания в Ultralytics YOLO26 сталкивается с той же проблемой фрагментации, когда скрипт передаёт видео в model.track() по одному кадру за раз: persist=True сохраняет состояние трекера между вызовами, поэтому каждый объект сохраняет один и тот же ID от кадра к кадру.

Относительное позиционное кодирование#

Повторное использование кэшированных состояний нарушает работу абсолютных позиций: первый токен кэшированного сегмента и первый токен текущего сегмента получали бы одинаковое позиционное эмбеддинг-представление. Вместо этого Transformer-XL вычисляет оценку внимания по содержимому запроса и ключа, а также синусоидальному кодированию расстояния между ними; для содержимого и позиции используются два обучаемых вектора смещения, общих для всех позиций. Поскольку оценки зависят только от расстояний, при оценивании модель может использовать кэш большей длины, чем во время обучения. Поворотные позиционные эмбеддинги (RoPE), используемые в большинстве современных больших языковых моделей (LLM), тоже делают оценки внимания зависимыми от относительной позиции.

Результаты и влияние#

Transformer-XL установил новые рекорды на пяти тестах: 0,99 бита на символ на enwik8, 1,08 на text8, а также перплексия 18,3 на WikiText-103, 21,8 на One Billion Word и 54,5 на Penn Treebank — без дообучения. Обученная только на WikiText-103 модель генерировала достаточно связные новые статьи длиной в тысячи токенов. XLNet, созданная несколькими из тех же авторов, построила предварительное обучение для понимания естественного языка на архитектуре Transformer-XL, а Compressive Transformer от DeepMind расширил кэш, сжимая старые воспоминания вместо их удаления.

Transformer-XL и родственные модели#

• Долгая краткосрочная память (LSTM): обновляет одно скрытое состояние по одному токену за раз, поэтому информация из далёкого прошлого должна пройти через все промежуточные шаги.
• Longformer: обрабатывает целый документ из 4 096 токенов за один раз, используя оконное и глобальное внимание, без рекуррентности.
• KV-кэш: хранит ключи и значения предыдущих токенов во время генерации, чтобы их не приходилось вычислять повторно. Память Transformer-XL также используется при обучении, а распространение градиентов останавливается на границе сегмента.
• Модели пространства состояний (SSM): сжимают всю историю в состояние фиксированного размера.