KV-PRM: Ускорение LLM в 37 раз за счет отказа от перекодировки текста

Проблема текстовых PRM

Процессуальные модели вознаграждения (Process Reward Models, PRM) критически важны для масштабирования тестового времени (Test-Time Scaling) в мультисистемах на базе больших языковых моделей. Однако существующие PRM работают исключительно с текстом: они заново кодируют всю траекторию рассуждений модели. В длинных сессиях с участием нескольких агентов стоимость такого скоринга растет квадратично относительно длины последовательности ($O(L^2)$), создавая узкое место, которое делает применение PRM в длинных контекстах вычислительно неэффективным.

Решение: KV-PRM

Авторы предлагают KV-PRM — архитектуру, которая устраняет необходимость повторного кодирования текста. Вместо этого метод напрямую считывает KV-кэш (Key-Value Cache), который естественным образом формируется в процессе генерации токенов основной LLM. Оценка производится путем обработки одного «верификационного токена» (verify token) против уже существующего кэша. Это снижает сложность вычислений до линейной ($O(L)$).

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

Результаты бенчмарков

Эксперименты проводились на стандартных наборах данных для математического мышления: MATH, GSM8K и AIME. KV-PRM демонстрирует результаты, сопоставимые или превосходящие текстовые PRM, при использовании различных методов масштабирования, таких как Beam Search, MCTS (поиск по дереву Монте-Карло) и Weighted Voting.

Значение для индустрии

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

Источник: arXiv cs.AI ↗