Rustban írták újra a cross-encoder inferenciát, de nem mindenhol lett gyorsabb

Új Rust-könyvtár készült a rangsoroláshoz

A Qdrant szeptember 25-i bejegyzése szerint a vállalat elkészítette a cross-encode-rs nevű Rust-könyvtárat. A cél a cross-encoder modellek következtetésének, vagyis az inferenciának a futtatása Rust környezetben.

A cross-encoder egy lekérdezést és egy dokumentumot egyszerre olvas be, majd egyetlen relevanciapontszámot ad a párhoz. Ezeket a modelleket jellemzően újrarangsorolásra használják. Egy gyors első lépés kiválasztja a lehetséges találatokat, a cross-encoder pedig újrarendezi őket, hogy a legjobb egyezések kerüljenek előre.

A módszer pontosabb lehet, mivel a modell a lekérdezés és a dokumentum minden szavát együtt hasonlítja össze. Ennek ára, hogy minden lekérdezés és dokumentum párosa teljes modellfuttatást igényel. Az embedding, vagyis bi-encoder modellek ezzel szemben külön vektorokká alakítják a lekérdezést és a dokumentumokat.

A Rust nem váltotta ki az alapul szolgáló motort

A Qdrant a Rust-könyvtárat a Python ökoszisztéma két elterjedt megoldásával, a fastembed és a sentence-transformers könyvtárral vetette össze. Mindhárom megoldás ugyanazt a C++-ban készült ONNX Runtime motort használja, ezért a futási idő nagy része közös, lefordított kódban telik.

Az ONNX egy betanított modellekhez használt fájlformátum. A modellek PyTorchból vagy TensorFlow-ból exportálhatók, az ONNX Runtime pedig több programozási nyelven, köztük Rustban is futtatja őket. A Qdrant ehhez az ort nevű Rust-csomagot használta.

A könyvtár fő optimalizációja a kötegelés. A dokumentumokat nem egyenként, egymás után kódolja, hanem együtt dolgozza fel. A bemeneteket tokenizált hosszuk szerint rendezi, és hasonló méretű dokumentumokat futtat egy kötegben, így kevesebb számítás megy el a kitöltő tokenekre. A fejlesztők a tokenizers v1 verzióját is használták, amely a bejegyzés szerint mérsékelt gyorsulást ad a 0.x verziókhoz képest.

A nagyobb modellen látszik igazán a különbség

A méréseket a Qdrant az mteb/scidocs-reranking adathalmaz tesztkészletén végezte, amely körülbelül 3,98 ezer lekérdezést tartalmaz. Lekérdezésenként átlagosan 30 lekérdezés és dokumentum párt futtattak két modellen, egy MacBook M4 Max számítógépen, 48 GB memóriával és mind a 14 CPU-mag használatával.

A kisebb, Xenova/ms-marco-MiniLM-L-6-v2 modellen a három könyvtár teljesítménye közel állt egymáshoz. A cross-encode-rs mediánban körülbelül 1,1-szer volt gyorsabb, a p99 értéknél pedig 1,4-szeres előnyt ért el. A leggyorsabb kéréseknél a Python-alapú megoldások enyhén jobbak voltak. A p99 mérésben egy kérés feldolgozása 36 ezredmásodpercet vett igénybe a cross-encode-rs, 49-et a fastembed és 51-et a sentence-transformers esetében.

A nagyobb, többnyelvű jinaai/jina-reranker-v2-base-multilingual modellen a Rust-könyvtár minden percentilisben 1,3 és 1,5 közötti gyorsulást mutatott a két Python-könyvtárhoz képest, az egyetlen leglassabb kérés kivételével. A medián késleltetés 131 ezredmásodperc volt a cross-encode-rs, 189 a fastembed és 191 a sentence-transformers esetében. A leglassabb kérésnél a fastembed 480, a cross-encode-rs 581, a sentence-transformers pedig 585 ezredmásodpercet ért el.

A modell betöltése továbbra sem lett gyorsabb

A Qdrant külön megvizsgálta a modell betöltését is, mivel ez fontos lehet hideg induláskor, automatikus skálázásnál és szerver nélküli környezetben. A mérés csak az ONNX inferenciamunkamenet létrehozására koncentrált, a Python indítását, a könyvtárak importálását és a tokenizáló betöltését kizárták.

A MiniLM esetében a fastembed átlagosan 33,5 ezredmásodperc alatt hozta létre a munkamenetet, ami 4,1 ezredmásodperccel, vagyis 12 százalékkal gyorsabb volt a cross-encode-rs megoldásánál. A Jina modellnél a két könyvtár gyakorlatilag döntetlen eredményt hozott, átlagosan 1 százalék volt köztük a különbség.

A teljes folyamatot mérve a fastembed mindkét modellen körülbelül 134 ezredmásodperccel tovább futott. A Qdrant szerint ez a Python értelmező indításához és a könyvtár importálásához kapcsolódó többlet, nem maga a fastembed modellfuttatása. A mérés így azt mutatja, hogy a Rust előnye elsősorban bizonyos inferenciahelyzetekben jelentkezik, a közös ONNX Runtime miatt pedig a modell betöltésénél kevésbé tud eltérni a két megoldás.