Дополнение · X-алгоритм · страница 4 из 11

Retrieval: две башни, хеши и semantic IDs

Как из 28 миллионов постов за считанные миллисекунды достать несколько сотен кандидатов. Разбираем двухбашенную модель Phoenix: почему у пользователя нет собственного эмбеддинга, как пост представляется без словаря идентификаторов, что такое semantic ID и чем обучение здесь отличается от того, как двухбашенные модели учат обычно.

Карта страницы

1. Зачем вообще две башни

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

Корпус кандидатов — 28 672 000 постов (константа max_posts в конфиге). На запрос надо отдать несколько сотен. Если бы модель считала совместную функцию \(f(\text{пользователь}, \text{пост})\), пришлось бы вызвать её 28 миллионов раз — это невозможно ни за какое разумное время.

Двухбашенность — это ограничение, наложенное сознательно: скор обязан раскладываться в произведение двух независимо вычисляемых векторов,

Плата известна: башни не видят друг друга, поэтому никакие «перекрёстные» признаки — совпадение языка пользователя и поста, время с момента подписки на автора — модель выразить не может. Выигрыш тоже известен: \(\mathbf{p}_i\) для всех постов считаются заранее, а на запросе остаётся посчитать \(\mathbf{q}_u\) один раз и найти ближайших соседей.

def _l2_normalize_candidates(embeddings: jax.Array) -> jax.Array:

2. У пользователя нет эмбеддинга

Самое неожиданное решение в этой модели, и в конфиге оно записано прямо:

"use_user_embedding": False,
"use_post_embedding": False,

Фрагмент из xrecsys_two_tower.py · код X, Apache 2.0, коммит 28e414f

Ни у пользователя, ни у поста нет собственной обучаемой строки в таблице. Пользователь представлен исключительно тем, с чем он взаимодействовал: последовательностью до 1023 событий истории плюс один токен с грубыми профильными признаками — страна, язык, локация, пол, возраст, установленные приложения.

Цена тоже есть, и о ней стоит сказать честно. Устойчивые долгосрочные предпочтения, которые не выражаются в недавних действиях, модель выразить не может. Человек, который годами любит астрономию, но последнюю неделю читал про футбол, для этой модели — любитель футбола. Обучаемый эмбеддинг такое бы запомнил. Компромисс выбран в пользу свежести, и для ленты, где контент живёт часы, он оправдан.

"empty_history_user_dropout_rate": 0.1,

3. Хеш-эмбеддинги: как обойтись без словаря

Автор поста, пользователь, IP — всё это категориальные признаки с огромным числом значений. Классический подход — словарь «значение → индекс строки» — здесь не работает: словарь надо строить, хранить, пересобирать, а новые значения появляются каждую секунду.

Вместо словаря используется универсальное хеширование. Функция в коде выглядит так:

raw = (ids[i] * scales[j] + biases[j]) % modulus

Фрагмент из recsys_embedding.py · код X, Apache 2.0, коммит 28e414f

то есть \(h_j(\text{id}) = \bigl((\text{id} \cdot a_j + b_j) \bmod M\bigr) \bmod m\), где \(a_j, b_j\) — фиксированные константы, \(M\) — большое простое, \(m\) — число строк таблицы. И ключевое: таких функций две.

Это ровно та схема, которую мы разбирали в главе «Unified Embedding» под названием Unified Embedding, и она же стоит за виджетом хеширования: вероятность того, что два значения окажутся неразличимы, равна не \(p\), а \(p^2\). При загрузке таблицы, скажем, в 5% это даёт падение с 5% до 0.25%.

4. Semantic ID: пост без идентификатора

Самая интересная часть. Раньше пост представлялся хешами своего идентификатора. Сейчас — semantic ID, и в конфиге это записано так:

"use_post_embedding": False,
"use_post_sid": True,
"sid_num_levels": 6,
"sid_codebook_size": 256,
"sid_embed_dim": 1024,

Фрагмент из xrecsys_two_tower.py · код X, Apache 2.0, коммит 28e414f

Как получается такой код

Механизм — остаточное квантование, ровно то, что мы разбирали в главе «Generative retrieval и semantic IDs». Берётся мультимодальный эмбеддинг поста (текст плюс изображение) размерности 1024 и кодируется шестью уровнями:

1. Первый кодбук из 256 векторов. Находим ближайший — это первый токен \(c_1\).
2. Вычитаем его из исходного вектора. Остаётся остаток — то, что первый уровень не объяснил.
3. Второй кодбук из 256 векторов, но уже для остатков. Ближайший — токен \(c_2\).
4. И так шесть раз.

Пост превращается в кортеж из шести чисел, каждое от 0 до 255: например, \((37,\ 210,\ 4,\ 118,\ 250,\ 91)\). Всего различимых кодов — \(256^6 \approx 2.8 \cdot 10^{14}\), с огромным запасом на любой каталог.

Ещё две строчки конфига стоят пояснения:

• sid_hash_level: True — коды разных уровней не делят одно пространство, а разводятся хешированием: код 37 первого уровня и код 37 второго — разные векторы. Логично, ведь они кодируют разное: первый — грубую тему, второй — уточнение остатка.
• sid_cross_attn: True — шесть токенов кода не просто складываются, а обрабатываются вниманием. Уровни могут смотреть друг на друга, и модель сама решает, какой уровень важнее для конкретного поста.

5. Обучение: sampled softmax с двумя видами негативов

Здесь мы встречаем в живом коде почти всё, что выводили в главе «Двухбашенные модели». Разберём по частям.

Что считается позитивом

Наивный ответ — «взаимодействие». Реальный оказывается строже:

valid_positive_mask = (
    has_positive_actions
    & ~has_hard_negative_actions
    & ~has_soft_negative_actions
    & ad_mask_candidates
    & candidate_padding_mask[:, :C]
)

Фрагмент из recsys_two_tower_model.py · код X, Apache 2.0, коммит 28e414f

Позитивом считается пост, на котором произошло положительное действие и не произошло ни одного отрицательного. Причём отрицательные разделены на два сорта:

Обратите внимание на первую строку: в базовой конфигурации retrieval обучается на лайках, и больше ни на чём. Не на репостах, не на кликах, не на времени просмотра. Это осознанное сужение: retrieval должен грубо очертить область интересов, а тонкие различия — работа ранжирования.

Два вида негативов

Негативы берутся из двух источников одновременно, и это отличается от того, что обычно рассказывают:

Зачем оба, если in-batch бесплатны? Потому что у них разное распределение, и мы это подробно разбирали в главе «Сэмплирование негативов».

• In-batch негативы приходят из распределения популярности: чем чаще пост лайкают, тем чаще он оказывается чьим-то позитивом и, значит, чужим негативом. Они автоматически «трудные» — популярное обычно и правда неплохое, — но систематически смещённые.
• Глобальные негативы сэмплируются отдельно и покрывают весь корпус, включая хвост, который в батч почти никогда не попадает. Они лёгкие, но именно они не дают модели схлопнуться на голове распределения.

В конфиге, кстати, видно интересное распределение ролей: num_negatives_per_example: 0 при num_global_negatives_per_example: 64. То есть «локальных» дополнительных негативов нет вовсе — работают in-batch плюс 64 глобальных.

LogQ-коррекция — в живом коде

И вот то, ради чего стоило сюда добираться. In-batch негативы смещены к популярному, и без поправки модель выучит «популярное значит плохое». В коде поправка есть:

def _apply_logq_correction(
    local_batch, local_neg, local_batch_correction, global_batch_correction
):
    if use_in_batch_negatives:
        local_batch_correction = local_batch_correction.reshape((1, -1))
    local_batch += local_batch_correction
    if N > 0:
        local_neg += global_batch_correction.reshape((1, -1))
    return local_batch, local_neg

Фрагмент из recsys_two_tower_model.py · код X, Apache 2.0, коммит 28e414f

Обратите внимание на две вещи. Первая: поправка прибавляется к логитам — это ровно \(s - \log Q\) из вывода LogQ-коррекции, с точностью до знака, зашитого в то, как считается сама поправка. Вторая: поправки две разные — своя для in-batch негативов и своя для глобальных, потому что распределения у них разные, и одной коррекцией их не выправить.

Масштаб коррекции — отдельный параметр:

logq_correction_scale: float = 2.0

Фрагмент из recsys_two_tower_model.py · код X, Apache 2.0, коммит 28e414f

То есть поправка применяется не как \(s - \log Q\), а как \(s - \lambda \log Q\) с \(\lambda = 2\). Теория говорит, что несмещённая оценка получается при \(\lambda = 1\); значение 2 означает сознательную перекоррекцию — популярное давится сильнее, чем требует математика.

Две скорости обучения

"learning_rate": 2e-3,
"emb_learning_rate": 0.1,

Фрагмент из xrecsys_two_tower.py · код X, Apache 2.0, коммит 28e414f

Скорость обучения эмбеддингов в 50 раз выше, чем у остальной сети. Это стандартная практика для рекомендательных моделей, и причина в разреженности градиента: плотные слои получают градиент на каждом примере батча, а конкретная строка таблицы эмбеддингов — только тогда, когда её айтем попал в батч, то есть в тысячи раз реже. Чтобы она успевала обучиться, шаг должен быть больше.

6. Индекс: сервинг ничего не считает

Как эти эмбеддинги попадают в поиск? Ответ в README неожиданный и красивый: индекс лежит внутри чекпоинта.

На каждом сохранении модели тренер прогоняет башню кандидатов по всему корпусу и записывает получившиеся 28.7 миллиона векторов в тот же файл, что и веса. Сервинг загружает чекпоинт — и индекс уже там.

7. Конфигурация продовой модели одним экраном

Последняя строка любопытна. У башни кандидатов две головы с разными определениями позитива: для ленты позитив — лайк, для полноэкранного видеорежима — лайк, ответ, цитата, репост, качественный просмотр, подписка, закладка или «поделиться». Одна модель, два разных представления о том, что такое «хорошо», в зависимости от поверхности. Это тот же аргумент про разные поверхности и разные целевые функции.

Частые ошибки и подводные камни

Вопросы с собеседований

Шпаргалка одним экраном

Первоисточники

• recsys_two_tower_model.py — башни, лосс, LogQ-коррекция, определение позитива.
• xrecsys_two_tower.py — продовые конфиги: длины, размерности, число негативов, параметры semantic ID.
• recsys_embedding.py — хеш-функции, размеры таблиц, смещения зон.
• phoenix/README.md — описание архитектуры от авторов.
• Курс: глава «Двухбашенные модели» целиком, глава «Индуктивный bias и контент» об индуктивном смещении, глава «Unified Embedding» о Unified Embedding, глава «Generative retrieval и semantic IDs» о semantic IDs.