Baidu открыла код и веса Unlimited OCR 2 для обработки длинных документов

Обычная схема обработки длинного документа постепенно накапливает в памяти KV-кеш – промежуточные данные внимания, к которым модель обращается при генерации следующих токенов. Если хранить визуальные данные сразу для всех страниц, потребление памяти растёт вместе с документом. Unlimited OCR 2 обрабатывает страницы последовательно и удерживает ограниченный набор данных, нужных на текущем шаге. Baidu выложила код и веса на GitHub и Hugging Face под лицензией MIT.

Первую часть задачи решает SR-SWA, механизм разреженного внимания со скользящим окном. Page Gate выбирает страницы, к которым модели полезно обратиться при разборе текущего участка; Token Sieve затем оставляет наиболее подходящие фрагменты их визуальных токенов. Вместе с окном недавнего контекста это ограничивает объём информации, который модель просматривает на каждом шаге. По данным Baidu, в тесте на документах примерно из ста страниц один шаг декодирования с SR-SWA выполнялся в семь раз быстрее, чем при полном внимании, и в 5,6 раза быстрее прежнего механизма.

Вторую часть решает Unlimited-Horizon Decoding: система делит документ на блоки страниц и обрабатывает их по очереди. Соседние блоки перекрываются одной страницей, чтобы сохранить связь между ними; ненужный визуальный KV-кеш предыдущего блока освобождается. В тесте на документах примерно из ста страниц Baidu сообщает о снижении пикового объёма KV-кеша с 47 до 1,77 ГБ, то есть примерно на 96%. Это сравнение касается только кеша: для запуска модели также нужна память под её веса и другие данные.

Baidu продемонстрировала непрерывную обработку семи книг о Гарри Поттере общим объёмом более 3600 страниц. В собственных оценках компании модель набрала 84,6 балла на MPDocBench для многостраничных документов и 77,1 на MarathonDocBench для более длинных документов. Последний набор Baidu выпустила вместе с моделью. Независимых оценок на момент публикации этих результатов ещё не было, поэтому демонстрацию длины обработки и показатели качества стоит рассматривать отдельно: способность пройти тысячи страниц сама по себе не подтверждает точность распознавания каждой из них.