Wielowektorowe modele embeddingowe, takie jak ColBERT, oferują wysoką precyzję wyszukiwania, ale liczba wektorów przechowywanych na dokument szybko zwiększa koszt indeksu. Technika grupowania, token pooling, redukuje ten rozmiar, a odpowiedni trening pozwala, według źródła, lepiej zachować jakość wyników przy równym budżecie przechowywania. Przy pool_factor=4, normalnie wytrenowany model zachowuje 94,6% swojego nieskompresowanego nDCG@10 na NanoBEIR, w porównaniu do 97,8% dla modelu wytrenowanego z uwzględnieniem poolingu.
- Modele wielowektorowe, takie jak ColBERT, zapewniają wysoką precyzję, ale zwiększają rozmiar indeksu.
- Token pooling zmniejsza liczbę wektorów, a trening z jego uwzględnieniem lepiej zachowuje jakość wyników.
- Przy pool_factor=4 model wytrenowany z poolingiem zachowuje 97,8% wyniku nieskompresowanego, wobec 94,6% zwykłego modelu.
- Źródło zaleca łączyć kwantyzację, token pooling i embeddingi Matryoshka, aby zmniejszyć indeks.
Dlaczego modele wielowektorowe przechowują tak wiele wektorów
Klasyczny model embeddingowy podsumowuje cały dokument w jednym wektorze. Modele wielowektorowe zachowują natomiast zestaw wektorów kontekstowych, jeden na token (jednostkę tekstu, często słowo lub część słowa), co pozwala różnym częściom zapytania dopasować się do różnych części dokumentu. Źródło podaje przykład zapytania o laptop o długiej żywotności baterii, które może znaleźć fragment o urządzeniu w jednej części dokumentu i fragment o żywotności baterii w innej.
Takie działanie ma swoją cenę. Według przedstawionego obliczenia, przy kwantyzacji 4-bitowej na współrzędną rezydualną, 200 wektorów o 128 wymiarach reprezentuje 12,8 KB na dokument, czyli 12,8 GB na milion dokumentów. Sprowadzenie każdego dokumentu do 50 wektorów redukuje tę sumę do 3,2 GB, przy identycznych wymiarach i precyzji.
Metoda opiera się na obserwacji: wiele wektorów z tego samego dokumentu może być bliskich sobie, ponieważ wiele tokenów często opisuje podobne pojęcia. HierarchicalTokenPooling, biblioteka Sentence Transformers, grupuje te wektory za pomocą clusteringu Warda, który na każdym etapie łączy dwie grupy powodujące najmniejszy wzrost rozproszenia wewnętrznego. Każda grupa jest następnie zastępowana swoją średnią, a potem normalizowana. Przy pool_factor=4, dokument zachowuje około jednej czwartej swoich początkowych wektorów, tak że dłuższy dokument zachowuje proporcjonalnie więcej wektorów.
Trenowanie modelu pod kątem kompresji
Im silniejsza kompresja, tym łatwiej stracić użyteczną informację. Trening uwzględniający pooling ma na celu właśnie zachowanie ważnej informacji, gdy wektory są łączone. Każda partia treningowa łączy dwie straty: wyszukiwanie z oryginalnymi embeddingami dokumentów oraz wyszukiwanie po kompresji do około jednej czwartej wektorów. Te dwie straty są uśredniane, a oba cele aktualizują ten sam enkoder.
Ponieważ clustering podejmuje dyskretne decyzje, źródło wskazuje, że trening wykorzystuje estymator straight-through, który przekazuje gradienty od zgrupowanych wektorów do tokenów, które je tworzą, bez różnicowania przypisania do klastrów. Klasa MultiVectorTokenPoolingLoss enkapsuluje istniejącą stratę wyszukiwania, tutaj CachedMultiVectorMultipleNegativesRankingLoss, z wagami 0,5 dla każdej wariantu. Modelem wyjściowym jest answerdotai/ModernBERT-base, wytrenowany na 120 000 tripletów MS MARCO przez jedną epokę.
Wyniki na NanoBEIR
Przy pool_factor=4, nDCG@10 wzrasta z 49,03 do 51,23 dla modelu wytrenowanego z poolingiem, jednocześnie redukując o 74,8% liczbę wektorów dokumentów w porównaniu do wyszukiwania nieskompresowanego. Model nieskompresowany również się poprawia, z 51,84 do 52,39. Dla średniej 29,58 wektorów na dokument, adaptacyjny pooling zachowuje 97,3% nieskompresowanego wyniku modelu wytrenowanego, w porównaniu do 93,3% dla normalnie wytrenowanego modelu.
Aby zmniejszyć rozmiar indeksu, źródło sugeruje najpierw kwantyzację embeddingów, następnie token pooling, ewentualnie z dostrojeniem uwzględniającym pooling, a na końcu embeddingi Matryoshka w celu redukcji wymiaru każdego wektora. Te podejścia można łączyć. Prace, na których opiera się artykuł, opisano w badaniu Learn to Pool autorstwa Stefana Josefa.
