OneRec разбор (часть 2): токенизация

Токенизация - небходимый элемент моделей на основе трансформеров. Задача токенизации разбить вход на небольшие кусочки, которые трансформер будет учиться комбинировать. В NLP рецепт уже более-менее общий: разновидности BPE, O(100k) токенов, небольшие отличия в инженерных трюках (как обрабатывать пробелы и пунктуацию, разбивать ли числа на отдельные цифры, какие спец. токены добавить), после обучаемый словарь эмбеддингов ([1], [2], [3]). В vision language models рецепт токенизации пока не устоялся. Изображение обычно разбивается на патчи, которые пропускают через предобученную визуальную модель (An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale). Далее визуальные токены либо квантизуют в дискретное представление, либо подают на вход LLM, пропустив через небольшой адаптер. Основные design choices: какой выбрать визуальный энкодер (архитектура, задача обучения, датасет), как сжать визуальные токены перед входом в LLM (Q-Former, Perciever, Windowed Attention), как (и надо ли) превратить их в дискретные представления. В audio моделях ситуация очень похожая: аудио дорожка нарезается на отрезки, кодируется, выход подаётся как есть, либо дискретизуется (Audio Flamingo 3).

Рекомендательные трансформеры устроены похожим образом. История пользователя естественным способом разбивается на "кусочки" - отдельные события, перед входом в модель они пропускаются через специальны адаптер, обучаемый вместе с трансформером. Энкодеры бывают как предобученные, так обучающиеся совместно с основной моделью. Проблема такого способа токенизации - он не подходит для генерации. В других областях также часто используют токены разных видов для входа и выхода модели. Первыми решение проблемы предложили в Deepmind в статье TIGER. Идея заключается в том, чтобы построить машинно обученное дерево категорий документов. Таким образом каждое событие распадается на несколько токенов, каждый из которых уточняет предыдущий. Идею подглядели в CV.

Некоторые плюсы и минусы semantic ids:
➕Не нужно использовать гигантские эмбеддинг таблицы для item id
➕Токены меньше переобучаются, отсутствует one epoch phenomenon
➕ Используется полный softmax loss, вместо сэмплированной версии
➖ Кодируется только сам документ, контекст игнорируется
➖ Практически вся мемоизация уносится в веса трансформера
➖ Не гарантируется попадание похожих документов в один кластер
➖ Процесс обучения RQ-VAE нестабилен, есть эффект "схлопывания" кластеров

Направление рекомендательной токенизации сейчас активно развивается ([1], [2], [3], [4]). В Kuaishou предлагают свой способ. Его основные идеи:
1. Использовать в качестве семантического энкодера предобученную VLM
2. Сжать выход с помощью QFormer, чтобы уменьшить размерность
3. Дообучить модель на коллаборативно близких парах документов, чтобы уменьшить проблему мемоизации в весах словаря
4. Дополнительно подать выход QFormer в LLaMA 3 и навесить loss на captioning задачу, чтобы модель не разучилась понимать семантический смысл документов
5. На выходах QFormer запусть RQ-Kmeans, вместо изначального RQ-VAE

Большинство идей уже были описаны в их предыдущей статье, однко в OneRec Technical Report рецепт значительно изменили. Что нам нравится и не нравится в их подходе:

➕ Добавлять коллаборативный сигнал в токенизацию точно нужно, причём делать это на уровне контентной модели кажется проще, чем в качестве дополнительного лосса кластеризации (как в LETTER).
➕ RQ-Kmeans выглядит интересно. Кластера в RQ-VAE не сбалансированы (как по количеству, так и по популярности документов), часто становятся пустыми. Kmeans позволяет избежать этих проблем.
➖В целом конструкция получилась довольно громоздкой, с большим количетсвом моделей, хочется попробовать её упростить. Начнём точно с того, что дообучим какую-то модель на коллаборативный сигнал, без использования QFormer и дополнительной LLM после.

Мы сейчас активно экспериментируем с токенизацией. Первые результаты (RQ-VAE над GME-Qwen2VL) получились неплохими, удалось обогнать хэшированные sparse ids. Расскажу об этом в следующих постах.