SpatialLM-Mini ускорил обучение на контексте 16K в 2,47 раза на RTX 3050

Длинный контекст обходится дорого ещё до того, как модель начинает приносить пользу. При полном внимании каждый новый токен может обращаться ко всем предыдущим, и объём вычислений растёт квадратично. В эксперименте SpatialLM-Mini этот доступ ограничили: оставили ближайшие позиции и несколько дальних связей. На RTX 3050 с 8 ГБ памяти прототип обучался в 2,47 раза быстрее контрольного Transformer при окне 16 384 токена.

Для разработчика здесь есть конкретная зацепка: сокращение числа связей может ускорить эксперименты с длинными последовательностями даже на недорогой видеокарте. Но результат относится к небольшой модели, одному корпусу и определённому вычислительному ядру. Польза именно «пространственного» способа расчёта внимания пока остаётся открытым вопросом.

Как токены выбирают, кого слушать

Представьте длинный документ. В обычном причинном внимании токен в конце может напрямую учитывать любую предыдущую позицию. В SpatialLM-Mini список доступных позиций задан заранее: сам токен, несколько ближайших предшественников и позиции на расстояниях 1, 2, 4, 8 и далее. Обращения в будущее запрещены.

Дальше модель решает, сколько внимания уделить каждому разрешённому соседу. Для этого скрытое состояние преобразуется в обучаемые трёхмерные координаты запроса и ключа. Чем меньше расстояние между ними, тем выше может быть вес связи; дополнительные обучаемые параметры корректируют этот расчёт.

Координаты здесь служат способом вычислить веса. Модель не ищет ближайшие токены по всему документу и не использует RT-ядра видеокарты. Список соседей фиксирован, а распределение внимания внутри этого списка зависит от текста.

Такой граф сокращает вычисления, но ограничивает прямой доступ к далёким фактам. Поэтому автор проверил и гибрид: два пространственных блока чередуются с двумя блоками полного внимания. Это позволяет части слоёв обращаться ко всему предыдущему контексту.

Что показало сравнение на 8 ГБ

Три модели примерно по 9,3 млн параметров обучали с нуля на «Гордости и предубеждении» Джейн Остин. У каждой было четыре слоя, одинаковый контекст и три запуска с разными начальными случайными состояниями. От GPT-2 использовали только токенизатор.

При бюджете 180 секунд измеряемого обучения контрольный Transformer обрабатывал 13 126 токенов в секунду, пространственная модель — 32 364, гибрид — 18 725. У гибрида получилась самая низкая средняя ошибка на тестовой части. Однако быстрые модели за это время успели обработать больше токенов и выполнить больше обновлений.

При сравнении после одинаковых 100 шагов разница в ошибке заметно сократилась. Пространственной модели понадобилось 50,64 секунды, Transformer — 125,10 секунды. Главный подтверждённый выигрыш этого опыта — скорость обновлений на длинной последовательности.

С памятью получилось иначе: пик выделенной PyTorch памяти вырос примерно с 499 до 671 МиБ. Разреженная реализация создаёт промежуточные тензоры при сборе соседей, а полное внимание использует оптимизированное ядро. Эти цифры не отражают всю занятую видеопамять компьютера.

Что проверять перед переносом в свой проект

Контрольный Transformer работал через PyTorch Efficient Attention. FlashAttention в использованной Windows-сборке отсутствовал. Поэтому коэффициент ускорения нельзя автоматически переносить на другую систему или сравнение с другими ядрами.

Есть и более существенный пробел: одновременно изменили набор связей и способ вычисления их весов. Чтобы понять вклад координат, нужен тот же разреженный граф с обычным скалярным произведением или фиксированными весами. Пока преимущество может объясняться прежде всего сокращением числа взаимодействий.

Приём 16 тысяч токенов также не доказывает, что модель умеет доставать нужный факт из начала документа. Для этого нужен отдельный тест. Скорость генерации не измеряли, KV-кэша в прототипе нет, а короткое обучение на одном романе не позволяет судить о пригодности для рабочих задач.

Практический следующий шаг — повторить сравнение на своём корпусе и оборудовании, затем добавить контроль с тем же графом без координат. По описанию автора, код, конфигурации и исходные измерения доступны в репозитории SpatialLM-Mini. Для команды, которая исследует дешёвое обучение длинного контекста, это полезная отправная точка для собственного эксперимента.