Pięć sposobów indeksowania wektorów w RAG, przetestowanych na milionie embedów z Wikipedii

System typu RAG żyje lub umiera w zależności od możliwości wyszukiwania, a to staje się coraz trudniejsze w miarę wzrostu korpusu danych. Ten przegląd porównuje pięć sposobów indeksowania wektorów, wykorzystując rzeczywisty eksperyment, dzięki czemu wybór metody opiera się na danych liczbowych, a nie na przyzwyczajeniu.

Czym jest indeksowanie wektorów i dlaczego jest ważne

Indeksowanie wektorów to sposób organizacji dużego zbioru wektorów, który umożliwia szybkie znajdowanie podobnych wektorów, bez konieczności porównywania zapytania z każdym z nich. Bazy danych relacyjnych mają odpowiedniki w metodach bucketing i Z-ordering, które zmniejszają ilość danych, które musi zostać przeskanowana. Bazy danych wektorowych stosują specjalistyczne struktury w tym samym celu – aby zmniejszyć przestrzeń wyszukiwania.

Dokładne wyszukiwanie o niskiej opóźnieniu stanowi podstawę każdej skalowalnej architektury RAG. Już duże platformy od niej zależą: Amazon wyszukuje w katalogach zawierających miliony produktów, Netflix w tysiącach tytułów dostosowanych do globalnej publiczności, a X w miliardach postów i kont. Dane przychodzą teraz w ilościach terabajtów dziennie, więc skalowanie nie może być kwestią drugorzędną – musi być częścią architektury od samego początku projektu.

Każda z poniższych metod jest oceniana według trzech kryteriów:

• Opóźnienie: szybkość, z jaką indeks zwraca najlepsze K wyników.
• Dokładność wyszukiwania: stopień zgodności wyników z tymi uzyskanymi przy dokładnym wyszukiwaniu.

Złożoność konserwacji: stopień trudności utrzymywania indeksu w aktualnym stanie. Dane używane w aplikacjach typu RAG rzadko pozostają nieruchome i mogą wymagać odświeżania co 24 godziny, dlatego indeks musi umożliwiać łatwe aktualizacje w tym czasie bez dużego wysiłku inżynieryjnego.

Ustawienie testowe

W eksperymencie wykorzystano pierwsze milion wierszy z zestawu danych CohereLabs/wikipedia-2023-11-embed-multilingual-v3 – skala wystarczająco duża, by pokazać różnice pomiędzy metodami. Zapytania były embedowane za pomocą modelu mxbai-embed-large, który generuje wektory o wymiarze 1024 odpowiadające embedom z tego zestawu danych. Indeksy zostały utworzone w PostgreSQL przy użyciu pgvector.

Metoda 1: prosty wyszukiwarka

Poszukiwanie liniowe, nazywane również siłą brutálną, jest najprostszą opcją. Oblicza się podobieństwo pomiędzy zapytaniem a każdym wektorem, a następnie zachowuje najlepsze wyniki.

• Czas reakcji: najgorszy w grupie, ponieważ porównywany jest każdy wektor. Czas rośnie liniowo wraz z liczbą wektorów.
• Dokładność: najlepsza, ponieważ poszukiwanie jest dokładne i zwraca prawdziwe wektory najbliższe.
• Konserwacja: żadna, ponieważ nie ma indeksu.

Średni czas reakcji w teście wyniósł 6,25 sekundy. Duży sklep detaliczny odrzuciłby to natychmiast, podobnie jak mała firma start-up.

Metoda 2: IVFFlat

Inverted File Flat to technika oparta na przybliżonym najbliższym sąsiedzie. Wykorzystuje K-Means do dzielenia przestrzeni wektorowej na klastry. Zamiast skanować wszystko, znajduje środki mas najbliższe zapytaniu i przeszukuje tylko te klasy, stosując dokładne obliczenia odległości w ich obrębie.

Przepływ działań wygląda następująco:

Query → closest centroids → search selected clusters → Top-K

W pgvector kontrolę nad tym sprawiają dwa ustawienia. Liczba klastrów, nazywana listami, jest wybierana z góry, a w tym teście użyto 1000 list. Podczas wykonywania zapytania drugie ustawienie, probes, decyduje, które z tych list zostaną przeszukane. Gdy probes jest ustawione na 10, sprawdzane są te dziesięć list, których centra leżą najbliżej zapytania.

Opóźnienie było akceptowalne. Przy milionie wektorów i 1000 klastrach zapytania kończyły się w maksymalnie sekundę.

Dokładność była niska z dwóch powodów. Po pierwsze, problem granic – bliskość jest mierzona w odniesieniu do centrów, a nie do samych wektorów, więc niektóre z najbliższych wektorów pozostają niewykryte, jeśli należą do klastra, którego centrum znajduje się dalej od zapytania. Po drugie, ustawienie probes ogranicza przeszukiwanie do określonej liczby klastrów, więc istotne wektory w nieskanowanych klastrach nigdy nie są brane pod uwagę.

Zwiększenie liczby sond z 10 do 500 powodowało wzrost opóźnienia, choć w tym teście tylko nieznacznie. Nie ma jednak sposobu, aby z góry stwierdzić, czy 10, 100, 500 czy 900 sond da najlepsze wyniki, ponieważ prawdziwe sąsiednie elementy mogą znajdować się w klastrze, który został pominąty. Użycie wszystkich 1000 sond unieważnia sens tego podejścia, ponieważ zamienia indeks z powrotem w metodę siłowej przeszukiwania.

Konserwacja również była kiepska. Tworzenie indeksu dla ponad miliona wektorów zajęło około 2,37 minuty, co jest akceptowalne. Problemy pojawiają się podczas aktualizacji – indeksy oparte na środkach geometrycznych mogą wymagać ponownego obliczenia lub przeskalowania wektorów za każdym razem, gdy przychodzi nowa partia danych. Przy miliardach wektorów taki nakład inżynieryjny może nie być opłacalny.

Dla tego przypadku użycia kompromis pomiędzy opóźnieniem a dokładnością po prostu nie był akceptowalny. Następnym pytaniem jest to, czy jakaś inna metoda może zmniejszyć opóźnienie i jednocześnie zapewnić lepsze wyniki.

Metoda 3: HNSW

Dobrym przykładem może być listonosz szukający adresu. Zamiast przeszukiwać każdą ulicę, kieruje się głównymi drogami do odpowiedniej okolicy, gdzie jest niewiele połączeń, ale pokonuje duże odległości. Następnie przechodzi na główne drogi, gdzie sieć jest gęstsza, a na końcu idzie lokalnymi uliczkami, gdzie sieć jest najgęstsza, aż dotrze do właściwych drzwi.

HNSW, skrót od hierarchical navigable small world, działa w ten sam sposób. Zaczyna od rzadkich, dalekosiężnych połączeń na wyższych poziomach i przechodzi przez coraz bardziej gęste połączenia na niższych poziomach, aż dotrze do najbliższych wektorów. Jeśli hierarchia dróg jest dobrze zorganizowana, wyszukiwanie prowadzi prosto do celu, a dobra konstrukcja zapewnia dobrą dokładność wyników.

Czas reakcji był ponownie akceptowalny – maksymalnie jedna sekunda przy przetwarzaniu miliona wektorów.

Dokładność była dobra, z zastrzeżeniami wynikającymi z parametrów konfiguracji. Indeks jest tworzony za pomocą instrukcji o następującym kształcie:

    CREATE INDEX wikipedia_embeddings_hnsw_idx
    ON wikipedia_embeddings
    USING hnsw (emb vector_cosine_ops)
    WITH (
        m = 16,
        ef_construction = 64
    );

To stwierdzenie buduje indeks na podstawie korpusu. Istotne są trzy parametry:

• m to liczba sąsiadów, z którymi każdy wektor może być połączony w grafie. Określa ona połączalność i gęstość.
• ef_construction to liczba kandydatów na sąsiadów branych pod uwagę podczas budowy. Kontroluje ona przestrzeń poszukiwań podczas tworzenia indeksu.
• ef_search ogranicza przestrzeń poszukiwań podczas wyszukiwania, ponieważ określa, ilu kandydatów jest sprawdzanych podczas odpowiadania na zapytanie.

W jaki sposób graf wybiera swoje m połączeń? Dla miliona wektorów każdy z nich ma 999,999 potencjalnych sąsiadów. Narzędzie do budowy nie wybiera ich losowo – ocenia kandydatów pod kątem odległości oraz stosuje heurystykę do wyboru sąsiadów, zachowując małą grupę, około 16 elementów, które są przydatne do nawigacji po grafie.

Dlaczego to zapewnia lepszą dokładność?

1. Niewielka liczba przydatnych sąsiadów na wektor stanowi „autostrady” grafu.
2. Wybierane są te połączenia, aby graf pozostał połączony i łatwy do nawigacji.
3. Podczas wyszukiwania algorytm podąża za nimi krok po kroku w kierunku wektorów, które najbardziej przypominają zapytanie.
4. Ponieważ graf oferuje efektywne ścieżki do odpowiedniej strefy, istnieje większe prawdopodobieństwo dotarcia do prawdziwych najbliższych sąsiadów.
5. Dlatego staranna konstrukcja i nawigacja zwiększają dokładność, unikając jednocześnie pełnej porównywania wszystkich możliwości.

Krótko mówiąc, graf jest budowany celowo, a nie losowo. Ustawienie wartości m i ef_construction określa sposób jego budowy, co z kolei wpływa na skuteczność wyszukiwania.

Konserwacja to obszar, w którym HNSW sprawia problemy. Stworzenie indeksu dla ponad miliona wektorów zajęło zaskakujące 37 minut. Podczas aktualizacji, jeśli początkowy zbiór danych jest duży i dobrze odzwierciedla ogólną rozkładność danych, nowe dane mogą być dodawane bez konieczności ponownego budowania indeksu. Jednak w miarę wzrostu ilości nowych danych w stosunku do oryginalnego zbioru rozkładność może ulec zmianie, graf staje się mniej optymalny, a w końcu konieczne staje się ponowne utworzenie indeksu. Zdarza się to rzadziej niż w przypadku IVF, ale faktycznie ma miejsce.

Ponieważ graf opiera się na zasadniczej strukturze, stanowi lepszą opcję niż IVF pomimo dodatkowej pracy inżynieryjnej. Równowaga pomiędzy opóźnieniem a dokładnością jest akceptowalna, choć nie wybitna.

Metoda 4: BM25 plus wyszukiwanie wektorowe

Czasami oczywista odpowiedź jest właściwa. Podczas gdy uwaga skupia się na coraz bardziej skomplikowanych metodach, od trzydziestu lat istnieje prosty rozwiązanie, które pozostaje niezauważone. BM25 to algorytm rankingu leksykalnego, który ocenia dokumenty na podstawie częstotliwości występowania terminów z zapytania, rzadkości tych terminów we całym korpusie oraz długości dokumentu.

Zasady działania:

1. Dokumenty są dzielone na terminy, a odwrotny indeks łączy każdy termin z dokumentami, które go zawierają.
2. Częstotliwość terminu mierzy, jak często termin ten pojawia się w danym dokumencie.
3. Odwrotna częstotliwość dokumentu mierzy, na jak rzadki lub powszechny termin jest w całej kolekcji.
4. Termin, który jest częsty w jednym dokumencie, ale rzadki ogólnie, otrzymuje wyższą wagę.
5. BM25 łączy te czynniki z długością dokumentu, aby ocenić jego istotność.
6. Podczas wyszukiwania terminy z zapytania są sprawdzane w odwróconym indeksie, oceniane w odniesieniu do pasujących dokumentów, a te o najwyższej ocenie są zwracane.

Sama formuła nie jest zbyt elegancka, ale idea jest prosta: znaleźć dokumenty najbardziej istotne dla słów zawartych w zapytaniu.

Gdy te dokumenty zostaną odzyskane, obszar wyszukiwania znacznie się zmniejsza. Wtedy prosty wyszukiwacz wektorowy w tym małym zbiorze znajduje najlepsze wyniki. To właśnie jest hybrydowe odzyskiwanie informacji.

Jedna uwaga: pierwszy etap szuka dokładnych słów kluczowych, więc błąd ortograficzny użytkownika może mieć negatywny wpływ. Krok normalizacji wspomagany przez sztuczną inteligencję, który poprawia błędy i warianty przed rozpoczęciem wyszukiwania, łatwo radzi sobie z tym problemem.

Opóźnienie. Część wektorowa odpowiada w oczekiwanych milisekundach, natomiast część leksykalna w tym ustawieniu zajmowała od 12 do 14 sekund. Cache może to rozwiązać, choć pierwsza prośba nadal będzie wolna. Cache zasługuje tu na uwagę, ponieważ szczególnie dobrze sprawdza się w przypadku BM25 – dokumenty są wyszukiwane poprzez dopasowanie słów kluczowych. W przypadku innych metod indeksowania cache musi być semantyczny, co wymaga progu podobieństwa i może kosztować nieco dokładności wyszukiwania.

Konserwacja. Brak, ponieważ nie tworzy się indeksu wektorowego.

W praktyce, dla większości produktów AI z maksymalnie około 50 000 dokumentami, hybrydowe wyszukiwanie z użyciem BM25 jest stałym wyborem, oferującym dobrą równowagę pomiędzy opóźnieniem a dokładnością na tym poziomie.

Metoda 5: DiskANN i graf Vamana

Ta metoda jest bardziej abstrakcyjna, podobnie jak same embeddingi. Proces składa się z ośmiu kroków:

1. Zacznij od dużej kolekcji wektorów, które wymagają efektywnego wyszukiwania.
2. Stwórz graf, łącząc każdy wektor z starannie wybranym zestawem sąsiadów. To jest graf Vamana.
3. Zapisz graf na dysku, a nie w pamięci, dzięki czemu można obsługiwać zbiory danych znacznie większe niż RAM.
4. Gdy przyjdzie zapytanie, rozpocznij od punktu wejścia w grafie.
5. Podążaj za połączeniami między wektorami, krok po kroku zmierzając w stronę obszarów bliskich zapytaniu.
6. W miarę postępu wyszukiwania ładuj potrzebne węzły z SSD.
7. Po zbadaniu wystarczającej części grafu zebraj mały zestaw obiecujących kandydatów.
8. Porównaj te kandydaty dokładniej i zwróć K najbliższych.

Kluczowym elementem jest to, że graf nawigacyjny znajduje się na SSD, a nie w RAM, co sprawia, że ograniczeniem jest raczej rozmiar niż koszt pamięci.

Najpierw filtrowaj metadane

Jedna technika zasługuje na szczególne wyróżnienie i powinna być wymieniona przed wszystkimi powyższymi metodami. Filtrowanie metadanych opiera się na tej samej zasadzie co grupowanie w bazie danych relacyjnej: przed przeprowadzeniem jakiegokolwiek wyszukiwania wektorowego zawęża się zbiór kandydatów za pomocą atrybutów takich jak kategoria, data czy właściciel.

Którą metodę wybrać

• Siła bruta: małe do średnich zbiory danych, do około 20 000 dokumentów. Najlepsza dokładność przy akceptowalnej opóźnieniu.
• BM25 w połączeniu z wyszukiwaniem wektorowym: średnie zbiory danych, około 20 000 do 100 000 dokumentów. Dobry balans między wyszukiwaniem leksykalnym a semantycznym.
• HNSW: duże zbiory danych zawierające miliony wektorów lub więcej. Dobra szybkość i dokładność, ale przy wyższej złożoności implementacji.
• DiskANN i Vamana: gdy pamięć RAM staje się droga. SSD oferuje bardzo dużą pojemność, więc ograniczeniem jest ilość pamięci, a nie jej rozmiar.
• IVFFlat: duże zbiory danych, liczące od jednego do dziesięciu milionów wektorów, gdy dane są dobrze sklastrowane. W przeciwnym razie pogarsza się dokładność wyników.