Apple navrhuje zrýchliť MIPS cez naučené podporné funkcie

Apple Machine Learning Research predstavil prácu o amortizovanom maximum inner product search, skrátene MIPS. Hoci názov znie akademicky, ide o jednu zo základných operácií moderných AI systémov: nájsť v databáze vektor, ktorý sa najlepšie zhoduje s dotazom. Takéto vyhľadávanie stojí za retrieval-augmented generation, odporúčacími systémami, vyhľadávaním dokumentov aj mnohými multimodálnymi aplikáciami. Ak sa dá MIPS zrýchliť bez veľkej straty kvality, dopad môže byť veľmi praktický.

Klasické vektorové vyhľadávanie používa indexy, ktoré sa snažia nájsť dobrých kandidátov bez prechodu celej databázy. Pri veľkých kolekciách dokumentov, obrázkov alebo produktov však stále vzniká kompromis medzi presnosťou, latenciou a počtom sondovaných častí indexu. Výskumníci Apple navrhujú iný pohľad: ak dotazy prichádzajú z relatívne známej distribúcie a databáza kľúčov je pevná, časť práce sa dá „amortizovať“ učením. Model sa môže naučiť predpovedať riešenia MIPS pre typické dotazy, namiesto toho, aby sa každý prípad riešil od nuly.

Kľúčová matematická myšlienka je, že hodnota MIPS sa dá chápať ako podporná funkcia množiny kľúčov. Podporná funkcia je konvexná funkcia a jej gradient ukazuje na optimálny kľúč. Z tejto interpretácie autori odvodzujú dva modely. SupportNet je input-konvexná neurónová sieť, ktorá regresne aproximuje podpornú funkciu. KeyNet je vektorová sieť, ktorá priamo predpovedá optimálny kľúč alebo reprezentáciu vhodnú na ďalšie vyhľadávanie.

Praktické použitie je zaujímavé najmä pri indexoch typu IVF, teda inverted file indexoch, ktoré rozdeľujú vektorový priestor do klastrov. SupportNet môže slúžiť ako router: namiesto slepého skúšania viacerých partícií pomáha nasmerovať dotaz do relevantnejších častí databázy. KeyNet sa dá použiť ako náhrada alebo transformácia pôvodného dotazu, ktorú následne spracuje existujúca vyhľadávacia pipeline. To je dôležité, pretože návrh nemusí nutne vyžadovať úplne nový vektorový engine.

Apple uvádza experimenty na benchmarku BEIR pre dokumentové embeddingy. Podľa zhrnutia práce naučené SupportNet a KeyNet zlepšujú IVF match rate, keď sa berie do úvahy výpočtová námaha, či už meraná vo FLOPoch, počte sond alebo reálnom čase. Inými slovami, cieľom nie je iba nájsť rovnaký výsledok ako drahšie vyhľadávanie, ale dosiahnuť lepší pomer medzi kvalitou a cenou. Presne tento pomer je pri produkčných RAG systémoch často dôležitejší než laboratórne maximum presnosti.

Výskum má však aj jasné hranice. Amortizácia dáva zmysel vtedy, keď je distribúcia dotazov aspoň čiastočne predvídateľná a databáza kľúčov sa nemení príliš rýchlo. Ak sa index neustále mení, alebo ak prichádzajú veľmi neobvyklé dotazy, naučený router môže starnúť alebo zavádzať. Produkčné nasadenie by preto potrebovalo monitoring driftu, pravidelné pretrénovanie a fallback na robustnejšie klasické vyhľadávanie. Naučený komponent by mal zrýchľovať bežné prípady, nie byť jedinou cestou k výsledku.

Pre RAG aplikácie je téma relevantná aj kvôli nákladom mimo samotného modelu. Veľa pozornosti sa venuje cene generovania odpovede, ale retrieval vrstva vie byť pri veľkých korpusoch rovnako významná. Ak systém pri každej otázke hľadá v miliónoch dokumentov, zlá konfigurácia indexu pridáva latenciu ešte predtým, než LLM začne písať. Learned routing nad MIPS je jedna z ciest, ako presunúť časť nákladov do tréningu a znížiť cenu každého ďalšieho dotazu.

Z pohľadu ekosystému je zaujímavé, že Apple prácu spája s kódom na GitHube. To umožňuje výskumníkom a infraštruktúrnym tímom overiť, či sa výsledky dajú zopakovať na ich vlastných embeddingoch a indexoch. Pri vektorovom vyhľadávaní totiž veľmi záleží na dátach: výsledky z akademického benchmarku nemusia automaticky platiť pre interné dokumenty, e-commerce katalóg alebo multimodálnu databázu. Otvorený kód uľahčuje aspoň prvotnú replikačnú kontrolu.

Najväčší význam práce je v tom, že ukazuje ďalší smer optimalizácie AI infraštruktúry. Po modeloch, ktoré sa učia odpovedať, a agentoch, ktoré sa učia používať nástroje, sa čoraz viac učia aj samotné komponenty vyhľadávacej pipeline. Ak sa takéto metódy osvedčia, budúcnosť vektorových databáz a RAG systémov nemusí byť iba o lepších indexoch písaných ručne, ale aj o indexoch doplnených malými modelmi, ktoré poznajú typické dotazy a rozloženie konkrétnej databázy.