In-Place Tokenizer Expansion for Pre-trained LLMs

Technische Samenvatting: In-Place Tokenizer Expansie voor Pre-trained LLM's

Probleemstelling
Tokenizers voor Large Language Models (LLM's) zijn doorgaans vastgesteld aan het begin van de pre-training, waarbij de vocabulaire-allocatie de distributie van de initiële corpus reflecteert. Wanneer de prioriteiten bij implementatie verschuiven naar ondervertegenwoordigde talen, lijden deze talen onder een hoge "token-fertility"—ze worden in aanzienlijk meer tokens gesplitst per woord vergeleken met de talen waarvoor de tokenizer geoptimaliseerd was. Deze fragmentatie verhoogt de decode-latentie, rekenkracht en het energieverbruik, met name voor on-device modellen waarbij de embedding- en LM-head-matrices een substantieel deel van de per-token decode-bandbreedte uitmaken. Terwijl cloud-modellen brede vocabularia kunnen veroorloven, moeten edge-modellen worden geleverd met compacte vocabularia, waarbij fragmentatie buiten een vaste taalset wordt geaccepteerd. Standaardoplossingen, zoals pre-training vanaf nul met een nieuwe tokenizer of het gebruik van zero-shot transfer methoden naar willekeurige derde-partij tokenizers, zijn ofwel computationeel verspillend of lossen een moeilijker probleem op dan noodzakelijk is.

Methodologie
De auteurs stellen Tokenizer Expansion voor, een in-place recept voor het upgraden van de tokenizer van een pre-trained model wanneer de producent van het model het ontwerp controleert. De aanpak bestaat uit drie hoofdonderdelen:

1. Tokenizer Constructie (Continued-BPE):
   In plaats van een nieuwe tokenizer onafhankelijk te trainen of een unie van vocabularia te nemen, zetten de auteurs de Byte-Pair Encoding (BPE) merge-training van de bron-tokenizer voort op een meertalige corpus die zwaarder gewogen is naar onder-getokeniseerde talen.

   • Eigenschappen: Bestaande merges blijven vastgezet; nieuwe merges worden toegevoegd. Dit zorgt voor een deterministische mapping waarbij de meeste bron-tokens als enkele tokens behouden blijven, en elke nieuwe token een exacte decompositie heeft in één of meer bron-tokens.
   • Schaal: Toegepast op een 8B-parameter Mixture-of-Experts (MoE) model, waarbij de vocabulaire werd uitgebreid van 65K naar 128K actieve entries.

2. Embedding Initialisatie:
   De embedding-matrix wordt opnieuw geïnitialiseerd om overeen te komen met de nieuwe vocabulairegrootte ($|V_{new}|$).

   • Overgenomen tokens: Rijen voor tokens die 1-op-1 van de bron afkomstig zijn, worden ongewijzigd gekopieerd.
   • Nieuwe tokens: Rijen voor nieuwe tokens worden geïnitialiseerd als het gemiddelde van de embedding-rijen van hun constituerende bron-subtokens. Dit zorgt ervoor dat de nieuwe rijen een norm hebben die vergelijkbaar is met de bestaande rijen, zonder dat willekeurige initialisatie nodig is.

3. Twee-fasen Adaptatie:
   Om de kwaliteit van de bron-checkpoint te herstellen, ondergaat het model een specifieke trainingspipeline:

   • Fase 1 (Alleen-Embedding Training): Het modellichaam is bevroren. Alleen de nieuwe embedding-rijen (en hun gekoppelde LM-head tegenhangers) worden getraind op 600B tokens van een meertalige mix gericht op onder-getokeniseerde talen. Cruciaal is dat de rijen die van de bron zijn gekopieerd, bevroren blijven om te voorkomen dat ze afwijken van het bevroren lichaam.
   • Fase 2 (Volledige Continued Pre-Training): Alle parameters worden ontdooid. Het model vervolgt de pre-training voor 400B tokens op een gebalanceerde meertalige corpus. Deze fase is essentieel voor het herbalanceren van het model en het integreren van de nieuwe vocabulaire zonder de prestaties op voorheen ondersteunde talen te laten degraderen.
   • Mid-Training/Post-Training: De aangepaste checkpoint gaat verder met standaard mid-training (context extensie naar 128K) en post-training fasen.

Belangrijkste Resultaten
De methode werd toegepast op de LFM2-8B-A1B checkpoint om LFM2.5-8B-A1B te produceren.

• Token Efficiëntie: De uitgebreide tokenizer vermindert de token-aantallen voor onder-getokeniseerde talen aanzienlijk, terwijl de pariteit voor Engels en code behouden blijft.
  • Reducties: Hindi (2,4× minder tokens), Vietnamees (2,6×), en Thais (~4,0×).
• Kwaliteitsherstel:
  • Zero-shot Swap: Het simpelweg wisselen van de tokenizer veroorzaakte een daling van ~5,8 geaggregeerde benchmarkpunten.
  • Fase 1: Herstelde ~4,8 punten, wat grotendeels de prestaties voor Engels en code herstelde.
  • Fase 2: Dicht de resterende kloof en overtrof de bron-model met 3,6 geaggregeerde punten, waarbij het de bron evenaarde of overtrof op de categorieën Kennis, Wiskunde, Code en Meertalig.
  • Meertalige Benchmarks: Voorheen ondersteunde talen (bijv. Arabisch, Duits, Japans) behielden of verbeterden hun prestaties licht, terwijl onder-getokeniseerde talen substantiële winsten zagen (bijv. Vietnamees +11,6 punten, Hindi +7,6 punten op Global-MMLU).
• On-Device Latentie:
  • Trade-off: Een grotere vocabulaire verhoogt de per-token decode-kosten (gemeten als een ~7–10% daling in tokens/sec op referentie-apparaten door geheugenbandbreedte-beperkingen).
  • Netto Versnelling: Door de combinatie van token-compressie met de per-token kosten, levert de uitgebreide tokenizer een 2,2× tot 3,7× per-karakter decode-versnelling op voor Thais, Bengalees, Vietnamees en Hindi op referentie-apparaten (M4 Max, Snapdragon 8 Elite).
  • Regressie: Voor talen met bijna-pariteit compressie (Engels, code, belangrijke Europese talen), regresseert de per-karakter decode-tijd met tot wel ~9%.

Negatieve Bevindingen en Ablaties
Het paper rapporteert twee kritieke negatieve bevindingen die de uiteindelijke recept vormden:

1. Het trainen van de volledige embedding in Fase 1: Het trainen van alle embedding-rijen (inclus inclusief gekopieerde bron-rijen) terwijl het lichaam bevroren was, veroorzaakte een degradatie in generatieve capaciteiten (bijv. instorting van MATH500, HumanEval+) en ernstige repetitie, zelfs toen de multiple-choice accuratesse intact bleef. Het beperken van de training tot alleen de nieuwe rijen was noodzakelijk.
2. Fase 2 Data-mix: Het gebruik van een overwegend Engelse mix voor de Fase 2 continued pre-training veroorzaakte een "generatieve instorting" (verlies van redeneren en meertalige generatie), ondanks het behouden of verbeteren van de multiple-choice scores. Een gebalanceerde meertalige mix was vereist om de generatieve kwaliteit te bewaren.

Betekenis en Claims
De auteurs claimen dat dit de eerste account is van een model-zijde recept voor in-place, continued-BPE tokenizer expansie op productieschaal (specifiek voor een 8B on-device MoE).

• Efficiëntie: Het demonstreert dat vocabulaire-expansie kan worden bereikt zonder de geaccumuleerde pre-training rekenkracht weg te gooien, mits de tokenizer wordt geconstrueerd als een voortzetting van de bron-merges.
• Deployment: Het biedt een praktische oplossing voor on-device modellen om extra talen efficiënt te ondersteunen, waarbij een kleine uniforme per-token kosten wordt geruild voor enorme efficiëntiewinsten op specifieke schriften.
• Evaluatie Inzicht: Het werk benadrukt dat standaard multiple-choice benchmarks (zoals MMLU) falen kunnen in het detecteren van degradatie in generatieve capaciteiten tijdens vocabulaire-adaptatie, wat de inclusie van generatieve benchmarks in evaluatie noodzakelijk maakt.

De auteurs publiceren de modelgewichten (LFM2.5-8B-A1B) en de uitgebreide tokenizer, waarbij zij opmerken dat het recept specifief is ontworpen voor scenario's waar de producent van het model de tokenizer-upgrade controleert, in plaats van het adapteren aan willekeurige derde-partij tokenizers.