Tom Aarsen kuvasi menetelmän Hugging Facen blogissa 11. marraskuuta ja julkaisi samalla koulutus- ja arviointiskriptit sekä kaksi esimerkkimallia. Aarsen kertoo nojaavansa Stefan Josefin Learn to Pool -työhön ja LightOnin tutkimukseen. Hänen oma tavoitteensa on selittää menetelmä ja tarjota toistettava esimerkki.

Monivektorimalli, kuten ColBERT, tallentaa asiakirjan jokaisesta tokenista oman vektorin. Haussa kukin kyselyn token valitsee asiakirjasta parhaiten vastaavan vektorin, ja pisteet lasketaan yhteen. Vektoreita kertyy paljon: Aarsenin laskelmassa 128-ulotteisten vektorien 4-bittisiksi kvantisoidut jäännösosat vievät miljoonan asiakirjan aineistossa 12,8 Gt, kun asiakirjaa kohden on 200 vektoria. 50 vektorilla niihin riittää 3,2 Gt.

Miten tiivistys ja koulutus toimivat

Indeksiä voi pienentää yhdistämällä asiakirjan samankaltaisia vektoreita (token pooling). Sentence Transformersin HierarchicalTokenPooling ryhmittelee ne Wardin hierarkkisella klusteroinnilla ja korvaa kunkin ryhmän sen keskiarvolla. Vektorimäärää voi rajoittaa suhteella (pool_factor), ylärajalla tai etäisyyskynnyksellä. Kynnyksen voi kalibroida tuottamaan tavoitellun keskimääräisen vektorimäärän. Kalibrointi tähtää keskiarvoon, joten yksittäisen asiakirjan vektorimäärä vaihtelee, ja toisessa aineistossa keskiarvokin voi muuttua. Oletuksena kyselyjä ei tiivistetä.

Tavallista mallia ei ole koulutettu toimimaan tiivistetyillä vektoreilla, joten yhdistäminen voi hävittää haun kannalta tärkeän eron. Tiivistykseen valmistavassa koulutuksessa (pooling-aware) häviöfunktio MultiVectorTokenPoolingLoss laskee hakutehtävän häviön valituilla tiivistystasoilla. Aarsenin esimerkissä tasoja oli kaksi, tiivistämätön ja noin neljäsosaan tiivistetty, ja ne painotettiin tasan. Ryhmäjakoa ei derivoida, vaan gradientit kulkevat yhdistetyltä vektorilta suoraan ryhmän jäsenille (straight-through estimator). Tiivistetyt versiot muodostetaan jo lasketuista vektoreista, joten mallia ei tarvitse ajaa niitä varten erikseen.

Tulokset yhdestä vertailuajosta

Aarsen koulutti ModernBERT-basesta kaksi mallia samoilla asetuksilla ja samoilla 120 000 MS MARCO -esimerkillä yhden epookin ajan. Vain häviöfunktio erosi, eikä hyperparametreja viritetty. Mittarina oli keskimääräinen nDCG@10 NanoBEIRin 13 aineistolla.

Taulukon mukaan tiivistykseen koulutettu malli sai paremman tuloksen kaikilla kokeilluilla tiivistystasoilla, ja suurin ero syntyi tiukimmalla rajalla.

Lähde: Aarsen, Hugging Face -blogi. NanoBEIR, 13 aineistoa, keskimääräinen nDCG@10 (%).

Taulukosta ei voi päätellä, että etu säilyisi muilla aineistoilla, koska mallit testattiin vain NanoBEIRillä. Kummastakin mallista on yksi koulutusajo, ja luvut ovat kirjoittajan omia.

Kertoimella 4 tiivistettynä tiivistykseen koulutettu malli säilytti 97,8 % ja tavallinen malli 94,6 % omasta tiivistämättömästä tuloksestaan. Tällöin nDCG@10-tulosten ero oli 2,2 pistettä (49,03 ja 51,23), ja vektoreita oli 74,8 % vähemmän kuin tiivistämättä. Etu näkyi myös ylärajoilla ja kalibroiduilla kynnyksillä, vaikka koulutuksessa käytettiin vain suhteellista tiivistystä.

Muun muassa zglg.work ja AI Breaking Wire toistavat blogin luvut. zglg.work:n mukaan on vielä selvittämättä, säilyykö etu paljon suuremmissa aineistoissa ja tuotantoindekseissä.

Menetelmä kuuluu Sentence Transformers -kirjastoon, ja tiivistystä voi käyttää myös tavallisesti koulutetuissa monivektorimalleissa. Blogi neuvoo ajamaan koulutusskriptin omalla tai Hugging Face Jobs -palvelun GPU:lla. Valinnaista GPU-kiihdytystä varten kirjasto valitsee yhteensopivat GPU-kernelit ja käyttää muualla tavallista toteutusta; esimerkiksi MaxSim-pisteytyksen oletuskerneli vaatii CUDA:n ja Tritonin.