Два метода оптимизировать embedding models

Рассмотрим два принципиально разных метода оптимизировать модели для эмбеддингов: Matryoshka Representation Learning (blogpost) и Binary and Scalar Embedding Quantization.

Что такое эмбеддинги и зачем они нужны

Эмбеддинги — ключевой инструмент в задачах машинного обучения, превращающие сложные объекты (текст, изображения, аудио) в численные векторы. Они позволяют оценить сходство объектов, что важно для систем рекомендаций, поиска, обнаружения аномалий и других задач.

Почему эмбеддинги нужно оптимизировать

Многие sota модели выдают векторные представления размерностью 1024, каждое число кодируется в формате float32, т.е., они требуют 4 байта на размерность. Для выполнения поиска по 250 миллионам векторов, нам, таким образом, понадобится около 1ТБ памяти!

Matryoshka Representation Learning (MRL)

Хотим иметь модель для эмбеддингов, которая может выдавать вектора разных размерностей, от 1024 до 64. Причем такие вектора должны иметь смысл.

Авторы статьи показали, что можно обучить модель так, чтобы она складывала наиболее важную информацию из эмбеддингов в начало вектора, а потом можно просто отбрасывать хвост и использовать эффективный размер вектора.

То есть во время обучения применяем лосс не только к полноразмерному эмбеддингу, но еще и к обрезанным частям. Заметного оверхэда по скорости обучения и памяти не возникает.

Например, для nomic-ai/nomic-embed-text-v1.5 получается сохранить 95.8% перфоманса при уменьшении эмбеддингов в 3 раза и 90% при 6 кратном уменьшении.

MRL в коде

Обучаем

from sentence_transformers import SentenceTransformer
from sentence_transformers.losses import CoSENTLoss, MatryoshkaLoss

model = SentenceTransformer("microsoft/mpnet-base")

base_loss = CoSENTLoss(model=model)
loss = MatryoshkaLoss(
    model=model,
    loss=base_loss,
    matryoshka_dims=[768, 512, 256, 128, 64],
    matryoshka_weight=[1, 1, 1, 1, 1],
)

model.fit(
    train_objectives=[(train_dataset, loss)],
    ...,
)

Используем

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim

model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka")

matryoshka_dim = 64
embeddings = model.encode(
    [
        "The weather is so nice!",
        "It's so sunny outside!",
        "He drove to the stadium.",
    ]
)
embeddings = embeddings[..., :matryoshka_dim]  # Shrink the embedding dimensions
print(embeddings.shape)
# => (3, 64)

# Similarity of the first sentence to the other two:
similarities = cos_sim(embeddings[0], embeddings[1:])
print(similarities)
# => tensor([[0.8910, 0.1337]])

Binary and Scalar Embedding Quantization

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

Binary quantization

Каждое число в нормализованных эмбеддингах отоброжаем с помощью простой пороговой функции в 0 или 1:

Для вычисления сходства используем расстояние Хэмминга - эффективный способ посчитать количество различных элементов в строке.

Получается, что каждое измерения эмбеддинга закодируется 1 битом, но на практике принято работать с байтами, поэтому биты закодируем в байты и получим из 1024 бит 128 байт, если изначально имеем эмбеддинг размерности 1024.

Многие вектор базы данных вроде faiss поддерживают бинарные вектора.

Scalar quantization

Процесс квантования эмбеддингов из float32 в int8.

Это включает в себя отображение непрерывного диапазона значений float32 в дискретный набор значений int8, которые могут представлять 256 различных значений (от -128 до 127). Для этого используется большой калибровочный набор данных (эмбеддингов). Мы вычисляем диапазон этих эмбеддингов, т.е. min и max каждого измерения эмбеддинга. Отсюда мы рассчитываем шаги (бакеты) для категоризации каждого значения.

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

Квантизация эмбеддинга размерности 1024 к uint8 или int8 приведет к 1024 байтам.

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

Никто также не запрещает пробовать объединять MRL и квантизацию.

Квантизация в коде

Бинарная:

from sentence_transformers import SentenceTransformer
from sentence_transformers.quantization import quantize_embeddings

# 1. Load an embedding model
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# 2b. or, encode some text without quantization & apply quantization afterwards
embeddings = model.encode(["I am driving to the lake.", "It is a beautiful day."])
binary_embeddings = quantize_embeddings(embeddings, precision="binary")

Скалярная:

from sentence_transformers import SentenceTransformer
from sentence_transformers.quantization import quantize_embeddings
from datasets import load_dataset

# 1. Load an embedding model
model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")

# 2. Prepare an example calibration dataset
corpus = load_dataset("nq_open", split="train[:1000]")["question"]
calibration_embeddings = model.encode(corpus)

# 3. Encode some text without quantization & apply quantization afterwards
embeddings = model.encode(["I am driving to the lake.", "It is a beautiful day."])
int8_embeddings = quantize_embeddings(
    embeddings,
    precision="int8",
    calibration_embeddings=calibration_embeddings,
)