Matryoshka embedding

Matryoshka embedding er en metode for å trene en embeddingmodell til å “strukturere” embedding dimensjonene på en slik måte at vi kan fjerne dimensjoner uten å tape informasjon ¹. Dette gjør at vi kan redusere dimensjonaliteten til en embeddingmodell uten å tape for mye ytelse.

Ved å redusere dimensjonaliteten til en embeddingmodell kan vi redusere behovet for lagringsplass i en vektordatabase. Dette kan være nyttig for å kunne utføre et semantisksøk hurtigere eller rett og slett spare penger på nødvendig lagringsplass.

Mange embeddingmodeller støtter i dag Matryoshka embedding, noe som gjør at det er greit å kjenne til hvordan man kan finjustere en slik embeddingmodell.

Viktig

Vi kommer til å gjenbruke data fra Grunnleggende og kommer ikke til å gjengi hvordan data er strukturert eller sammenstilt her.

Hvis du ønsker å forstå hvordan vi har sammenstilt treningsdataene og hva som er innholdet, anbefaler vi at du gjør det først.

Nødvendige avhengigheter

Vi kommer til å bruke de samme pakkene som vi brukte i Grunnleggende.

Vi har følgende i pyprojects.toml:

dependencies = [
    "datasets>=3.3.2",
    "polars[pyarrow]>=1.23.0",
    "rich>=13.9.4",
    "sentence-transformers[train]>=3.4.1",
    "transformers[torch]>=4.49.0",
]

Samt flash-attn for CUDA-akselerasjon:

uv add flash-attn --no-build-isolation

Laste inn treningsdata

Vi begynner med å laste inn treningsdata før vi gjør klar modellen og gjør endringene som trengs for Matryoshka embedding.

import polars as pl

try:
    from rich import print
except ModuleNotFoundError:
    pass

# Les inn datasett
test_dataset = pl.read_parquet("test_dataset.parquet")
train_dataset = pl.read_parquet("train_dataset.parquet")
# Kombiner for å kunne arbeide med hele datasettet
dataset = pl.concat([test_dataset, train_dataset])
# Skriv ut raske tall
print(f"Antall elementer [bold magenta]totalt[/]:\t{len(dataset)}")
print(f"Antall elementer i [bold green]trening[/]:\t{len(train_dataset)}")
print(f"Antall elementer i [bold blue]test[/]:\t{len(test_dataset)}")

Vi lager oss deretter et corpus og et sett med relevante dokumenter for mulige søk.

# Merk at vi bruker `dataset` for å bruke _alt_ innhold
corpus = dict(dataset.select(["id", "positive"]).rows())
# For "queries" bruker vi det vi har plukket ut i test
queries = dict(test_dataset.select(["id", "anchor"]).rows())

relevant_docs = {}
for qid in queries.keys():
    # Hvert "spørsmål" vil være knyttet til en tittel fra Nav.no, vi henter ut
    # denne tittelen og henter alle rader i datasettet med samme tittel som
    # relevant dokument
    q_pos = (
        dataset.filter(pl.col("id") == qid)
        .unique("positive")
        .item(row=0, column="positive")
    )
    relevant_docs[qid] = set([qid])
    relevant_docs[qid].update(
        set(dataset.filter(pl.col("positive") == q_pos).get_column("id"))
    )

Embeddingmodell og Matryoshka evaluering

Det neste vi gjør er å definere embeddingmodell på samme måte som vi gjorde i Grunnleggende.

import torch
from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "Alibaba-NLP/gte-modernbert-base",
    # NOTE: Vi velger `device` tilpasset CUDA, for Mac kan man bruke `mps` og
    # `cpu` vil alltid være tilgjengelig
    device="cuda" if torch.cuda.is_available() else "mps",
    model_kwargs=dict(
        attn_implementation="flash_attention_2"
        if torch.cuda.is_available()
        else "sdpa",
    ),
    tokenizer_kwargs=dict(padding="max_length", truncation=True),
)

Først nå vil koden endre seg fra tidligere. Når vi skal definere hvordan modellen skal evalueres så må vi få med reduksjon av dimensjonene som en del av evalueringen.

Vi starter med å definere dimensjonene vi ønsker å benytte med modellen. Her er det viktig at dimensjonene er strukturert fra størst til minst.

matryoshka_dimensions: list[int] = [768, 512, 256, 128, 64]

from sentence_transformers.evaluation import (
    InformationRetrievalEvaluator,
    SequentialEvaluator,
)
from sentence_transformers.util import cos_sim

# Opprett liste med evalueringer per dimensjon
sub_evaluators = []
for dim in matryoshka_dimensions:
    evaluator = InformationRetrievalEvaluator(
        queries=queries,
        corpus=corpus,
        relevant_docs=relevant_docs,
        name=f"dim_{dim}",
        truncate_dim=dim,
        score_functions={"cosine": cos_sim},
        batch_size=64
        if torch.cuda.is_available()
        else 4,  # Skru denne ned eller opp avhengig av tilgjengelig minne, høyere gir raskere evaluering
    )
    sub_evaluators.append(evaluator)
# Vi lager så en sekvensiel evaluator som kjører alle evalueringene etter
# hverandre
evaluator = SequentialEvaluator(sub_evaluators)

base_results = evaluator(model)

for dim in matryoshka_dimensions:
    key = f"dim_{dim}_cosine_ndcg@10"
    print(f"NDCG@10 for [bold magenta]{dim}[/] dimensjoner:\t{base_results[key]}")

Treningsmetode (loss function)

For at vi skal kunne finjustere modellen må vi definere treningsmetode og oppsett spesielt tilpasset matryoshka embedding.

from sentence_transformers.losses import MultipleNegativesRankingLoss, MatryoshkaLoss

# Først definerer vi hvordan hver dimensjon skal rangeres
inner_train_loss = MultipleNegativesRankingLoss(model=model)
# Før vi slår disse sammen med en "meta"-trener som trener én modell med flere
# dimensjoner
train_loss = MatryoshkaLoss(
    model, loss=inner_train_loss, matryoshka_dims=matryoshka_dimensions
)

Treningsoppsett

Vi definerer deretter resten av treningsoppsettet som vi gjorde i Grunnleggende.

from sentence_transformers import SentenceTransformerTrainingArguments
from sentence_transformers.training_args import BatchSamplers

# Definer hvordan trening skal foregå
train_args = SentenceTransformerTrainingArguments(
    output_dir="gte-modernbert-navno-matryoshka",
    num_train_epochs=4,  # Antall epoker å trene, flere er bedre
    per_device_train_batch_size=128,  # Bestemt av maskinvare, høyere trener raskere
    per_device_eval_batch_size=32,
    warmup_ratio=0.1,
    learning_rate=2e-5,
    lr_scheduler_type="cosine",
    optim="adamw_torch_fused",
    tf32=True,  # Kjekt å sette til `True` hvis maskinvare støtter (krever nyere Nvidia GPU)
    fp16=False,  # Sett til `True` hvis man ikke kan bruke `bf16`
    bf16=True,  # Kjekt å sette på hvis maskinvare støtter (støttes av Mac og Nvidia GPU-er)
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # Veldig praktisk å fjerne duplikater når man har Positiv Pair
    eval_on_start=True,
    eval_strategy="epoch",  # Evaluer etter hver X steg
    save_strategy="epoch",  # Lagre modell etter X steg
    logging_steps=50,
    save_total_limit=3,  # Bare spar på de 3 siste modellene
    load_best_model_at_end=True,
    # NOTE: Vi optimaliserer hele modellen for best mulig NDCG@10 med 128
    # dimensjoner, dette er en endring fra hvordan vi gjorde det i
    # "Grunnlegende"
    metric_for_best_model="dim_128_cosine_ndcg@10",
)

Vi laster deretter inn treningsdataene i datasets før vi kan utføre selve treningen.

from datasets import Dataset

train_ds = Dataset.from_polars(train_dataset.select(["anchor", "positive"]))
train_ds

from sentence_transformers import SentenceTransformerTrainer

trainer = SentenceTransformerTrainer(
    model=model,
    args=train_args,
    train_dataset=train_ds,
    loss=train_loss,
    evaluator=evaluator,
)

Utføre finjustering

Tilslutt er det bare å utføre finjusteringen av modellen.

# Utfør trening
trainer.train()

# Pass på at vi lagrer modellen
trainer.save_model()

Vi evaluerer deretter en siste gang for å sammenligne hvordan ytelsen er for hver dimensjon sammenlignet med før finjustering.

final_eval = evaluator(model)

from rich.table import Table
from rich.console import Console

table = Table(title="Sammenligning av NDCG@10")
table.add_column("Før/etter finjustering")
for dim in matryoshka_dimensions:
    table.add_column(f"{dim}")
table.add_row(
    "Før",
    *[str(base_results[f"dim_{dim}_cosine_ndcg@10"]) for dim in matryoshka_dimensions],
)
table.add_row(
    "Etter",
    *[str(final_eval[f"dim_{dim}_cosine_ndcg@10"]) for dim in matryoshka_dimensions],
)

console = Console()
console.print(table)

Resultat av finjustering

Før/etter finjustering	768	512	256	128	64
Før	0.224	0.210	0.195	0.175	0.151
Etter	0.377	0.369	0.368	0.356	0.305

Som vi kan se fra tabellen over så gjør finjusteringen vår at selv så lite som 64 dimensjoner kan gi en god ytelse (bedre enn den originale modellen ved 768 dimensjoner)!

Fotnoter

1. Litt ytelse blir borte, men det skal ikke være katastrofalt som hvis man reduserer dimensjoner i en modell som ikke er trent på denne måten