Mixture of Depths (MoD)

Mixture of Depths (MoD) — это метод для трансформеров, в котором обучаемый маршрутизатор на каждом блоке решает, какие токены обрабатывать, а какие пропускать через блок по остаточному соединению. Благодаря этому разные токены получают разный объём вычислений. Дэвид Рапосо, Адам Санторо и их коллеги из Google DeepMind представили метод в апреле 2024 года в статье «Mixture-of-Depths: динамическое распределение вычислений в языковых моделях на основе трансформеров». Фиксированная ёмкость k задаёт количество токенов, обрабатываемых каждым маршрутизируемым блоком, поэтому бюджет вычислений и формы тензоров известны заранее, хотя выбор токенов зависит от входных данных. Большие языковые модели (LLM) MoD показали качество на уровне базовых моделей при одинаковых обучающих FLOPs и времени выполнения, используя лишь часть FLOPs за прямой проход. При генерации их скорость была более чем на 50% выше. При фиксированном бюджете обучения лучшая модель MoD была больше лучшей базовой модели и достигла меньшего значения функции потерь.

Как работает маршрутизация MoD#

Перед маршрутизируемым блоком линейный маршрутизатор присваивает каждому токену скалярную оценку. Токены с k наибольшими оценками проходят через самовнимание и MLP блока; их выходы умножаются на оценки маршрутизатора и добавляются обратно в остаточный поток. Умножение на оценку включает маршрутизатор в путь градиента, поэтому обратное распространение ошибки может научить его определять, каким токенам полезна обработка. Остальные токены проходят без изменений и не участвуют во внимании этого блока, поэтому последующие токены не могут учитывать их в этом блоке. Затраты на внимание уменьшаются пропорционально квадрату ёмкости: обработка половины токенов требует четверти произведений «запрос — ключ». В статье в качестве лучшего параметра указана маршрутизация каждого второго блока с ёмкостью 12,5% последовательности. Минимальная реализация одного маршрутизируемого блока, здесь только с MLP, на PyTorch:

import torch
from torch import nn

torch.manual_seed(0)
batch, seq_len, dim = 2, 64, 128
capacity = int(0.125 * seq_len)  # 12.5% of tokens, the paper's best setting
x = torch.randn(batch, seq_len, dim)
router = nn.Linear(dim, 1)
block = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))

# The router scores every token, and only the top-k tokens enter the block
weights, idx = router(x).squeeze(-1).topk(capacity, dim=1)
idx = idx[..., None].expand(-1, -1, dim)
update = block(x.gather(1, idx)) * weights[..., None]  # router weight keeps the router trainable

# Selected tokens get the residual update; the rest skip the block unchanged
out = x.scatter_add(1, idx, update)
unchanged = (out == x).all(dim=-1).sum().item()
print(f"Block ran on {batch * capacity} of {batch * seq_len} tokens; {unchanged} skipped it unchanged")

Он выводит Block ran on 16 of 128 tokens; 112 skipped it unchanged: MLP блока выполнила одну восьмую обычного объёма работы, а выходной тензор сохранил полную форму последовательности для следующего блока.

Маршрутизация во время генерации#

Выбор top-k не является причинным: попадание токена в top-k последовательности зависит от оценок токенов, расположенных после него, которых ещё нет, когда модель генерирует текст по одному токену за раз. Авторы сделали маршрутизацию причинной двумя способами: добавили вспомогательную двоичную кросс-энтропийную функцию потерь, которая обучает выход маршрутизатора предсказывать принадлежность к top-k и увеличивает функцию потерь языкового моделирования всего на 0,2–0,3%, либо использовали небольшую отдельную MLP-модель-предсказатель на входах с остановленным градиентом, которая быстро достигла точности 97–99%. В обоих случаях качество авторегрессионной генерации снижалось незначительно. Обученные маршрутизаторы также демонстрируют закономерность: одни токены задействуют каждый маршрутизируемый блок на всех уровнях модели, а другие при возможности пропускают блоки.

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

• Смесь экспертов (MoE): Направляет каждый токен к одной или нескольким экспертным MLP в пределах слоя, увеличивая количество параметров при фиксированных вычислительных затратах на токен. В статье также рассмотрено сочетание двух методов под названием MoDE: токен можно направить эксперту или пропустить без изменений. Эффект от этих методов суммировался.
• Досрочный выход: Токен выходит после одного из слоёв и пропускает все оставшиеся. MoD может пропустить один блок, продолжив обработку в последующих блоках.
• Объединение токенов (ToMe): Объединяет похожие токены в визуальном трансформере, чтобы последующие слои обрабатывали меньше токенов; обучаемого маршрутизатора нет.
• Обрезка модели: Безвозвратно удаляет веса, поэтому для каждого входа используется одна и та же уменьшенная сеть.
• Маршрутизация моделей: Направляет целый запрос к одной из нескольких моделей. Это та же идея, применённая на уровне полных моделей.

Большинство сетей имеют фиксированную глубину. Ultralytics YOLO26 задаёт её для каждого масштаба модели: коэффициент глубины в её файле модели YAML равен 0,50 для масштабов n, s и m и 1,00 для l и x. Поэтому этап C3k2, указанный с двумя повторениями, выполняет один блок в YOLO26n и два — в YOLO26x, причём каждое изображение проходит через все эти блоки.