EmbeddingGemma 2: как перейти с 768 на 512 координат и сэкономить память на RTX 5090

Пять миллионов 768-мерных эмбеддингов в FP32 занимают 15,36 ГБ. На RTX 5090 с её 32 ГБ VRAM это почти половина, причём только под координаты. Поисковому индексу нужна ещё служебная память ^[1], а если на той же карте работает модель, место потребуется и ей.

У EmbeddingGemma 2 есть возможность сократить этот объём. Последний линейный слой модели преобразует 512 чисел в 768. В точной арифметике его выходы лежат в 512-мерном подпространстве: для их хранения можно выбрать другой базис и оставить 512 координат, сохранив косинусное сходство.

Мы проверили, что из этого получается на практике. На пяти миллионах документов пик памяти процесса точного поиска снизился с 15,88 до 11,11 ГиБ, примерно на 30%. При этом полный запрос ускорился гораздо скромнее: с 36,19 до 34,80 мс в режиме одного запроса. Ниже разберём преобразование, проверим качество и посмотрим, куда уходит выигрыш во времени.

Для экспериментов арендовали стенд с RTX 5090 на gpugo.ru ^[2]. Конфигурацию выбрали в каталоге.

Почему вектор из 768 чисел можно хранить в 512 координатах

В зафиксированном checkpoint EmbeddingGemma 2 последний линейный слой текстовой части модели имеет матрицу W размера 768×512 и работает без bias. Дальше SentenceTransformers усредняет состояния токенов без padding и нормирует полученный вектор до единичной длины:

512D состояния токенов → token RMSNorm → Linear 512→768

                      → masked mean pooling → L2-нормализация

У этой схемы два полезных свойства. Во-первых, линейную проекцию можно перенести после усреднения: в точной арифметике mean(W h) = W mean(h). Такая возможность уже отмечена в реализации Transformers ^[3]. Token RMSNorm при этом остаётся на своём месте, перед усреднением: для него равенство не работает.

Во-вторых, линейный слой не может превратить 512 независимых координат в 768 независимых направлений. Проверка весов дала ранг 512 и число обусловленности около 11,64. Выходы записаны в 768 числах, но всё ещё лежат в подпространстве размерности 512.

Чтобы хранить их компактнее, используем сокращённое QR-разложение W = Q R. Матрица Q задаёт ортонормальный базис выходного подпространства, а R позволяет получить координаты в этом базисе.

W = Q R

Q: 768×512, ортонормальные столбцы

R: 512×512

полное представление: e = normalize(W x)

представление 512D:   z = normalize(R x)

Сама факторизация здесь стандартная. Нас интересовало, сохранится ли результат в реальном вычислительном пути модели и сколько памяти удастся освободить.

Какие варианты сравнивали

Новый базис можно использовать двумя способами. QR-convert преобразует уже рассчитанные Full768: тексты повторно кодировать не нужно. QR-direct сразу получает 512D-вектор, применяя R после pooling.

Для сравнения взяли и штатное уменьшение размерности. EmbeddingGemma 2 обучена с Matryoshka Representation Learning, или MRL: полезный сигнал сохраняется в первых координатах. Можно оставить первые 512 или 256 чисел штатного выхода и заново нормализовать вектор.

Полный набор вариантов выглядел так:

Pool768 нужен, чтобы отдельно проверить последствия переноса линейного слоя после pooling. В Pool768 и QR-direct этот слой считается в FP32, тогда как штатная голова работает в BF16. Поэтому меняется не только размерность, но и численный путь. Все сохраняемые векторы мы одинаково приводили к FP32 и L2-нормализовали.

У QR-convert, QR-direct и MRL512 одинаковый объём хранения. Однако MRL отбрасывает часть исходных координат, а QR меняет базис полного выходного подпространства. Штатные MRL-префиксы после такой замены не сохраняются: первые 256 координат QR-вектора нельзя считать MRL256.

Память: треть массива и 30% всего процесса

Во всех вариантах мы храним координаты в FP32, по четыре байта на число. Вычисления модели идут с BF16, но тип хранения в этом эксперименте одинаковый. Квантизацию не применяли.

Размер плотного массива можно посчитать заранее:

байты = число векторов × размерность × байты на координату

Для пяти миллионов документов получаются такие значения:

Реальный процесс занимает больше: к координатам добавляются CUDA context, служебные структуры FAISS и временная память. Поэтому мы измеряли и размер массива, и пик памяти всего процесса. Для второго использовали NVML: torch.cuda.memory_allocated() и memory_reserved() не охватывают всю память FAISS.

Для замеров взяли MS MARCO. Использовали вложенные префиксы исходного корпуса: 100 тысяч, 500 тысяч, один и пять миллионов документов. Все шесть представлений рассчитали для настоящих текстов; маленький массив до нужного размера не размножали.

Искали через GPU IndexFlatIP из FAISS ^[4]. Это точный перебор всех векторов по скалярному произведению, которое для векторов единичной длины совпадает с косинусным сходством. Каждый индекс строили в отдельном процессе, без загруженной модели кодирования текста.

На пяти миллионах документов медиана наблюдаемых пиков трёх раундов составила 15,88 ГиБ у Full768 и 11,11 ГиБ у QR-direct 512D. Массив координат уменьшился на 33,3%, а пик памяти процесса снизился на 30,0%. Служебные расходы не сокращаются пропорционально размерности.

На графике ниже видно, как менялась память с ростом корпуса:

В таблице объём указан в ГБ, то есть в 10⁹ байт. Графики используют ГиБ, то есть 2³⁰ байт: например, 15,36 ГБ равны 14,305 ГиБ. Единицы здесь нужно различать, иначе часть «экономии» получится просто из пересчёта.

Мы проверили корпуса до пяти миллионов документов. Предельную вместимость RTX 5090 и границу OOM в этом опыте ^[5] не искали.

Что произошло с качеством

В доказательстве выше мы всё время оговаривали точную арифметику. Модель работает иначе: в BF16 каждое умножение и pooling округляют значения. Штатный выход уже не обязан строго лежать в подпространстве W, поэтому QR-convert не даёт побитово точного восстановления.

На всех query-document парах SciFact и FiQA максимальная ошибка ^[6] cosine у QR-convert была около 1,5×10⁻⁵. Для QR-direct относительно Full768 она доходила до 2,3×10⁻³: здесь меняются и порядок операций, и точность головы.

Чтобы проверить влияние этих расхождений на поиск, взяли полные тестовые наборы BEIR ^[7]:

• SciFact: 5183 документа и 300 запросов;

• FiQA: 57 638 документов и 648 запросов.

Поиск был точным, по нормализованным FP32-векторам. Основная метрика nDCG@10 оценивает порядок релевантных документов в первых десяти результатах. Дополнительно считали Recall@100, долю известных релевантных документов в первой сотне выдачи.

У QR-convert значения nDCG@10 и Recall@100 совпали с Full768 по каждому из 948 запросов. Упорядоченный top-10 тоже совпал во всех случаях. В top-100 встречались различия около границы отбора: полного совпадения векторов, scores и всей выдачи мы не получили.

На графиках показана разница nDCG@10 относительно Full768. Точки обозначают оценку разницы, отрезки обозначают парные 95% доверительные интервалы. Их получили bootstrap по запросам: 10 тысяч раз пересобирали выборку с возвращением. Слева показаны все варианты, справа увеличена область около нуля. Шкалы соответствующих панелей у обоих наборов одинаковые; точные значения есть в таблицах под графиками.

У MRL512 интервалы включают ноль на обоих наборах. Этого недостаточно, чтобы объявить QR более качественным. Отдельная проверка noninferiority тоже не дала определённого ответа: мы не смогли исключить для MRL512 потерю больше 0,005 nDCG@10, поскольку нижняя граница интервала оказалась ниже выбранного порога. Само по себе это не доказывает ухудшение. У MRL256 интервалы разницы целиком лежали ниже нуля на обоих наборах.

Pool768 и QR-direct совпали друг с другом по всем per-query метрикам. Их небольшое отличие от Full768 связано с изменением численного пути, его нельзя приписать одному QR. Нулевая точка и наблюдаемый нулевой интервал QR-convert относятся к проверенным запросам; на будущие запросы эта гарантия не распространяется.

Зачем оставили контекст в 8192 токена

Для оценки качества мы сохранили документы целиком. Лимит в 256 токенов уменьшил бы расходы на кодирование, но заодно изменил бы входные тексты.

Реальные длины считали с prompts, заголовками и спецтокенами, без обрезания. У MS MARCO медиана составила 76 токенов, максимум достиг 962; за границу 256 выходят только 0,0668% passages. Для SciFact доля документов длиннее 256 токенов уже 61,14%, для FiQA она равна 19,96%. Один короткий лимит по-разному обрезал бы эти корпуса.

Основной лимит явно установили в 8192 токена, согласно карточке checkpoint ^[8]. Максимумы SciFact и FiQA составили 1901 и 3545 токенов, поэтому все документы трёх корпусов поместились целиком. При динамическом padding 8192 задаёт верхнюю границу, а не длину каждого батча. Отдельные контроли 256/512/8192 подтвердили влияние контекста на качество. Все результаты выше относятся к полным текстам.

Поиск быстрее. Полный запрос тоже, но не на треть

Замеры скорости проводили на том же арендованном стенде GPUGO ^[2], на котором проверяли расход памяти. Сначала измерили поиск по индексу, который уже находится в памяти GPU. В таймер вошли перенос векторов запросов с CPU на GPU, вызов FAISS и возврат scores и идентификаторов на CPU. Перед остановкой таймера синхронизировали устройство. Кодирование текста и построение индекса в этот замер не входили.

Для пяти миллионов документов получили:

Это p50, медианная задержка. В каждом раунде считали свою p50, затем брали медиану трёх значений. Под графиками есть и минимальные и максимальные p50 раундов. Эти диапазоны не являются доверительными интервалами. По оси X отложено число документов в линейном масштабе, линии соединяют соседние измерения. Везде, где указано B=32, приведено время целого батча.

MRL512 при тех же пяти миллионах документов дал 13,95 и 19,79 мс. По стоимости поиска два 512D-представления практически не различаются. Но даже здесь сокращение числа координат на треть не превращается в такое же снижение задержки: остаются другие операции, включая выбор top-k.

Полный путь от текста до ответа

Затем измерили запрос целиком:

текст → prompt и tokenizer → H2D → encoder → нормализация

      → GPU Flat → возврат top-10 → source IDs и JSON

Модель уже загружена, индекс уже находится на GPU. Их загрузку и построение в запрос не включали. Использовали 256 реальных запросов и пять повторов: по 1280 измерений для B=1 и 40 полных батчей для B=32 на вариант.

На одном миллионе документов переход с Full768 на QR-direct снизил p50 на 1,1% для B=1 и на 3,5% для B=32. На пяти миллионах выигрыш вырос:

На этом этапе заметна цена кодирования текста. Encoder продолжает выполнять свою работу, а сокращение индекса ускоряет только часть пути. Поэтому для одного запроса мы получили примерно 30% экономии памяти процесса поиска и 3,8% снижения полной задержки.

Полный запрос измеряли одним запуском на размер корпуса, с перемешанным порядком вариантов и повторением ^[9] запросов. Это последовательный локальный эксперимент: без сети, очередей и конкурентного обслуживания. Переносить его проценты на SLA работающего сервиса без отдельного замера нельзя.

Как перевести уже рассчитанную базу

Для готовой FP32-базы подходит QR-convert. Эмбеддинги можно читать порциями, умножать на Q, нормализовать и записывать в 512D-файл. После этого потребуется пересобрать индекс. Повторно прогонять тексты через модель не нужно.

Базис строится один раз для конкретной версии модели. Пример для загруженного SentenceTransformer:

import torch

head = st[0].auto_model.language_model.embedding_projection

assert head.bias is None

W = head.weight.detach().cpu().double()

Q, R = torch.linalg.qr(W, mode="reduced")

# Зафиксируем соглашение о знаках.

signs = torch.where(torch.diagonal(R) < 0, -1.0, 1.0)

Q = (Q * signs.unsqueeze(0)).float()

R = signs.unsqueeze(1) * R

def compress(vectors768, Q):

    assert Q.device == vectors768.device

    return torch.nn.functional.normalize(

        vectors768.float() @ Q, dim=-1

    )

На GPU матрицу Q нужно перенести заранее, один раз. Вместе с ней сохраняем revision модели и хеш весов. Документы и запросы должны использовать один базис: 512D-запрос в QR-базисе нельзя отправлять в старый 768D-индекс. При переключении сервиса нужно обновить и индекс, и преобразование запросов.

Сколько заняла конвертация пяти миллионов векторов

В отдельном CPU-замере конвертация прочитала 15,36 ГБ и записала 10,24 ГБ. Наблюдаемый пик RSS составил около 1,30 ГБ.

Всего получилось 126,78 с с учётом записи, flush и финального fsync. На умножение и нормализацию ушло 11,58 с. Основную часть, 108,95 с, заняла запись с flush. Чтение и копирование потребовали 4,32 с, ещё около 1,94 с ушло на создание выходного массива и служебную работу вне отдельных таймеров.

Входной файл заранее прочитали, поэтому он мог оставаться в файловом кэше ОС. Кэш не очищали, и этот результат не характеризует чтение с холодного NVMe. В 126,78 с также не включены prefetch (4,21 с), отдельная проверка результата (199,87 с) и пересборка индекса. Максимальная ошибка координаты относительно сохранённого GPU QR-convert составила 3,13×10⁻⁷.

Что из этого можно применить у себя

Для готовой FP32-базы мы бы начали с QR-convert. На SciFact и FiQA он сохранил nDCG@10, Recall@100 и упорядоченный top-10 по всем 948 запросам, а тексты для конвертации повторно кодировать не пришлось. На своей базе всё равно нужно проверить собственные запросы, особенно близкие оценки у границы top-k.

В новом сервисе можно проверить QR-direct, который сразу выдаёт 512D-вектор. В нашем опыте он совпал с Pool768 по per-query метрикам, но отличается от штатного численного пути. Рядом стоит оставить MRL512: у него тот же объём хранения и почти та же стоимость поиска. Превосходства QR по релевантности наши данные не доказывают.

У результата есть и границы. Мы проверили FP32-хранение текстовых эмбеддингов и плотный точный поиск. Веса encoder от смены базиса меньше не становятся, расходы на его промежуточные активации тоже остаются. Качество оценено на двух англоязычных retrieval-задачах, мультимодальный путь не измерялся. MS MARCO использовали для масштаба и скорости, без выводов о релевантности на выбранных префиксах.

Покоординатной квантизации и product quantization (PQ), которая кодирует части вектора отдельно, понадобятся собственные эксперименты: поворот базиса может изменить их поведение ^[10].

Если задача состоит в том, чтобы уместить индекс, смотрите на пик памяти всего процесса. Если нужно ускорить сервис, измеряйте полный запрос. В нашем случае 512 координат вместо 768 дали примерно 30% экономии памяти процесса поиска, а выигрыш в полной задержке зависел от размера корпуса и батча.

Подобрать GPU для своих экспериментов можно в каталоге GPUGO ^[2]. Инструкции по запуску и настройке — в документации ^[11].

Канал GPUGO в Telegram ^[12].

По вопросам сотрудничества и аренды VPS пишите на info@gpugo.ru ^[13]. Для прямой обратной связи — Telegram ^[14].