Tekstgeneratie voorbij discrete token sampling

Text Generation Beyond Discrete Token Sampling

May 20, 2025

Auteurs: Yufan Zhuang, Liyuan Liu, Chandan Singh, Jingbo Shang, Jianfeng Gao

cs.AI

Samenvatting

In standaard autoregressieve generatie voorspelt een LLM de volgende-tokenverdeling, samplet een discreet token, en verwijdert vervolgens de verdeling, waarbij alleen het gesampelde token als nieuwe invoer wordt doorgegeven. Om de rijke informatie van deze verdeling te behouden, stellen we Mixture of Inputs (MoI) voor, een trainingsvrije methode voor autoregressieve generatie. Na het genereren van een token volgens het standaardparadigma, construeren we een nieuwe invoer die het gegenereerde discrete token combineert met de eerder verwijderde tokenverdeling. Specifiek gebruiken we een Bayesiaanse schattingsmethode die de tokenverdeling behandelt als de prior, het gesampelde token als de observatie, en de conventionele one-hot vector vervangt door de continue posterior verwachting als de nieuwe modelinvoer. MoI stelt het model in staat om een rijkere interne representatie te behouden gedurende het generatieproces, wat resulteert in verbeterde tekstkwaliteit en redeneervaardigheden. Op het gebied van wiskundig redeneren, codegeneratie en PhD-niveau QA-taken verbetert MoI consistent de prestaties van meerdere modellen, waaronder QwQ-32B, Nemotron-Super-49B, Gemma-3-27B en DAPO-Qwen-32B, zonder extra training en met verwaarloosbare rekenkundige overhead.

English

In standard autoregressive generation, an LLM predicts the next-token distribution, samples a discrete token, and then discards the distribution, passing only the sampled token as new input. To preserve this distribution's rich information, we propose Mixture of Inputs (MoI), a training-free method for autoregressive generation. After generating a token following the standard paradigm, we construct a new input that blends the generated discrete token with the previously discarded token distribution. Specifically, we employ a Bayesian estimation method that treats the token distribution as the prior, the sampled token as the observation, and replaces the conventional one-hot vector with the continuous posterior expectation as the new model input. MoI allows the model to maintain a richer internal representation throughout the generation process, resulting in improved text quality and reasoning capabilities. On mathematical reasoning, code generation, and PhD-level QA tasks, MoI consistently improves performance across multiple models including QwQ-32B, Nemotron-Super-49B, Gemma-3-27B, and DAPO-Qwen-32B, with no additional training and negligible computational overhead.