A multi-vektoros beágyazó modellek minden lekérdezés-token számára külön lehetőséget adnak arra, hogy dokumentumban találjanak releváns egyezést. Ez rugalmas, de a dokumentumokra eső vektormennyiség gyorsan nő: milliós korpuszok esetén százmilliós nagyságrendű vektorok gyűlhetnek össze. A token pooling a hasonló tokenvektorok összevonásával csökkenti ezt a tárolási igényt, de túl agresszív tömörítésnél fontos információ veszhet el. A pooling-aware (poololás-tudatos) tréning célja, hogy az enkódert arra tanítsa, mely információkat kell megőriznie a vektorok összevonása után, ezáltal jobb lekérdezési pontosságot érve el ugyanazon tárhely mellett.
Ebben a cikkben bemutatom a pooling működését, egy Sentence Transformers alapú példát a pooling-aware finomhangolásra, és a minőség–tárhely kompromisszum mérőszámait. A szerző kísérleteiben pool_factor=4 mellett a hagyományos modell a NanoBEIR nDCG@10 értékének 94.6%-át tartja meg a poololás után, míg a pooling-aware tréning ezt 97.8%-ra növeli ugyanazon 4x-es vektorcsökkentés mellett.
Miért kell sok vektor multi-vektoros lekérdezéshez
Egy single-vector modell egy dokumentumot egyetlen vektorral reprezentál. Ezzel szemben a multi-vektoros megközelítés (például ColBERT-stílusú modellek) minden tokenhez kontextuális vektort rendel, így egy lekérdezés különböző részei külön dokumentumrészletekhez illeszthetők. MaxSim pontozásnál minden lekérdezésvektor a legjobban illeszkedő dokumentumvektort választja, és ezek pontszámai összeadódnak (MaxSim(Q,D) = sum_{q in Q} max_{d in D} q^T d).
Még kvantálással (például Fast-PLAID 2- vagy 4-bites maradékkvantálása) is nagy a költség: például 128 dimenzió, 4 bit per koordináta, és 200 vektor dokumentumonként ~12.8 KB-ot jelent dokumentumonként, ami egy millió dokumentumnál ~12.8 GB. Ha dokumentumonként 50 vektorra poololunk, ez 3.2 GB-ra csökken.
Kulcsmegfigyelés: sok dokumentumvektor hasonló lehet, mert több token közel kapcsolódó fogalmakat ír le — ezt használja ki a token pooling.
Token pooling: HierarchicalTokenPooling és működése
A HierarchicalTokenPooling Ward hierarchikus klaszterezést használ a hasonló token-beágyazások összegyűjtésére. Kezdetben minden token saját csoport, majd ismétlődően azok a csoportok egyesülnek, amelyek a belső szórás minimális növekedését okozzák. A belső szórás mérőszáma az egyes vektorok és a csoportátlag négyzetes távolságainak összege, így Ward tömör, hasonló vektorokból álló csoportokat preferál.
A hierarchia lehetővé teszi egy tetszőleges tömörítési szint kiválasztását; a csoportok átlagával reprezentáljuk a csoportot, majd normalizáljuk őket MaxSim-hez. A pool_factor például arányos tömörítést jelent:
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
pool_factor=4 esetén körülbelül az eredeti vektorok negyedét tartjuk meg, így hosszabb dokumentumok arányosan több kapacitást őriznek meg. Alapértelmezésben az első token védett, és a poololás csak a dokumentum-beágyazásokat érinti, a lekérdezésvektorok változatlanok maradnak.
Pooling-költségvezérlők
A HierarchicalTokenPooling három vezérlési módot kínál:
- max_tokens=16: dokumentumonként legfeljebb 16 vektor (védett tokenek is beleszámítanak).
- pool_factor=4: dokumentumonként körülbelül 1/4 marad.
- distance_threshold=...: Ward egyesítési távolságküszöb, amely dokumentumonként változó vektorszámot eredményezhet.
A választás fontos, ha a dokumentumhossz változó: pool_factor hosszabb dokumentumoknak arányosan több vektort enged meg, míg distance_threshold akkor növeli a vektorszámot, ha több különböző beágyazást kell megőrizni.
Adaptív pooling kalibrálása
Az adaptív pooling egy közös Ward-távolságküszöböt használ, amely agresszívebben egyesít hasonló beágyazásokat, de megtart több vektort azoknak a dokumentumoknak, amelyekben sok megkülönböztethető tartalom van. Bármely multi-vektoros modellhez alkalmazható, akár poololás-tudatos cél nélkül tanult modellhez is.
A küszöb kézi beállítása helyett kalibrálhatjuk azt egy reprezentatív dokumentumsampelre, hogy egy cél-átlagvektorszámot érjünk el. A kalibrációhoz elegendő a dokumentumbeágyazások mintája; a visszaadott küszöb újrahasználható indexeléskor. Az elért átlag korpuszonként változhat.
Poololás-tudatos tréning elve
A probléma akkor jelentkezik, ha egy dokumentumban van egy lekérdezés szempontjából fontos megkülönböztetés, amely poololáskor eltűnik. Az eredeti tanítási cél nem kényszerítette a modellt, hogy jól működjön tömörített reprezentációkkal. A pooling-aware tréning ezt a hiányt pótolja: a tanítási célot kiterjeszti a tömörített dokumentumokra is.
A megközelítés hasonlít a Matryoshka embeddingekre: több reprezentációszélességen adunk lekérdezési célokat. A bemutatott példában minden tréningbatch két veszteséget ad hozzá:
- Lekérdezési veszteség az eredeti (unpooled) dokumentumbeágyazásokkal.
- Lekérdezési veszteség dokumentumok poololása után (kb. 1/4 vektor).
A végső veszteség az átlaguk: L = 1/2 L_unpooled + 1/2 L_pooled. Az első tag megtanítja a modellt teljes reprezentációval működni, a második pedig arra, hogy a tokenbeágyazásokat olyan módon rendezze el, hogy a poololás után is megmaradjanak a hasznos egyezések. Mindkét cél ugyanazt az enkódert frissíti.
Mivel a Ward-klaszterezés diszkrét döntéseket hoz, a HierarchicalTokenPooling tréning közben straight-through becslőt használ, hogy a poololt vektorokra visszafelé a tagtokenekhez gradienseket továbbítson anélkül, hogy explicit differenciálna a klaszter-assignációkon. Így az enkóder megtanulhat a poololt reprezentációk hibáiból.
Implementáció Sentence Transformers-ben
A MultiVectorTokenPoolingLoss egy wrapper, amely meglévő multi-vektoros lekérdezési veszteséget alkalmaz a kiválasztott poolerekkel. Példa:
from sentence_transformers import MultiVectorEncoder from sentence_transformers.multi_vector_encoder.losses import ( CachedMultiVectorMultipleNegativesRankingLoss, MultiVectorTokenPoolingLoss, ) from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
model = MultiVectorEncoder("answerdotai/ModernBERT-base", similarity_fn_name="maxsim") pooling = HierarchicalTokenPooling(pool_factor=4)
base_loss = CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=64) loss = MultiVectorTokenPoolingLoss( model, base_loss, token_poolings=[None, pooling], weights=[0.5, 0.5], )
CachedMultiVectorMultipleNegativesRankingLoss a lekérdezésekre vonatkozó rangsorolási célt szolgáltatja, a gradient caching pedig lehetővé teszi kis mini-batch feldolgozást, miközben nagyobb in-batch negatív készletet tartunk meg (például mini-batch 64, tényleges tréningbatch 256 esetén). A poolerek újrafelhasználják a tokenbeágyazásokat, így minden további pooler klaszterezési és pontozási lépést ad hozzá, de nem igényel külön enkóder-futást.
A poolerek helyébe bármilyen BaseTokenPooling-implementáció is behelyettesíthető, amely képes gradienseket továbbadni.
Gyorsítás (kernelek)
A MaxSim pontozás, Ward klaszterezés és egyes alap-háttérműveletek felgyorsíthatók kernel-implementációkkal (CUDA/Triton). Telepítési javaslat:
pip install -U "sentence-transformers[train,kernels]"
A példából a model inicializációjánál use_kernels=True beállítással a kompatibilis kernel-ek automatikusan kiválasztódnak (például egy default MaxSim kernel, ami memóriahatékonyan számol és Tritont/CUDA-t igényel, valamint egy gyors Ward kernel, amely 30–150x gyorsulást hozhat SciPy-hoz képest a beállástól függően).
A betanított modell használata
Betölthetjük a pooling-aware modellt és indexelhetünk poololt dokumentumvektorokat. A kernel-beállítások nincsenek elmentve a modellhez, ezért ha szeretnénk gyors kernel-eket inference alatt is, adjuk meg use_kernels=True-t a betöltéskor.
Példa lekérdezés és indexelés:
model = MultiVectorEncoder("tomaarsen/mve-ModernBERT-base-token-pooling") pooling = HierarchicalTokenPooling(pool_factor=4)
queries = ["Which planet is known as the Red Planet?"] documents = [ "Venus is often called Earth's twin because of its similar size and proximity.", "Mars, known for its reddish appearance, is often referred to as the Red Planet.", "Jupiter, the largest planet in our solar system, has a prominent red spot.", ]
query_embeddings = model.encode_query(queries) document_embeddings = model.encode_document(documents, token_pooling=pooling) print([len(embedding) for embedding in document_embeddings])
[5, 6, 5] (poololva)
A poololás az enkóder pipeline utáni tömörítést végzi: a példában az összes dokumentumvektor száma 58-ról 16-ra csökkent, miközben a lekérdezési rangsor megtartotta Mars-t legjobb találatként.
A pool_factor inference során is módosítható a tárhely-költség felfedezéséhez, vagy a poolert runtime alapértelmezetté tehetjük a MultiVectorEncoder inicializációjánál.
Értékelési eredmények
A szerző két modellt tanított az answerdotai/ModernBERT-base kiinduló modellből, mindkettőn friss 128-dimenziós tokenprojekcióval:
- tomaarsen/mve-ModernBERT-base-token-pooling-baseline — csak a cached contrastive veszteséggel.
- tomaarsen/mve-ModernBERT-base-token-pooling — a cached contrastive veszteség csomagolva az unpooled és factor-4 poolozott célokkal.
Mindkét futás egy epoch MS MARCO 120,000 tripletjén történt, seed=42, training token cap 300, és a rövid lekérdezések legalább 32 tokenre bővültek. Csak a veszteség volt eltérő; nem futott hyperparaméter-keresés.
Értékelés NanoBEIR@13-on, mean nDCG@10 (%) szerint. A táblázat több beállítást hasonlít össze (pool_factor 2,3,4,6; max_tokens és adaptív distance_threshold). Fontosabb eredmények:
- pool_factor=4 mellett a pooling-aware modell nDCG@10 49.03-ról 51.23-ra javult.
- Ugyanezen a beállításon a pooling-aware modell 74.8%-kal kevesebb dokumentumvektort tárol az unpooled állapothoz képest.
- A score retention (poololt nDCG@10 osztva az adott modell unpooled nDCG@10-ével) pool_factor=4 esetén a hagyományos modellnél 94.6%, pooling-aware esetén 97.8%.
- Az unpooled teljesítmény is javult: 51.84-ről 52.39-re nőtt.
A javulás kiterjed fix caps és adaptív pooling esetére is, még akkor is, ha a tréning csak teljes és arányos poololást használt. Például adaptív poolingkal, ~29.58 vektor/dokumentum átlag mellett a pooling-aware modell 97.3%-át tartja meg az unpooled pontszámának, miközben 84.7%-kal kevesebb vektort tárol. A hagyományos modell ugyanazon a költségen 93.3%-ot tartott meg.
A táblázat részleteiben kisebb és nagyobb pool_factoroknál, illetve max_tokens korlátoknál is konstans javulás látható a pooling-aware modell javára, különösen erős előnyökkel nagyon alacsony vektorszámoknál (például max_tokens=4 esetén a gain 5.78 pont volt a kísérletben).
További index- és modelkompressziós lehetőségek
Ha az index mérete a legfontosabb szempont, a szerző javasolt sorrendje:
- Embedding kvantálás (például Fast-PLAID 4/2/1 bitre) a koordinátákon, a tárolási igény csökkentésére.
- Token pooling a dokumentumvektorok számának csökkentésére, opcionálisan pooling-aware finomhangolással.
- Matryoshka token embeddingek a tokenvektorok dimenziójának csökkentésére, opcionálisan finomhangolással, hogy kisebb szélességen is megmaradjon a lekérdezési minőség.
Ezek kombinálhatók: például MultiVectorMatryoshkaLoss-szal dimenzióredukcióra és MultiVectorTokenPoolingLoss-szal a poololásra is együttesen lehet tanítani.
Összegzés
Néhány extra sor a tréningcélban lehetővé teszi, hogy a modell aktívan megtanulja használni a kisebb dokumentumreprezentációt. A pooling-aware finomhangolás jelentősen javítja a lekérdezési pontosság megőrzését ugyanazon tárhely-büdzsében, így több lehetőség nyílik arra, hogy a multi-vektoros megközelítést gyakorlati, skálázható indexekben alkalmazzuk.
Referenciák és források
A bejegyzés hivatkozik több munkára, köztük Stefan Josef "Learn to Pool" és LightOn pooling regularization tanulmányaira (Benjamin Clavié, Antoine Chaffin, Griffin Adams), valamint a Matryoshka embeddingek dokumentációjára. A kód és példák elérhetők a Sentence Transformers repositoryjában, és a cikk mellé több előre betanított checkpoint is csatolva (például tomaarsen/mve-ModernBERT-base-token-pooling és ...-baseline).
