Можно ли ускорить обучение небольшой языковой модели на длинных последовательностях, если заменить полное внимание разреженным графом, а веса связей определять расстояниями между обучаемыми координатами?

Для проверки появился SpatialLM-Mini — прототип на Python и PyTorch с тремя архитектурами: пространственным вниманием, обычным Transformer и гибридом. На RTX 3050 с 8 ГБ VRAM пространственный вариант обучался в 2,47 раза быстрее нашего Transformer при контексте 16 384 токена. Но он использовал больше выделенной памяти, а отдельную пользу геометрии этот опыт пока не доказывает.

Ниже — устройство модели, условия сравнения и результаты, включая те, которые ограничивают выводы.

Что здесь называется пространством

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

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

В каждой голове скрытое состояние проецируется в две трёхмерные координаты — запроса и ключа. Оценка разрешённой связи рассчитывается через квадрат евклидова расстояния:

q, k = tanh(обучаемые проекции скрытого состояния)
score(t, j, h) = -exp(clamp(log_scale[h])) * ||q[t,h] - k[j,h]||²
                 + edge_bias[h, offset(t,j)]
weights = softmax(score по разрешённым соседям)

Затем значения соседних токенов смешиваются с этими весами, как в обычном внимании. Масштаб и смещение для типа связи обучаются вместе с остальной моделью.

Здесь важно различать выбор соседей и их взвешивание. Граф фиксирован, а веса связей зависят от содержимого текста. Поиска ближайших соседей среди всех токенов нет. Координаты — внутреннее представление для вычисления весов; физического пространства, полигонального рендеринга или RT-ядер в реализации нет.

Если число соседей равно K, геометрическое смешивание требует порядка O(TKD) операций вместо O(T²D). В текущем графе K растёт как размер локального окна плюс логарифм максимальной длины. Проекции, полносвязные слои и выходной слой при этом продолжают выполнять обычные матричные умножения.

Зачем нужен гибрид

Разреженный граф ограничивает прямой доступ к произвольной далёкой позиции. Информация может пройти через несколько слоёв, но это ещё не означает, что модель научится извлекать нужный факт из начала длинного текста.

Поэтому в сравнении есть гибрид из четырёх блоков:

пространственный → полное внимание → пространственный → полное внимание

У пространственной модели все четыре блока используют геометрическое внимание. У контрольного Transformer все четыре блока используют полное причинное внимание. Нормализация, остаточные связи и полносвязная часть присутствуют во всех трёх вариантах.

Условия эксперимента

Все модели обучались с нуля на романе Джейн Остин «Гордость и предубеждение» из Project Gutenberg. Предобученные веса не использовались. От GPT-2 взят только фиксированный токенизатор через tiktoken 0.12.0, со словарём 50 257 токенов.

В этом опыте 16 384 токена — именно subword-токены GPT-2, а не байты и не число параметров. Среднее отношение в обучающей части — около 3,96 UTF-8 байта на токен.

Текст разделён на последовательные части 80/10/10 по символам, после чего каждая часть токенизирована независимо. Получилось 141 979 обучающих токенов, 18 537 validation и 18 225 test. Окна не пересекают границы частей. Контрольные суммы текста и последовательностей ID сохранены в репозитории.

Основные настройки:

Параметр	Значение
Видеокарта	NVIDIA RTX 3050, 8192 МиБ VRAM
Окружение	Windows, Python 3.13.14, PyTorch 2.10.0+cu128
Ширина / слои / головы	128 / 4 / 4
Контекст / batch	16 384 токена / 1 последовательность
Оптимизатор	AdamW, learning rate 0,0006, weight decay 0,01
Ограничение нормы градиента	1,0
Dropout / TF32	0 / выключен
Seed	42, 43, 44

Число параметров почти одинаково: 9 323 248 у пространственной модели, 9 323 392 у Transformer и 9 323 320 у гибрида. Для этого ширина FFN различается: 628, 512 и 570 соответственно. Это дополнительное архитектурное различие, которое нужно учитывать при интерпретации.

Для каждого seed все архитектуры получают одинаковую последовательность обучающих окон. Порядок запуска архитектур меняется между seed. Настройки основного опыта зафиксированы после проверки работоспособности и памяти. Test не используется для выбора настроек, остановки или лучшего checkpoint.

Как удалось уложить 16K в память

Во всех моделях используются одинаковые способы экономии памяти: FP16 autocast при FP32 весах и состояниях AdamW, пересчёт блоков при обратном проходе и расчёт выходного слоя порциями по 512 позиций.

Последний пункт существенен при словаре из 50 тысяч токенов. Вместо хранения выходной матрицы для всего контекста сразу loss считается порциями. При backward эти порции пересчитываются. В каждой позиции softmax по-прежнему охватывает весь словарь, а итоговая ошибка учитывает все 16 384 целевых токена. Эквивалентность ошибки и градиентов обычному расчёту проверена тестами.

Контрольный Transformer использует PyTorch EFFICIENT_ATTENTION с запрещённым переходом на math-backend. Это полное причинное внимание с экономией промежуточной памяти. FlashAttention в использованной Windows-сборке не был скомпилирован. Поэтому результат относится к доступному на этой машине Efficient Attention; переносить соотношение скоростей на другие ядра и видеокарты нельзя.

Сравнение при одинаковом времени

Каждая модель получила 180 секунд накопленного времени шагов обучения. В измерение входят выбор окна, передача на GPU, прямой и обратный проходы, пересчёт промежуточных данных, ограничение градиента и обновление параметров. Прогрев, оценка, сохранение весов и генерация исключены. Это время обучения, а не полная продолжительность запуска программы.

Validation и test оцениваются на одинаковых 16 384 целевых позициях: по одному полному блоку на каждую часть. Ошибка — средняя отрицательная логарифмическая вероятность следующего GPT-2 токена; меньше означает лучше. В таблице указаны среднее и стандартное отклонение по трём seed.

Модель	Обучающих токенов/с	Шагов за бюджет, в среднем	Test loss	Пик выделенной памяти, МиБ
Transformer	13 126	145	5,2176 ± 0,0087	499,1
Пространственная	32 364	356	4,8771 ± 0,0148	671,1
Гибрид	18 725	206	4,8689 ± 0,0125	657,2

Пространственный вариант обучается в 2,47 раза быстрее, гибрид — в 1,43 раза. При одинаковом времени оба показывают меньшую test loss во всех трёх парных запусках. Минимальная средняя ошибка получилась у гибрида, максимальная скорость — у пространственного варианта.

Однако более быстрые модели успевают сделать больше обновлений и обработать больше обучающих токенов. Это часть сравнения при одинаковом времени. Чтобы видеть разницу отдельно от числа обновлений, сохранена ещё одна контрольная точка.

Сравнение после одинаковых 100 шагов

За 100 обновлений каждая модель обработала 1 638 400 целевых токенов.

Модель	Время, с	Test loss, среднее ± SD
Transformer	125,10	5,6905 ± 0,0155
Пространственная	50,64	5,6527 ± 0,0111
Гибрид	87,55	5,6394 ± 0,0139

Здесь ошибки заметно ближе. Основной практический выигрыш текущего опыта — возможность быстрее выполнять обновления на длинной последовательности, а не большое улучшение качества на каждом шаге.

Чего результаты не показывают

Экономии размера модели нет: число параметров специально почти выровнено. Экономии памяти тоже не получено: пространственный вариант использовал больше выделенной PyTorch памяти. Причина в реализации: сбор соседей материализует промежуточные тензоры, тогда как полное внимание использует оптимизированное ядро. Указанные значения не равны полной занятости VRAM драйвером и рабочим столом.

Скорость генерации не сравнивалась. В прототипе пока нет KV-кэша: при каждом следующем токене контекст пересчитывается. Из ускорения обучения нельзя выводить такое же ускорение ответа пользователю.

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

Корпус небольшой: один роман, один длинный блок оценки на часть, три seed и короткое обучение. Стандартное отклонение здесь не является проверкой статистической значимости. В генерации есть повторы и ошибки связности; это не готовый диалоговый ассистент.

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

Что можно воспроизвести

В репозитории SpatialLM-Mini есть код, исходные измерения, конфигурации, данные с контрольными суммами и инструкции на русском и английском.

После подготовки CUDA-окружения по README сначала можно выполнить короткую проверку — по три шага для каждой архитектуры:

python -m spatiallm.subword16k pilot --steps 3 --seeds 42 --out runs/pilot_repeat

А затем запустить основное сравнение:

python -m spatiallm.subword16k run --steps 100 --seconds 180 --seeds 42 43 44 --out runs/subword16k_repeat

На моей RTX 3050 основная серия занимает около 27 минут измеряемого обучения плюс время на оценку и сохранение. Нужны CUDA и поддерживаемый Efficient Attention. Для каждого опыта следует выбирать новый каталог вывода и запускать GPU-проверки последовательно, чтобы они не мешали замерам друг друга.

Веса создаются при обучении и в Git не включены. Для знакомства с кодом есть небольшое CPU-демо, но оно не воспроизводит опыт на 16K.

Локально прошли все 25 тестов, включая проверки причинности и градиентов. CPU-проверки GitHub Actions тоже завершились успешно. Это подтверждает работоспособность реализации, но не заменяет независимого повторения.

Что хочется проверить вместе с сообществом

Пока этот эксперимент оставил больше вопросов, чем ответов. На моей конфигурации пространственная модель обучалась быстрее контрольного Transformer. Но что именно дало преимущество: обучаемые 3D-координаты или сам разреженный граф связей?

Буду благодарен за помощь с проверкой нескольких гипотез:

Нужна ли геометрия? Сохранится ли качество, если оставить тот же граф, но заменить расстояния обычным скалярным произведением или фиксированными весами?
Используется ли далёкий контекст? Может ли модель извлечь нужный факт из начала длинной последовательности? Возможность принять 16 тысяч токенов ещё не отвечает на этот вопрос.
Повторяется ли ускорение? Что получится на другой видеокарте, другом корпусе и при более длительном обучении?
Достаточно ли сильный контрольный вариант? Особенно интересно сравнение с оптимизированным локальным вниманием и другими разреженными схемами.

Сама идея сочетания локальных и глобальных связей уже исследовалась в Longformer, BigBird, Routing Transformer, LongNet и Gemma 2. С ними прямого сравнения пока не было. Поэтому сейчас хочется понять, есть ли у этой конкретной реализации полезное свойство, которое выдержит более строгую проверку.

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

Пространственное внимание на RTX 3050: эксперимент с контекстом 16K