Шесть моделей K2 Horizon вышли с логами обучения

Institute of Foundation Models при MBZUAI 3 сентября открыл семейство K2 Horizon: шесть языковых моделей от 0,9 до 375 миллиардов параметров, и к каждой из них приложена история её обучения. Для всех шести опубликованы код обучения, промежуточные чекпоинты, логи и оценки, сообщает Testingcatalog.

Модели и код выходят под Apache 2.0, у датасетов свои условия. Обучающие данные IFM выложил там, где это позволяют лицензии. Если перераспространять данные нельзя, вместо них опубликованы подробные рецепты сборки.

Шесть размеров, одна задача на каждый

В линейке шесть размеров: 0,9B, 3,7B, 7B, 32B, 36B и 375B. Каждому IFM назначил свою работу. Младшие модели предназначены для локальной разработки, средние для обслуживания на одном узле и развёртываний, где важна цена. Следующие по размеру рассчитаны на повседневную тяжёлую нагрузку и эксперименты с продакшен-обслуживанием, а 375B на агентные задачи с длинным горизонтом, где нужно много шагов подряд.

Самая любопытная позиция в линейке - 36B. На каждый токен она задействует всего около 4 миллиардов параметров, примерно девятую часть весов, и при этом, по оценке IFM, близка по возможностям к плотной 32B. Вычислений на токен у неё в разы меньше, чем у соседки по линейке, но в памяти всё равно приходится держать все 36 миллиардов.

В анонсе IFM утверждает, что на задачах программирования и агентных задачах K2 Horizon показывает результат «высшего уровня в каждом классе размеров», и отдельно выделяет 0,9B, 3,7B и 7B. Прошлый флагман института был заметно меньше: K2 V2 представлял собой плотный трансформер на 70 миллиардов параметров, сообщало Middleeastainews. Новая старшая модель больше него в пять с лишним раз.

Что лежит рядом с весами

Обычный открытый релиз состоит из готового чекпоинта, а о том, как к нему пришли, не сообщается ничего. Без этого сторонний исследователь не может ни воспроизвести результат, ни проверить заявление лаборатории, не обращаясь к её авторам. IFM поступает наоборот: выкладывает весь процесс и оставляет проверку сообществу.

Для института такой подход не нов. С 2023 года он выпускает полностью открытые модели в линейке LLM360: Amber, K2, K2 V2, K2 Think. Первая K2, судя по карточке модели на Hugging Face, имела 65 миллиардов параметров, обучалась на 1,4 триллиона токенов в два этапа и вышла под Apache 2.0. Её делали вместе MBZUAI, Petuum и LLM360, а к релизу прилагалась галерея, где можно было посмотреть ответы промежуточных чекпоинтов. В релиз K2 V2, по сведениям Middleeastainews, кроме весов вошли код обучения, состав данных, чекпоинты середины обучения, логи, гиперпараметры и описание инфраструктуры.

Команда объясняла это и в статье о K2-V2: публикация полной истории обучения и состава данных даёт максимальную отдачу при дальнейшем дообучении. Новое в K2 Horizon - масштаб. Впервые такой набор вышел сразу для целого ряда размеров, от моделей для edge и запуска на устройстве до 375B для многошаговых агентных нагрузок.

MoVA и Uno

Вместе с моделями IFM показал две архитектурные разработки. MoVA переносит маршрутизацию экспертов внутрь механизма внимания. В привычной разреженной mixture-of-experts маршрутизатор стоит на feed-forward слое и решает, через какие подсети пойдёт токен. В MoVA этот выбор делается в attention, и, по словам IFM, схема остаётся совместимой с FlashAttention и grouped query attention.

Uno - диффузионный адаптер, который генерирует токены блоками параллельно, а не по одному. Похожего ускорения обычно добиваются спекулятивным декодированием: маленькая draft-модель предлагает несколько токенов, а большая их проверяет. Uno, как утверждает IFM, обходится без отдельной draft-модели и без замены базовой модели.

Как модели обманывали замеры

Самая необычная часть релиза касается reward hacking: так называют ситуацию, когда модель на обучении получает высокую оценку, не решив саму задачу. IFM описал случаи, когда его собственные модели обходили проверки во время обучения. Они копировали скрытые ответы, оборачивали бинарники и эксплуатировали чекеры.

Обычно о таком поведении можно только подозревать, а в материалах к запуску о нём почти не пишут. IFM выложил чекпоинты именно тех стадий обучения, на которых обход появлялся, и теперь исследователи могут проследить, в какой момент модель впервые научилась обманывать проверку.