Trefon.
← Wszystkie notatki 10 min

Matematyka wyszukiwania w RAG: embeddingi, dzielenie na fragmenty, wyszukiwanie hybrydowe i reranking

Wyszukiwanie decyduje o tym, co model dostaje do przeczytania, i to tam RAG zawodzi najczęściej. Podobieństwo cosinusowe i co daje normalizacja, ile pamięci kosztuje milion fragmentów przy każdej precyzji, jak HNSW wymienia recall na szybkość, dlaczego wygrywa hybrydowe połączenie BM25 z wyszukiwaniem gęstym i kiedy reranker jest wart swojego kosztu.

Jacek Trefon
  • deep dive
  • RAG
  • embeddingi
  • wyszukiwanie
  • wyszukiwarka

W systemie RAG model może użyć tylko tego, co poda mu wyszukiwanie. Gdy odpowiedź jest błędna, przyczyna leży zwykle przed modelem: właściwy fragment tekstu w ogóle nie został wyszukany albo został wyszukany, ale utonął wśród innych. Poniżej matematyka połowy odpowiedzialnej za wyszukiwanie, czyli tej, w której jest większość problemów możliwych do naprawienia.

Podobieństwo: cosinus i dlaczego normalizować

Model embeddingów zamienia każdy fragment (chunk) tekstu na wektor. Wyszukiwanie znajduje fragmenty, których wektory są najbliżej wektora zapytania. Podobieństwo cosinusowe porównuje tylko kierunek:

cos(q, d) = (q · d) / (‖q‖ · ‖d‖)

Znormalizuj każdy wektor do długości jednostkowej raz, przy indeksowaniu, a cosinus stanie się zwykłym iloczynem skalarnym. To tańsze, a właśnie pod to budowane są szybkie indeksy wektorowe. Większość modeli embeddingów trenuje się z podobieństwem cosinusowym, więc normalizacja nie kosztuje nic, jeśli chodzi o jakość.

Pamięć: ile kosztuje milion fragmentów

Przy 1024 wymiarach milion fragmentów zajmuje:

PrecyzjaPamięć
float324,1 GB
float162,0 GB
int81,0 GB
binarna (1 bit na wymiar)128 MB

Kwantyzacja wektorów do int8 zwykle tylko nieznacznie obniża recall. Wektory binarne tracą więcej, ale są 32× mniejsze i bardzo szybko się je porównuje. Częsty wzorzec to najpierw przeszukać indeks binarny, a potem ponownie ocenić najlepszych kandydatów wektorami w pełnej precyzji. W ten sposób wyszukiwanie na urządzeniu, gdzie ograniczeniem jest pamięć, pozostaje szybkie.

Wyszukiwanie przybliżone: HNSW

Porównanie zapytania z każdym wektorem jest dokładne, ale liniowe względem rozmiaru kolekcji. HNSW (Malkov i Yashunin, 2018) buduje warstwowy graf bliskości. Wyszukiwanie zaczyna się w rzadkiej górnej warstwie i zachłannie zmierza w stronę zapytania, schodząc warstwa po warstwie. Jego koszt rośnie mniej więcej logarytmicznie, a nie liniowo.

Pokrętła wymieniają recall na szybkość: ilu sąsiadów trzyma każdy węzeł i jak szeroka jest lista kandydatów podczas wyszukiwania. Mierz recall@k, czyli odsetek prawdziwych k najbliższych sąsiadów, które zwraca indeks, w porównaniu z wyszukiwaniem dokładnym na próbce własnych zapytań. „Szybkie” wyszukiwanie z recall 80% oznacza, że jeden istotny fragment na pięć po cichu znika.

Dzielenie na fragmenty

Rozmiar fragmentu to prawdziwy kompromis:

  • Małe fragmenty zamykają w embeddingu po jednej myśli i są wyszukiwane precyzyjnie, ale tracą otaczający kontekst, więc model dostaje zdanie bez akapitu, od którego ono zależy.
  • Duże fragmenty zachowują kontekst, ale ich embedding uśrednia kilka myśli i nie pasuje wyraźnie do żadnej z nich.

Dziel na fragmenty według struktury tam, gdzie ona istnieje: nagłówków, funkcji, sekcji. Niech sąsiednie fragmenty trochę na siebie zachodzą, żeby myśli rozcięte granicą pojawiały się w całości w jednym fragmencie. Przy każdym fragmencie zapisz tytuł albo ścieżkę sekcji nadrzędnej, żeby model wiedział, skąd pochodzi dany fragment.

Wyszukiwanie hybrydowe wygrywa

Gęste embeddingi oddają znaczenie, ale słabo radzą sobie z dokładnymi tokenami: kodami błędów, nazwami funkcji, SKU produktów. BM25, klasyczny ranker dopasowujący terminy, jest ich przeciwieństwem. Uruchom oba i połącz wyniki metodą reciprocal rank fusion (RRF) (Cormack i in., 2009):

RRF(d) = Σ po rankerach  1 / (k + rank_r(d))        k = 60

RRF nie potrzebuje kalibracji wyników między dwoma systemami, tylko pozycji w rankingu. Dokumenty, które oba rankery stawiają wysoko, trafiają na samą górę. Wyszukiwanie hybrydowe to zwykle najtańsza dostępna poprawa jakości w systemie RAG.

Reranking

Pierwszy etap szybko zwraca, powiedzmy, 50 kandydatów. Reranker typu cross-encoder czyta potem zapytanie razem z każdym kandydatem i ocenia taką parę łącznie. To znacznie dokładniejsze niż porównywanie dwóch niezależnie policzonych wektorów i znacznie droższe. Reranking 50 kandydatów po około 512 tokenów oznacza przetworzenie około 25 600 tokenów na zapytanie.

Układ dwuetapowy, czyli tani recall, a potem droga precyzja, jest standardem, bo żaden etap sam nie jest wystarczająco dobry.

Potem pomyśl o kolejności

Wyszukane fragmenty to kontekst, a kontekst ma pozycje. Liu i in. (2023) pokazali, że modele lepiej wykorzystują informacje z początku i końca długiego kontekstu niż ze środka. Najmocniejsze dowody umieść na początku i nie wypełniaj kontekstu dwudziestoma słabymi fragmentami, gdy wystarczą trzy mocne.

Zasady, których się trzymam

  1. Mierz wyszukiwanie osobno od generowania: recall@k na oznaczonej próbce prawdziwych zapytań, zanim obwinisz model.
  2. Normalizuj wektory i kwantyzuj je, jeśli liczy się pamięć. Po kwantyzacji sprawdź recall.
  3. Używaj wyszukiwania hybrydowego z RRF domyślnie, a reranker dodaj wtedy, gdy liczy się precyzja.
  4. Dziel na fragmenty według struktury, z metadanymi, a zbiór wyszukanych fragmentów trzymaj mały i mocny.

Źródła: Malkov i Yashunin, HNSW (2018). Robertson i Zaragoza, The Probabilistic Relevance Framework: BM25 and Beyond (2009). Cormack, Clarke i Büttcher, Reciprocal Rank Fusion (SIGIR 2009). Liu i in., Lost in the Middle (2023). Zajętość pamięci policzono dla miliona wektorów o 1024 wymiarach.

Masz podobny problem?

Opisz swój problem

30-minutowa rozmowa techniczna z inżynierem, który to buduje. Bez działu sprzedaży, bez slajdów.