On the Distillation Loss Functions of Speech VAE for Unified Reconstruction, Understanding, and Generation

Abstract: Continuous speech representations based on Variational Autoencoders (VAEs) have emerged as a promising alternative to traditional spectrogram or discrete token based features for speech generation and reconstruction. Recent research has tried to enrich the structural information in VAE latent representations by aligning with self-supervised learning (SSL) features, aiming for better generation performance. However, it remains unclear whether the widely-used alignment approach based on time-axis distillation is optimal when considering more tasks. To address this problem, this paper systematically explores different alignment approaches and analyzes their impact on the performances over three axes: reconstruction, understanding, and generation. We investigate various design choices in the distillation loss. Extensive experiments show that the joint-marginal alignment approach with adaptive weighting can achieve the best overall performance while allowing for a controllable balance.

• The paper introduces a novel joint-marginal alignment loss with adaptive weighting to balance reconstruction, understanding, and generation tasks.
• It details a dual-axis methodology combining T-axis and D-axis cosine similarities to align VAE latents with foundation model features.
• Empirical evaluations on LibriSpeech show that adaptive weighting significantly stabilizes training, enhancing overall performance across speech tasks.

Distillation Loss Functions for Speech VAEs: A Unified Perspective

Introduction

The paper "On the Distillation Loss Functions of Speech VAE for Unified Reconstruction, Understanding, and Generation" (2604.12383) provides a comprehensive analysis and methodological advancement in the development of continuous speech representations using Variational Autoencoders (VAEs). The motivation arises from the intrinsic limitations of discrete representations—especially information loss in quantization—that impede downstream speech generation and understanding tasks. This work specifically scrutinizes the role of knowledge distillation loss functions in aligning VAE latents with speech foundation model features, proposing novel loss designs and adaptive weighting schemes to achieve a more balanced and unified representation for reconstruction, understanding, and generation tasks.

Problem Formulation and Existing Approaches

Present VAE-based speech representations often leverage self-supervised learning (SSL) foundation models by distilling their semantic knowledge into the latent space via T-axis alignment losses. While beneficial for generation and moderately for understanding, these approaches encounter a dilemma: improving one aspect (generation or understanding) typically degrades the others, largely due to the structural incongruence between VAE latents and SSL features.

Figure 1: T-axis Aligned Semantic VAE (TAS-VAE) distills semantic knowledge from speech foundation models via T-axis cosine loss, improving TTS performance but remaining suboptimal for speech understanding tasks such as ASR.

Distillation Loss Function Design Space

T-axis and D-axis Alignment

The canonical approach employs a T-axis (temporal) point-wise cosine similarity loss to distill representations, which outperforms standard regression losses (MAE, MSE) in generation. The D-axis alignment, orthogonal to the T-axis, targets dimension-wise correlations and is inspired by similar approaches in token-based discrete speech representation learning.

Joint-Marginal Alignment

The core novelty is the introduction of joint-marginal aligned semantic VAE (JMAS-VAE), which incorporates both frame-wise (marginal) cosine distances and pairwise sequence similarities (joint-marginal) between the VAE latent and the SSL foundation model features. This enforces not only local semantic alignment but also consistency in global latent space structure, accounting for long-range dependencies critical in speech.

Adaptive Weighting

To further refine the training dynamics, the paper formalizes adaptive loss weighting based on the gradient norm ratios of reconstruction versus distillation alignment terms. This enables dynamic balancing of loss magnitudes during optimization, in contrast to static manually-set weights, thus stabilizing training and preventing dominance of any single objective.

Figure 2: The design space of distillation loss functions includes T-axis, D-axis, and joint-marginal alignment, with the possibility of adaptive loss weighting for improved optimization.

Empirical Evaluation

A broad empirical investigation is performed using the Libriheavy and LibriSpeech datasets, benchmarking speech VAEs across speech reconstruction (PESQ, STOI), understanding (SUPERB suite: ER, PR, ASR, KS, SID, ASV, SD, IC), and generation (WER, speaker SIM for TTS). The baselines include Mel/Fbank features, EnCodec, Vanilla VAE, and previously proposed Semantic-VAE models.

Key Findings

1. TAS-VAE achieves strong generation but limited understanding; ASR WER is markedly higher than Fbank, underscoring the inadequacy of naive T-axis alignment for comprehensive speech processing.
2. DAS-VAE enhances understanding scores with moderate generation degradation.
3. JMAS-VAE with adaptive gradient-based weighting yields the highest geometric mean of overall performance, signifying effective balancing of the three objectives.
4. Adaptive weighting is critical; it rescales distillation losses orders of magnitude higher than typical static settings, driving improved understanding without catastrophic reconstruction collapse.

Figure 3: The adaptive weighting mechanism dynamically increases the relative importance of distillation losses—$\omega_{\text{mcos}}$ and $\omega_{\text{mdss}}$—throughout training for JMAS-VAE, visualized on a log-scale.

Margins and Trade-offs: Ablation of JMAS-VAE

A fine-grained grid search over margin parameters in the joint-marginal losses elucidates the trade-off curves:

• Lower margins intensify semantic alignment, improving understanding but degrading reconstruction and TTS similarity.
• Higher margins relieve excessive regularization, preserving acoustic information at the expense of semantic discrimination.

Distinct behaviors are observed: frame-wise (T-axis) margins primarily tune semantic/understanding capacity, while joint-marginal margins regulate acoustic fidelity. Pearson correlations between distance metrics (cosine and distributional) and performance scores affirm these axes' orthogonality.

Figure 4: Performance as a function of margin parameters for JMAS-VAE reveals diverging optimal regimes for reconstruction, understanding, and generation—underscoring the critical balance needed for a unified representation.

Theoretical and Practical Implications

The comprehensive analytic results establish that endowing speech VAEs with more structurally-aware distillation objectives—especially via joint-marginal alignment—enables these models to serve as viable compact continuous representations for Speech LLMs. These representations can underpin both generative (TTS) and interpretive (ASR, speaker verification) tasks within a single embedding space, streamlining model pipelines and facilitating unified speech AI systems.

Adaptive weighting is shown to be indispensable for effective multi-objective training, suggesting its integration as a default strategy in future speech representation learning architectures.

The paradigm shift from rigid static alignment to context- and distribution-aware objectives (and their dynamic weighting) opens avenues for even tighter integration with foundation models, texture-aware speech editing, and next-generation Speech LLMs with controllable trade-off frontiers across orthogonal capabilities.

Conclusion

This work rigorously characterizes the distillation loss function landscape for speech VAEs, empirically and theoretically justifying the superiority of joint-marginal alignment with adaptive weighting over conventional alignment schemes. The findings lay a robust foundation for designing speech representations that unify reconstruction, understanding, and generation, with substantial implications for multi-modal LLMs and future unified speech processing systems.