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.
- 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:
| Precyzja | Pamięć |
|---|---|
| float32 | 4,1 GB |
| float16 | 2,0 GB |
| int8 | 1,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
- Mierz wyszukiwanie osobno od generowania: recall@k na oznaczonej próbce prawdziwych zapytań, zanim obwinisz model.
- Normalizuj wektory i kwantyzuj je, jeśli liczy się pamięć. Po kwantyzacji sprawdź recall.
- Używaj wyszukiwania hybrydowego z RRF domyślnie, a reranker dodaj wtedy, gdy liczy się precyzja.
- 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.