Organizations: Faculty of Information Engineering and Automation, Kunming University of Science and Technology, Kunming, China · Key Laboratory of Artificial Intelligence in Yunnan Province, Kunming University of Science and Technology, Kunming, China
Diffusion and flow models have recently been extended to language modeling, generating text through continuous denoising in token embedding spaces. Pretrained models provide token representations and contextual prediction capabilities, but improving generation with limited additional training requires an appropriate connection between supervision and continuous generation. Continuous MSE evaluates a predicted vector obtained by weighting token embeddings with predicted probabilities; different probability distributions can yield the same vector and are therefore indistinguishable under this objective. We propose Token Identity Brier Flow (TIBF), which applies squared supervision directly to token identities when continuing pretrained language flows. TIBF combines network scores with token-matching scores under a Gaussian observation model, supervises token probabilities using the Brier loss, and uses probability-weighted embeddings for continuous generation. With fixed token embeddings, we show that the squared error between predicted token probabilities and the true posterior bounds the squared errors of continuous denoising predictions and flow velocities. A risk decomposition separates uncertainty inherent in the observation from posterior fitting error, motivating continued training that preserves pretrained embeddings and focuses on probability prediction. We also establish convergence to token embeddings under ideal predictions and sufficient conditions for local convergence with prediction error. From common pretrained initializations and under matched continuation budgets, frozen-embedding TIBF achieves Gen.PPL values of 33.96 on OWT and 81.17 on LM1B, outperforming continuous-MSE and CE continuation controls. It attains the lowest Gen.PPL at every tested sampling budget on both datasets, together with the lowest repetition rate on LM1B.
Figures & tables
Figure 1: The same continuous prediction can conceal different identity judgments. Four embedding vectors lie in two dimensions. Predictions A and B are (0.4,0.1,0.4,0.1) and (0.1,0.4,0.1,0.4) , with token 1 as the true identity. Both probability-weighted vectors are zero, giving normalized continuous MSE 1 under Eq. ( 9 ), but half-Brier losses 0.27 and 0.57. Positions are shown as ev/d , with ∥ev∥2=d . This analytical example illustrates the many-to-one readout.
Figure 2: Generation quality and diversity under the common sampling settings. Left: OWT; right: LM1B. The top row plots Gen.PPL against token entropy; the bottom row plots Gen.PPL against repeated 4-gram fraction. Points correspond to the models in Tables 5 and 6 .
Figure 3: Empirical trajectory convergence of frozen-embedding TIBF. (a) Noise-normalized distances to the nearest and second-nearest embeddings. Solid curves show midpoints of the 1%–99% quantile intervals; shaded regions show the intervals. (b) Endpoint agreement of probability argmax with state and continuous-prediction nearest neighbors. Error bars are 95% sequence-bootstrap intervals. The dashed line marks the final schedule’s low-noise boundary.
Department of Statistics and Data Science, Chinese University of Hong Kong, Hong Kong · Department of Industrial and Operations Engineering, University of Michigan, Ann Arbor, USA