Efficiënte pre-training met token-superpositie

Samenvatting

Pre-training van grote taalmodellen is vaak buitensporig duur en inefficiënt op schaal, en vereist complexe en ingrijpende aanpassingen om een hoge doorvoersnelheid van data te bereiken. In dit werk presenteren we Token-Superposition Training (TST), een eenvoudige, inplugbare methode die de datadoorvoer per FLOP tijdens pre-training aanzienlijk verbetert zonder de parallellisatie, optimizer, tokenizer, data of modelarchitectuur te wijzigen. TST vindt plaats in twee fasen: (i) een zeer efficiënte superpositiefase waarin we veel aaneengesloten tokens combineren tot één bundel en trainen met een multi-hot kruisentropie-doelstelling (MCE), en (ii) een herstelfase waarin we terugkeren naar standaard training. We evalueren TST uitgebreid op schaal van 270M en 600M parameters en valideren op 3B en een 10B A1B mixture of experts-model, wat aantoont dat de methode zeer robuust is in verschillende instellingen. Uiteindelijk presteert TST consequent beter dan de baseline-verliesfunctie en downstream-evaluaties, en onder gelijkwaardige verliesinstellingen levert TST een reductie van maximaal 2,5x in de totale pre-trainingstijd op de 10B A1B-schaal.

English

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.