Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Языковая модель получила внешнюю память, которую можно держать в оперативной памяти хоста и менять под нужную предметную область без дообучения. В работе University of Washington и Allen Institute for AI модель на 1 млрд параметров достигла той же ошибки, что и обычная модель с дополнительными 57% вычислений, хотя препринт не рецензирован и числа получили сами авторы. Для команд, которые обучают собственные модели, это новый способ обменять свободную оперативную память на экономию GPU.
COUNT превращает частоты продолжений во входной вектор
COUNT хранит не тексты и не обученные представления, а таблицу частот для n-грамм — коротких последовательностей токенов. Для каждой последовательности таблица перечисляет наиболее частые следующие токены и присваивает им вес: например, после одного и того же фрагмента кода одни продолжения встречаются чаще других.
При обработке очередного токена система находит самый длинный известный фрагмент текущего контекста. Если точного совпадения нет, она переходит к более коротким последовательностям. Все частоты считают заранее по корпусу, поэтому во время обучения таблица остаётся неизменной и не требует градиентов или состояния оптимизатора.
Затем COUNT берёт из выходной матрицы модели векторы вероятных продолжений, складывает их с весами из таблицы и получает дополнительный входной вектор. Его прибавляют к обычному вектору текущего токена перед первым блоком трансформера. Новых обучаемых параметров при этом не появляется: метод повторно использует представления, которые модель уже выучила для предсказания следующего токена.
Это отличает COUNT от обучаемых n-граммных представлений. Последним нужны параметры, быстрые чтение и запись, градиенты и состояние оптимизатора, поэтому они конкурируют с основной моделью за память GPU. Таблица COUNT доступна только для чтения и может оставаться в обычной оперативной памяти.
Свободная память хоста компенсирует часть вычислений
Метод проверяли на трансформерах размером от 60 млн до 1 млрд параметров, обученных на Dolma 3 при трёх соотношениях числа токенов и параметров. COUNT сравнивали с обычным трансформером, обучаемыми n-граммными представлениями и объединением частот с предсказаниями модели только на выходе. Качество измеряли по ошибке предсказания текста, ошибке на прикладных заданиях и точности ответов.
Самая крупная таблица занимала 797 ГБ оперативной памяти. На прикладных заданиях обычному трансформеру требовалось на 6–56% больше вычислений, чтобы сравняться с COUNT по ошибке. Точность ответов росла менее равномерно, но COUNT превосходил стандартную модель почти при всех проверенных бюджетах.
При одинаковом числе обучаемых параметров COUNT чаще попадал на границу лучших конфигураций: добавленная память не вытесняла параметры основной модели из GPU. Обучаемые n-граммные представления иногда давали сопоставимое качество, но занимали память, которую можно было отдать более крупному трансформеру.
Результат зависит не только от ёмкости, но и от пропускной способности памяти и стоимости поиска по таблице. В работе использовали узлы с большим объёмом оперативной памяти и отдельные режимы обучения и запуска; перенос вывода на серверы с меньшей пропускной способностью требует собственных замеров задержки.
Доменные таблицы меняют планы только при собственном обучении модели
Таблицу COUNT можно заменить после обучения, поскольку все варианты используют общий словарь и одну выходную матрицу модели. Для математики, биомедицины и права исследователи построили отдельные таблицы, а затем смешали доменные частоты с исходными в пропорции 1:3. Ошибка снизилась во всех трёх областях, тогда как точность отдельных заданий менялась неодинаково и особенно смешанные результаты дала биомедицина.
Для продуктовой команды это означает, что обновление терминологии или подключение закрытого отраслевого корпуса не обязательно требует нового цикла обучения. Можно пересчитать агрегированные частоты, проверить их на целевых заданиях и подключить другую таблицу при запуске. Таблицу также проще изучать и исправлять вручную, чем знания внутри параметров модели: каждая строка показывает возможные продолжения конкретного фрагмента.
Но COUNT нельзя просто добавить к уже развёрнутой обычной модели и получить результаты статьи. Основную модель заранее обучали принимать дополнительный входной вектор. Если используется готовая сторонняя LLM без доступа к её обучению, практичнее остаются поиск по документам или объединение частот с выходными вероятностями, которое в экспериментах уступило COUNT.
Работа меняет планы прежде всего у команд, которые обучают компактную базовую модель, ограничены памятью GPU, но контролируют серверы и корпус. В таком проекте стоит сравнить две конфигурации: более крупную основную модель с COUNT в оперативной памяти и меньшую модель с обучаемыми дополнительными представлениями. Для API-продукта на закрытой модели архитектурный вывод слабее: переключаемая таблица требует поддержки со стороны самой модели.
Источники
Иллюстрация: рисунок из статьи «The COUNT: Counting N-grams Unlocks Scaling on Host Memory and Training-Free Adaptation», Ian Magnusson, Noah A. Smith, Pang Wei Koh, CC BY 4.0
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
