Длина контекста — это статья расходов
Задача звучала прозаично: разбирать поток тяжелых полуструктурированных выгрузок от контрагентов. Не десятки страниц — сотни тысяч токенов на документ, и таких документов много.
Проблема здесь арифметическая. Механизм внимания в трансформере сопоставляет каждый элемент последовательности с каждым: удвоили длину документа — учетверили работу. На коротком тексте это незаметно. На длинной выгрузке разница между квадратом и линией — это разница между стойкой GPU и одной картой.
Именно из этой арифметики выросло проектное решение. Мода на архитектуру тут ни при чем.
Что дают модели пространства состояний
Архитектуры класса SSM, к которым относится Mamba, обрабатывают последовательность иначе: они идут по ней, поддерживая сжатое состояние, и стоимость растет линейно с длиной. Память при этом не раздувается пропорционально контексту.
Для нашей задачи это оказалось попаданием в цель. Разбор выгрузки — это последовательное чтение с накоплением: границы записей, повторяющиеся блоки, реквизиты в шапке, которые действуют до конца раздела. Такая работа хорошо ложится на модель, идущую по потоку.
Где SSM проигрывает, и мы это знаем
Честная граница, о которой в статьях про Mamba пишут неохотно: сжатое состояние — это потеря. Когда нужно точно вспомнить редкую деталь, встреченную сто тысяч токенов назад, внимание трансформера справляется лучше — оно смотрит на всю историю, а SSM смотрит в свое состояние.
Поэтому конвейер получился гибридным. Поточный разбор и структурная разметка идут на SSM, а точечные задачи, где нужно сопоставить далеко разнесенные фрагменты, отданы отдельному этапу с другой архитектурой и коротким контекстом. Каждая часть занята тем, что у нее получается.
Пропускная способность считается в документах, а не в токенах
Разница в семь раз замерена так, как это ощущает заказчик: сколько реальных документов из его потока обрабатывается за час на одной и той же карте, включая подготовку и постобработку.
Токены в секунду — метрика для сравнения моделей между собой, и она регулярно вводит в заблуждение. Модель может выигрывать по токенам и проигрывать по документам, если ей нужно несколько проходов или если она не держит документ целиком и приходится резать его с перекрытием.
Резать документ дорого во всех смыслах
Обычный способ обойти ограничение контекста — нарезать документ на куски с перекрытием. Это добавляет работы (перекрытия обрабатываются дважды) и ломает смысл: реквизиты из шапки не доезжают до последних записей, а запись, попавшая на границу, разбирается дважды и по-разному.
Возможность взять документ целиком убрала весь класс проблем сшивки. Это тот случай, когда архитектурное решение экономит не только вычисления, но и месяцы отладки логики склейки.
Что стоит проверить до перехода
- Реальное распределение длин. Если 95% документов короткие, а длинных единицы, овчинка не стоит выделки: дешевле обработать хвост отдельно.
- Нужна ли точная память о деталях. Задачи вида «найди упоминание конкретного пункта где-то в документе» дешевле решать поиском по индексу.
- Зрелость инструментов. Экосистема вокруг SSM моложе трансформерной: часть привычных библиотек и приемов недоступна, и это закладывается в сроки.
- Стоимость эксперимента. Замер на своих данных занимает недели, а не кварталы. Делать выбор архитектуры по чужим бенчмаркам — самый дорогой способ ошибиться.
Как это проверялось до внедрения
Смена базовой архитектуры — решение, которое дорого откатывать, поэтому оно принималось по собственному замеру на своих данных. Мы собрали стенд: одна и та же выборка реальных документов заказчика, одно и то же железо, одинаковая постобработка. Дальше замерялись три вещи: сколько документов проходит за час, сколько памяти занимает обработка самого длинного документа в выборке и как меняется качество разбора на записях у границ.
Первые два замера дали разницу, ради которой все затевалось. Третий оказался важнее: на прежней схеме с нарезкой около части ошибок приходилось именно на стыки кусков, и эта доля исчезла целиком. Такой аргумент понятен и техническому директору, и финансовому.
Что осталось у заказчика
Работающий конвейер разбора, стенд для замеров на собственных данных и понимание, за какие деньги обрабатывается один документ. Последнее оказалось самым востребованным: имея цену документа, финансовый блок сам считает, какой объем потока имеет смысл автоматизировать дальше.
Смежные направления — высоконагруженные модели и эксплуатация ML в продакшне.