Soft Token Prediction Loss

• Soft token prediction loss is a family of loss functions that replaces one-hot targets with semantically enriched distributions to capture token-level uncertainty.
• It generalizes cross-entropy by incorporating temperature-controlled smoothing and mixes soft and hard targets to better handle paraphrases and overconfidence.
• The approach is applied in language modeling, visual prediction, and time-series forecasting, leading to improvements in BLEU scores, RMSE, and reduced exposure bias.

Soft token prediction loss refers to a family of loss functions in which token-level supervision is provided via distributions ("soft targets") rather than strict one-hot targets. This approach generalizes the classical cross-entropy loss in neural sequence models by replacing pointwise (Dirac) supervision with token- or embedding-level distributions that encode semantic proximity, probabilistic uncertainty, or structural relationships in the output space. Soft token prediction losses have been developed in language modeling, autoregressive visual prediction, and risk-aware forecasting, often motivated by issues such as overconfidence, poor generalization to paraphrases or synonyms, and exposure bias in teacher-forced training. Recent advances combine these techniques with sequence-level smoothing, dense representation-based supervision, or trajectory-level backpropagation for manifold regularization and stability.

1. Motivation and Conceptual Foundations

The standard maximum likelihood estimation (MLE) objective in neural sequence modeling minimizes the negative log-likelihood of the true token at each step, using a one-hot (Dirac) target. This leads to the loss

$$\mathcal{L}_\text{MLE} = -\sum_{t=1}^T \ln p_\theta(y^*_t|h^*_t) = \sum_{t=1}^T D_{\mathrm{KL}}(\delta_{y^*_t} \| p_\theta(\cdot|h^*_t))$$

where $p_\theta(\cdot|h^*_t)$ is the model's prediction conditioned on the history, and $\delta_{y^*_t}$ is the one-hot vector for the ground-truth token (Elbayad et al., 2018).

Such "hard" supervision implies that any incorrect token, regardless of semantic similarity to the correct word, is penalized equally. This results in brittle, overconfident token distributions and inadequate modeling of semantic or contextual closeness, especially in settings where paraphrases or near-synonyms are meaningful. By constructing "soft" target distributions—assigning probability mass to near-misses according to a reward, similarity, or risk metric—soft token prediction losses align the training signal with the structure in the output space, mitigate overconfidence, and facilitate representational generalization (Elbayad et al., 2018, Ning et al., 2023, Namazi et al., 10 Dec 2025).

2. Mathematical Formulation and Construction

In token-level smoothing for LLMs, the hard target at each time step is replaced with a reward-informed distribution $\tilde r(\cdot | y^*_t)$. The construction is: $$\tilde r(w|y^*_t) = \frac{ \exp(s(w, y^*_t)/\tau)}{ \sum_{v\in\mathcal V} \exp(s(v, y^*_t)/\tau)}$$ where $s(w, y^*_t)$ is a similarity function (e.g., cosine similarity in a pretrained embedding space), $\tau$ is a temperature parameter governing sharpness, and $\mathcal{V}$ is the vocabulary (Elbayad et al., 2018). The loss per token position is then the KL-divergence between the soft target and the model’s prediction: $$\mathcal{L}_\text{Tok} = \sum_{t=1}^T D_{\mathrm{KL}}(\tilde r(\cdot|y^*_t) \| p_\theta(\cdot|h^*_t))$$ A weighted interpolation with the classic MLE loss can be used: $$\mathcal{L}_{\text{Tok},\,\alpha} = \alpha \mathcal{L}_\text{Tok} + (1-\alpha)\mathcal{L}_\text{MLE}$$ where $p_\theta(\cdot|h^*_t)$0 balances soft and hard targets.

In discrete VQ-VAE-based visual prediction, a soft token is a probability vector $p_\theta(\cdot|h^*_t)$1 (simplex over the codebook), derived as: $p_\theta(\cdot|h^*_t)$2 where $p_\theta(\cdot|h^*_t)$3 is the encoder output and $p_\theta(\cdot|h^*_t)$4 are codebook embeddings. The token prediction loss can then be written either as categorical cross-entropy (if hard targets are used) or as a KL divergence (if soft assignments are targets): $p_\theta(\cdot|h^*_t)$5 where $p_\theta(\cdot|h^*_t)$6 is the model’s predicted distribution at position $p_\theta(\cdot|h^*_t)$7 (Ning et al., 2023).

In time-series forecasting, the soft-token prediction loss is the average cross-entropy between the model's predicted distribution and the true one-hot at each step of the unrolled autoregressive trajectory. However, at each future step, the model feeds back its own predicted soft embedding rather than the hard ground-truth, enabling differentiable trajectory fine-tuning (Namazi et al., 10 Dec 2025): $p_\theta(\cdot|h^*_t)$8 with all embeddings and hidden states fully differentiable throughout the unroll.

3. Practical Implementation and Pseudocode

Below is an overview of the core algorithmic steps for training with soft token prediction loss in the token-level smoothing scenario (Elbayad et al., 2018):

$\delta_{y^*_t}$4 Adjustments for sequence-level or embedding-level smoothing, as in RAML or NITP, modify the construction of targets, the loss accumulation, and the nature of the auxiliary objectives (Elbayad et al., 2018, Zhang et al., 24 May 2026).

For VQ-VAE visual models and autoregressive decoders with soft tokens, the steps include soft assignment calculation, KL or CE prediction loss, auxiliary reconstruction via the detokenizer, and standard backpropagation through all differentiable modules (Ning et al., 2023).

In fully differentiable unrolled trajectory forecasting, soft tokens serve as the feedback at every autoregressive step, and the loss is summed cross-entropy at each future time point, without the need for reinforcement or sampling (Namazi et al., 10 Dec 2025).

4. Integration with Sequence-Level and Continuous Representation Losses

Soft token prediction can be combined with sequence-level smoothing, such as Reward Augmented Maximum Likelihood (RAML), wherein entire sampled sequences are weighted by global task metrics (e.g., BLEU, CIDEr), and token-level smoothing is applied to each sampled sequence (Elbayad et al., 2018). This two-parameter mixture can be tuned to optimize for both local semantic proximity and global sequence quality.

In Next Implicit Token Prediction (NITP), soft supervision is further extended: in addition to predicting the next one-hot token, the model predicts the continuous representation of the next token derived from its own shallow layers, using a cosine similarity loss: $p_\theta(\cdot|h^*_t)$9 where $\delta_{y^*_t}$0 is the target shallow feature with stop-gradient, and $\delta_{y^*_t}$1 is a projection head (Zhang et al., 24 May 2026). This dense loss regularizes the representation geometry, closing the under-constrained orthogonal subspace left by pure cross-entropy supervision.

5. Applications and Empirical Results

Soft token prediction loss has demonstrated empirical gains across multiple domains:

• Language Modeling and Machine Translation (token-level smoothing): On MS-COCO, token-level smoothing yields BLEU-4 improvements (e.g., 30.14 → 31.27), and further gains when combined with sequence-level smoothing (e.g., up to BLEU-4 31.39). Gains are also observed in CIDEr and WMT/IWSLT translation benchmarks (Elbayad et al., 2018).
• Visual Tasks with VQ-VAE (soft assignment): In "All in Tokens," soft tokens, used in either task-solver or detokenizer stages, lead to consistent improvements over hard token baselines: reductions in depth estimation RMSE (0.3174 → 0.3080) and gains in mask mAP (31.1 → 33.2). Stacking soft tokens in both places and adding an auxiliary reconstruction loss yields further improvements (e.g., mask mAP 34.2) (Ning et al., 2023).
• Risk-Aware Time Series Forecasting: In Soft-Token Trajectory Forecasting (SoTra), the soft-token prediction loss (computed over unrolled trajectories) reduces average zone-based risk in glucose forecasting by 18% and clinical risk in blood-pressure forecasting by 15% (Namazi et al., 10 Dec 2025). The differentiable trajectory-level fine-tuning directly mitigates exposure bias, as the training graph matches inference.
• Implicit Token Prediction in LLM Pre-Training: The NITP approach demonstrates 5.7% absolute improvement on MMLU-Pro, 6.4% on C3, and 4.3% on CommonsenseQA (for 9B MoE), with only a 2% increase in training FLOPs and no inference penalty (Zhang et al., 24 May 2026). It also measurably increases the effective rank of hidden representations and improves transfer to downstream embedding tasks.

6. Theoretical and Practical Considerations

Soft token prediction loss addresses key defects in standard cross-entropy objectives by:

• Distributing the training signal according to semantic, metric, or probabilistic similarity, rather than only exact token matches.
• Providing a mechanism for regularizing hidden states and token predictions, which can reduce overconfidence and improve model calibration.
• Enabling differentiable feedback over multi-step or autoregressive trajectories, directly addressing exposure bias without reliance on sampling or reinforcement gradients.
• Integrating easily with existing workflows by interpolating with MLE objectives, controlling the trade-off via hyperparameters such as $\delta_{y^*_t}$2 (temperature) and $\delta_{y^*_t}$3 (mixing weight).

However, care must be taken in the construction of the reward/similarity function, the tuning of temperature and mixing, and the maintenance of strong signal on true tokens to avoid underweighting the ground-truth.

In summary, soft token prediction losses represent a general class of improvements to supervision in sequential prediction, capturing uncertainty, structure, and continuity at the token level with demonstrated gains in both generative quality and robustness (Elbayad et al., 2018, Ning et al., 2023, Namazi et al., 10 Dec 2025, Zhang et al., 24 May 2026).