Tom Aarsen beskrev metoden på Hugging Faces blogg den 11 november och publicerade samtidigt skript för träning och utvärdering samt två exempelmodeller. Aarsen uppger att han bygger på Stefan Josefs arbete Learn to Pool och forskning från LightOn. Hans eget mål är att förklara metoden och erbjuda ett reproducerbart exempel.

En flervektormodell, till exempel ColBERT, lagrar en egen vektor för varje token i dokumentet. Vid sökning väljer varje token i frågan den vektor i dokumentet som matchar bäst, och poängen summeras. Det blir många vektorer. I Aarsens beräkning upptar de 4-bitars kvantiserade residualerna av 128-dimensionella vektorer 12,8 GB i ett dataset med en miljon dokument, med 200 vektorer per dokument. Med 50 vektorer räcker 3,2 GB.

Så fungerar pooling och träning

Indexet kan krympas genom att likartade vektorer i ett dokument slås samman (token pooling). HierarchicalTokenPooling i Sentence Transformers grupperar dem med Wards hierarkiska klustring och ersätter varje grupp med dess medelvärde. Antalet vektorer kan begränsas med en kvot (pool_factor), ett tak eller en avståndströskel. Tröskeln kan kalibreras för att ge ett önskat genomsnittligt antal vektorer. Kalibreringen siktar på medelvärdet, så antalet varierar mellan enskilda dokument, och i ett annat dataset kan även medelvärdet ändras. Som standard poolas inte frågor.

En vanlig modell har inte tränats för att fungera med poolade vektorer, så sammanslagningen kan sudda ut en skillnad som är viktig för sökningen. Vid poolingmedveten träning (pooling-aware) beräknar förlustfunktionen MultiVectorTokenPoolingLoss sökuppgiftens förlust på valda poolingnivåer. I Aarsens exempel fanns två nivåer, opoolad och poolad till ungefär en fjärdedel, och de viktades lika. Gruppindelningen deriveras inte. I stället förs gradienterna från den sammanslagna vektorn direkt till gruppens medlemmar (straight-through estimator). De poolade versionerna bildas av redan beräknade vektorer, så modellen behöver inte köras separat för dem.

Resultat från en jämförelsekörning

Aarsen tränade två modeller utifrån ModernBERT-base med samma inställningar och samma 120 000 exempel från MS MARCO under en epok. Bara förlustfunktionen skilde sig, och hyperparametrarna justerades inte. Måttet var genomsnittlig nDCG@10 på 13 dataset i NanoBEIR.

Enligt tabellen fick den poolingtränade modellen bättre resultat på alla testade poolingnivåer, och störst var skillnaden vid den snävaste gränsen.

Källa: Aarsen, Hugging Faces blogg. NanoBEIR, 13 dataset, genomsnittlig nDCG@10 (%).

Av tabellen går det inte att dra slutsatsen att fördelen består i andra dataset, eftersom modellerna bara testades på NanoBEIR. Varje modell har tränats en gång, och siffrorna är författarens egna.

Med faktor 4 behöll den poolingtränade modellen 97,8 procent och den vanliga modellen 94,6 procent av sitt eget opoolade resultat. Skillnaden i nDCG@10 var då 2,2 poäng (49,03 och 51,23), och antalet vektorer var 74,8 procent lägre än utan pooling. Fördelen syntes även med tak och kalibrerade trösklar, trots att träningen bara använde relativ pooling.

Bloggens siffror återges bland annat av zglg.work och AI Breaking Wire. Det återstår att utreda om fördelen består i betydligt större dataset och produktionsindex, skriver zglg.work.

Metoden ingår i Sentence Transformers, och pooling kan också användas med vanligt tränade flervektormodeller. Bloggen ger anvisningar för att köra träningsskriptet på en egen GPU eller på en GPU i Hugging Face Jobs. För valfri GPU-acceleration väljer biblioteket kompatibla GPU-kernlar och använder i övrigt standardimplementationen. Standardkerneln för MaxSim-poängsättning kräver till exempel CUDA och Triton.