Nvidia ускорила смену ИИ-моделей простой математикой

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

Почему смена модели в середине сессии обходится так дорого

Когда большая языковая модель (LLM) получает запрос, она сначала проходит этап «префилла»: вычисляет ключи и значения для всех входных токенов и заполняет KV-кэш. Затем начинается этап декодирования, на котором модель читает из кэша и генерирует следующие токены. При длинных диалогах или агентных сессиях контекст разрастается, и стоимость префилла растёт линейно с длиной и размером модели.

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

Перенос памяти между моделями без перезапуска

Исследователи из Nvidia изучили, можно ли трансформировать KV-кэш одной модели в формат другой без повторного префилла. Перенос от маленькой модели к большой повышает качество ответов, а от большой к маленькой — снижает затраты. Например, мощная модель может обработать сложный документ в начале сессии, а затем передать кэш более дешёвой модели для быстрых ответов.

Ранее предпринимались попытки решить эту проблему, но они требовали либо дорогого обучения с градиентами, либо жёстких архитектурных ограничений. Новое исследование сосредоточено на внутрисемейственных переносах — например, между моделями семейств Qwen, Llama или Ministral разного размера. Они имеют общий токенизатор и схожую архитектуру, но различаются числом слоёв и параметров.

Ключевой вывод: структура KV-кэша оказалась линейной. Это позволяет выполнять маппинг с помощью простых алгебраических операций. Например, при переносе от Qwen3 14B к Qwen3 32B линейная регрессия восстанавливает до 56% вариации ключей и 32% значений на одном слое. При комбинировании нескольких слоёв показатели выросли до 79% и 65% соответственно.

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

Испытание линейного маппера

Чтобы проверить метод, исследователи оценили перенос на шести парах моделей с совпадающим числом KV-голов и размерностью. Среди них были скачки от Llama 3.1 8B до 70B (увеличение параметров в 8.8 раза). Тесты включали пять бенчмарков точности, а также оценку перплексии и многодиалоговую задачу CoQA. Для подгонки линейного маппера использовался крошечный калибровочный набор из 500 последовательностей.

Для четырёх из шести пар линейный метод сохранил от 73% до 98% точности целевой модели после полного префилла. При этом скорость переноса оказалась в 2.7–25 раз выше, чем повторный префилл. Например, перенос кэша из 32 768 токенов от Qwen3 14B к 32B занял всего 278 мс против почти 7 секунд на обычный префилл. Многодиалоговые тесты показали минимальный дрейф точности на протяжении 10 ходов.

Однако для двух пар Ministral линейный метод деградировал. Для их восстановления применили нелинейный многослойный перцептрон (MLP) с двумя скрытыми слоями. Это добавило сложности и затрат на обучение, но точность вернулась к уровню выше 90%.

Ключевые особенности решения Nvidia:

• Использование закрытой формы гребневой регрессии для каждой головы внимания.
• Выбор наиболее предсказательных слоёв исходной модели для каждого целевого слоя.
• Отбрасывание RoPE-кодировок для обобщения на более длинные последовательности.

Nvidia также активно развивает другие методы сжатия кэша: динамическое разрежение памяти (DMS) снижает затраты на рассуждения до 8 раз, а KV Cache Transform Coding (KVTC) сжимает память в 20 раз. Исследователи из MIT предложили алгебраическое сжатие Attention Matching с уменьшением кэша в 50 раз без потери качества.

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

Читайте также: Как понять, что ваш компьютер заражён вирусом

Источник: VentureBeat