Leveraging Soft Distributions of SSL-Derived Discrete Speech Tokens for Downstream Inference

Abstract: Discrete speech tokens obtained from self-supervised learning (SSL) models provide efficient data compression while maintaining strong performance, and have been widely used as intermediate representations in various tasks. However, discretization inevitably causes information loss, leading to degraded performance compared with continuous SSL features. In this work, we propose to apply soft token assignment only during downstream inference. This approach preserves the efficiency of hard discretization during training while enhancing the expressiveness of the tokens at inference. The proposed method outperforms conventional hard assignment on both ASR and speech synthesis tasks, and exhibits particularly strong generalizability to out-of-domain data. For ASR of non-native speech, it even surpasses models using continuous SSL features. Moreover, analysis of the resulting representations shows they align more accurately with phonemes compared with conventional hard assignment.

• The paper introduces inference-time soft assignment, replacing hard token indices with a softmax-weighted embedding sum to mitigate information loss.
• It demonstrates enhanced ASR and speech synthesis performance along with improved robustness to non-native and out-of-domain speech scenarios.
• Empirical results reveal that soft assignment yields more compact phoneme clusters, optimizing downstream accuracy without increasing training costs.

Inference-Time Soft Assignment of Discrete Speech Tokens for Robust Downstream Modeling

Introduction and Problem Statement

Self-supervised learning (SSL) models have become standard front-ends for speech processing, yielding highly informative representations for downstream automatic speech recognition (ASR) and speech synthesis. Recently, converting SSL continuous features into discrete token sequences (via k-means or related quantizers) has enabled efficient data compression and has acted as a critical interface for speech LMs and codec-based generation pipelines. However, hard discretization by selecting the nearest codebook centroid introduces information loss, adversely affecting downstream task performance compared to the usage of raw SSL continuous features. Existing approaches to bridge this gap—such as codebook multiplicity, deeper clustering, or layerwise ensembles—typically forfeit compression efficiency, increase model complexity, or pose optimization challenges.

This work introduces an inference-time posterior-based soft assignment method: during training, hard tokenization is retained for maximal data compression and architectural simplicity, but at inference, instead of a hard token index, the model consumes a softmax-weighted sum of embeddings based on the posterior probability over all token centroids. The temperature parameter ($\tau$) governs the assignment sharpness. This strategy selectively restores much of the lost information for inference, improves task accuracy, and boosts out-of-domain robustness—especially for domain shifts or non-native speech—without sacrificing training-time efficiency or increasing storage budgets.

Methodology

The approach leverages standard k-means discretization to cluster SSL feature frames into $K$ centroids. For training, the downstream task (e.g., ASR, TTS) consumes discrete token IDs as inputs, with each token mapped to a trainable embedding vector. At inference, instead of using the index of the nearest centroid, the model computes posterior weights:

$$p(k|\mathbf{x}) = \frac{\exp(-D_k(\mathbf{x})/\tau)}{\sum_{j=1}^{K} \exp(-D_j(\mathbf{x})/\tau)}$$

and inputs the expectation $$\mathbf{z} = \sum_{k=1}^{K} p(k|\mathbf{x})\,\mathbf{E}_k$$ to the downstream model.

This formulation requires no change in the training procedure and does not require SSL model fine-tuning or additional codebook training as in HuBERT-Soft or residual vector quantization approaches. The method can generalize to multiple codebooks or token streams with per-layer temperature tuning.

Experimental Analysis

ASR Performance and Domain Robustness

Application of inference-time soft assignment in end-to-end ASR produces consistent improvements over conventional hard assignment on both in-domain (LibriSpeech) and out-of-domain (TED-LIUM v2, CHiME4, ERJ—non-native speech) benchmarks. Notably, even though continuous SSL features outperform all discrete variants in most in-domain cases, for non-native speech (ERJ), the soft-assigned models can outperform the continuous representations. The effect is most pronounced with small codebooks, indicating better regularization and noise suppression in challenging settings.

Ablations on the softmax temperature ($\tau$) reveal that moderate values maximize information recovery while maintaining the discrete token's robustness (Figure 1).

Figure 1: Impact of varying the softmax temperature $\tau$ for inference-time soft assignment on ASR WER across test-clean, TED-LIUM v2, CHiME4, and ERJ benchmark splits.

Speech Synthesis and Voice Conversion

Speech synthesis experiments using HiFi-GAN vocoders trained on discrete tokens (LJSpeech input) show that soft assignment yields improved Mel-Cepstral Distortion (MCD), F0 RMSE, and UTMOS, closely approaching the performance of models based on continuous features. Voice conversion from out-of-domain TIMIT speech shows a favorable balance between preserving speaker-invariant content and accurately reconstructing linguistic information. In some cases, the proposed soft-assignment method achieves superior speaker similarity and phonetic fidelity relative to both hard tokens and continuous features.

Representation Structure and Phoneme Alignment

Embedding space analysis via intra/inter-class variance indicates that token representations derived from soft assignment cluster phoneme classes more compactly and increase class separability (Fisher ratio), consistent with improved modeling of segmental information. This effect is robust across different cluster sizes and for both ASR and TTS-trained embeddings.

Multi-Codebook (Layer Aggregation) Extensions

Combining soft assignment with layerwise tokenization and aggregation (multi-stream input) further improves accuracy on both matched and mismatched domain ASR conditions, especially when layer-specific temperature tuning is applied. However, the benefit saturates or requires careful temperature scheduling to avoid overfitting to in-domain-specific acoustic details.

Comparison to Prior Work

Compared to prior approaches such as HuBERT-Soft (which requires costly retraining/reprojection of SSL models, yielding "soft" but non-discrete features), the present method preserves both codebook-driven compression efficiency and the representational advantages of soft assignments. Unlike solutions based on multiple codebooks or residual quantization, it does not multiply storage or compute costs during training and is compatible with down-streaming in large-scale spoken LLMs (2606.06806).

Implications and Future Directions

The findings substantiate that inference-time soft assignment is a principled solution to the information loss-compression tradeoff in discrete speech modeling. It offers a simple, parameter-free post-processing component that can be retrofitted to any k-means–based token pipeline. Of particular interest is the result that soft discrete tokens enable generalization to diverse domains (non-native, noisy, or otherwise mismatched speech) superior to both hard tokens and, in certain cases, continuous SSL features.

The methodology is compatible with future extensions, including deduplicated tokens, acoustic BPE, or hybrid multi-modal inputs for LLM-driven speech systems. Moreover, its implications for robust spoken language understanding, cross-lingual generalization, and codebook interpretability serve as potential areas for further investigation.

Conclusion

Applying posterior-based soft assignment to SSL-derived discrete speech tokens during inference demonstrably closes much of the information gap between hard discrete tokens and continuous features. The proposed method enables more accurate and robust downstream ASR and synthesis, especially under domain shift or with non-native input, while retaining compression efficiency during training. The work presents empirical and analytical evidence for improved phoneme separability and reveals new directions for efficient, robust speech representation learning using lightweight, retrofittable soft assignment techniques (2606.06806).