PaperlayerView exact revision on arXiv

Source-linked AI summary

Drax: Speech Recognition with Discrete Flow Matching

Aviv Navon, Aviv Shamsian, Neta Glazer, Yael Segal-Feldman, Gill Hetz, Joseph Keshet, Ethan Fetaya

arXiv:2510.04162v1eess.AScs.LGcs.SD

TL;DR

Autoregressive ASR decoding has latency that scales with sequence length, and generative non-autoregressive ASR remains comparatively underexplored. Drax uses an audio-conditioned discrete flow path; it reports competitive accuracy and improved runtime efficiency against large-scale ASR models, with favorable accuracy–efficiency trade-offs.

  • Problem

    Autoregressive ASR latency scales with output length, while simple noise-to-target paths do not reflect the acoustically plausible imperfect hypotheses encountered during inference.

    Source: §1 · p. 1

  • Method

    Drax uses a discrete flow matching path that mixes a uniform source, an audio-conditioned middle distribution, and the target transcription.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    Drax achieves competitive performance with state-of-the-art large-scale ASR models while offering improved runtime efficiency and favorable accuracy–efficiency trade-offs.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The findings support discrete flow matching as a promising foundation for future non-autoregressive ASR research.

    Source: §6 · p. 10

    Experiments use curated public multilingual datasets, and the audio-conditioned middle distribution is only one of many largely unexplored ASR path designs.

    Source: §6 · p. 10

Abstract

from arXiv · show

Diffusion and flow-based non-autoregressive (NAR) models have shown strong promise in large language modeling, however, their potential for automatic speech recognition (ASR) remains largely unexplored. We propose Drax, a discrete flow matching framework for ASR that enables efficient parallel decoding. To better align training with inference, we construct an audio-conditioned probability path that guides the model through trajectories resembling likely intermediate inference errors, rather than direct random noise to target transitions. Our theoretical analysis links the generalization gap to divergences between training and inference occupancies, controlled by cumulative velocity errors, thereby motivating our design choice. Empirical evaluation demonstrates that our approach attains recognition accuracy on par with state-of-the-art speech models while offering improved accuracy-efficiency trade-offs, highlighting discrete flow matching as a promising direction for advancing NAR ASR.

1 Introduction

Drax applies discrete flow matching to non-autoregressive ASR, using an audio-conditioned middle distribution to better align training with inference. The paper reports competitive accuracy and favorable runtime–accuracy trade-offs, motivated by the latency of autoregressive decoding and the limited development of generative NAR ASR.

1 Introduction

  • Autoregressive decoding generates tokens sequentially, so latency scales with sequence length, while generative NAR ASR remains empirically underexplored.The authors argue that principled methods are needed to balance efficiency and accuracy.Source: §1 · p. 1
  • Drax achieves competitive accuracy with state-of-the-art ASR baselines while offering favorable runtime–accuracy trade-offs.The framework also benefits from candidate scoring and integrates with speculative decoding.Source: §1 · p. 2
  • An audio-conditioned middle distribution creates a tri-mixture path that exposes the decoder to acoustically consistent but imperfect hypotheses.The path is designed to mitigate the training–inference domain gap.Source: §1 · p. 2, Figure 1 · p. 2
  • The theoretical analysis relates generalization in flow-based models to divergence between training and inference occupancies, motivating the path design.The introduction describes the design as addressing the discrepancy between training transitions and acoustically plausible inference states.Source: §1 · pp. 1–2

2 Related Work

Related work contrasts established autoregressive ASR and CTC with emerging generative non-autoregressive approaches. The authors identify limited empirical study of generative NAR ASR, particularly on large-scale or multilingual benchmarks.

2 Related Work

  • CTC has been widely adopted, but its conditional-independence assumption limits long-range dependencies and helped motivate autoregressive approaches.Autoregressive systems decode sequentially, retaining an inference-speed bottleneck.Source: §2 · pp. 2–3
  • Diffusion and flow matching extend generative modeling to discrete sequences, while prior speech applications include discrete-latent speech modeling and speech-to-speech translation.Discrete flow models define paths over discrete state spaces through continuous-time Markov chains.Source: §2 · p. 3
  • Prior generative NAR ASR studies use multinomial diffusion or diffusion-based decoding, but mainly target English and lack large-scale or multilingual evaluation.The section names Transfusion, FFDM, and Whisfusion as examples.Source: §2 · p. 3

3 Method

Drax uses discrete flow matching to generate audio-conditioned text through a probability path with an audio-conditioned middle distribution. The method specifies a flow-based training formulation and sampling procedure, then considers candidate scoring to select among stochastic outputs.

3.1 Preliminaries

  • Discrete flow matching learns a time-dependent probability path from a source token distribution to the target data distribution, with a velocity field generating that path.The framework uses mixture conditional paths and trains with cross-entropy loss.Source: §3.1 · pp. 3–4

3.2 Speech Recognition with DFM

  • The audio-bridged path mixes a uniform source, an audio-conditioned middle distribution, and the ground-truth target to represent plausible intermediate hypotheses.The middle component is concentrated around t = 0.5; at inference, the method sets its velocity contribution to zero and samples as in the two-way case.Source: §3.2 · pp. 4–5
  • The method uses a pretrained Whisper audio encoder and a DiT decoder with audio cross-attention, while training the middle-distribution network jointly with the decoder.Training uses Gumbel-Softmax samples and an auxiliary cross-entropy loss on middle-distribution logits.Source: §3.2 · p. 5
  • Sampling integrates the marginal velocity in parallel across token positions, and candidate scoring can select among multiple stochastic transcriptions.Scoring options include sample frequency, minimum Bayes risk, external-model likelihood, and an ELBO-based internal score.Source: §3.2 · pp. 5–6
  • The figure presents RTFx by sequence length and a WER–RTF Pareto front, with Drax variants offering favorable trade-offs and control over the operating point.RTF is expressed as 100/RTFx.Source: Figure 2 · p. 6

4 Theoretical Analysis: Occupancy-Based Bound

The analysis bounds the training–inference generalization gap through discrepancies between their occupancies. It connects this gap to velocity-field error and motivates probability paths that reduce the discrepancy.

4 Theoretical Analysis: Occupancy-Based Bound

  • Reducing the training–inference occupancy discrepancy can tighten the total-variation generalization bound.Source: §4 · p. 6
  • Velocity-field error can compound as sampling trajectories diverge from the desired path, further widening the occupancy gap.Source: §4 · p. 6
  • The stability setup assumes p_0 = q_0 on a finite state space.Source: §4 · p. 6
  • The instantaneous bound separates intrinsic model error under p_t from an added contribution due to mismatch between q_t and p_t.Source: §4 · p. 7
  • The result is only an upper bound, and a deeper investigation of the connection is needed.Source: §4 · p. 7

5 Experiments

Experiments evaluate Drax on public multilingual ASR benchmarks and examine its accuracy–efficiency trade-offs and training paths. Drax performs on par with or better than strong baselines, while path design and decoding choices affect performance.

5 Experiments

  • Baselines span encoder-decoder, LLM-augmented, CTC, and diffusion ASR; evaluation reports WER/CER for accuracy and RTF/RTFx for runtime.Source: §5 · pp. 7–8

5.1 Multilingual Speech Recognition

  • Across broad English and multilingual public benchmarks, Drax is on par with or surpasses strong ASR baselines.Source: §5.1 · p. 8, Table 2 · p. 8

5.2 Accuracy-Efficiency Trade-off

  • NAR sampling is parallel across positions with fixed NFE; increasing NFE improves WER, while fewer steps yield faster decoding.Source: §5.2 · p. 9
  • MBR consistently achieves strong results comparable with Whisper scoring, while ELBO-based scoring is unstable; 8 NFE and 16 candidates offer a selected trade-off.Source: §5.2 · p. 9

5.3 Speculative Decoding

  • As a draft model, Drax achieves RTFx in the range of ∼20x to ∼42x and produces more matched tokens; it yields substantial speedups over Whisper-Turbo, especially on non-English languages.Source: §5.3 · p. 9

5.4 Training Path Design

  • The uniform-source, audio-conditioned-middle path has the lowest generalization WER throughout training; middle distributions improve over direct source-target paths.Source: Figure 3 · p. 10, §5.4 · p. 10

6 Discussion

Drax uses a tri-mixture probability path with an audio-conditioned middle distribution, and its analysis motivates that design. The conclusion reports competitive recognition performance with improved runtime efficiency, alongside limitations in evaluation scope and path-design exploration.

6 Discussion

  • Experiments use curated public multilingual datasets, and larger or more diverse corpora may reveal additional robustness and generalization challenges.The audio-conditioned middle distribution is one path instantiation; alternative or adaptive constructions remain for future investigation.Source: §6 · p. 10
  • Drax combines a tri-mixture probability path with an audio-conditioned middle distribution, supported by theoretical analysis motivating the path design.Source: §6 · p. 10
  • Drax achieves competitive performance with state-of-the-art large-scale ASR models while offering improved runtime efficiency.The framework also integrates naturally with candidate scoring and speculative decoding strategies.Source: §6 · p. 10

A Proofs

The appendix provides proof details for a generalization analysis based on training and generation distributions, derives velocity fields for the tri-mixture path, and specifies its scheduler.

A Proofs

  • The bounded-loss generalization analysis connects training–generation occupancy discrepancy to intrinsic model error and an additional domain-gap contribution.The appendix presents a path-marginal TV result and an occupancy-TV generalization bound.Source: A Proofs · pp. 18–20

B Tri-mixture Velocities

  • The appendix gives unconditional and conditional probability velocity fields for the tri-mixture path, with the scheduler construction dropping source-component terms.Source: B Tri-mixture Velocities · pp. 20–21

C.1 Tri-Mixture Scheduler

  • The factorized scheduler interpolates source, middle, and target distributions; with p = 2 and q = 2/3, the middle weight peaks at t = 0.5 and the path has pure-source and pure-target endpoints.Figure 4 presents the tri-mixture sampling scheduler.Source: §C.1 · p. 21, Figure 4 · p. 21

C.2 Training and Optimization Details

Training uses frozen Whisper large-v3 audio encoders and DiT decoders, public multilingual data, and specified optimization and conditioning procedures; sampling uses a fixed sequence length and cached audio computations.

C.2 Training and Optimization Details

  • Drax uses a frozen Whisper large-v3 encoder and a 16-block, 580M-parameter decoder; Drax-flash uses a 4-layer, 250M-parameter decoder.Both use a 28M-parameter audio-conditioned middle-distribution module.Source: §C.2 · p. 22
  • Models use AdamW with a 2500-step warmup, peak learning rate 3e−4, and batch size 240; Drax and Drax-flash train for 800K and 250K iterations.Audio conditioning is randomly dropped with probability 0.1, and language-token replacement follows the path with probability 0.2.Source: §C.2 · p. 22
  • Training combines public datasets spanning eight languages and approximately 15K hours of speech.Source: §C.2 · p. 22

C.3 Sampling

  • Inference uses a linear scheduler, fixed sequence length L = 144, and cached audio projections and cross-attention keys and values across generation steps.Sampling uses τ = 0.01 for a single transcript and τ = 0.1 for multiple candidates.Source: §C.3 · p. 23

C.4 Runtime Results

Runtime comparisons use a controlled single-GPU evaluation setup, while the accuracy-efficiency analysis grids NFE and ensemble size on MLS.

C.4 Runtime Results

  • Runtime metrics are measured on one L40s GPU with batch size 1, full precision, and no compilation.Source: §C.4 · p. 23

D.1 Accuracy-efficiency trade-off

  • The MLS accuracy-efficiency evaluation grids NFE over 4, 8, and 16 and ensemble size over 1, 8, and 16 across scoring methods.The passage identifies Table 4 as reporting WER and RTFx for these configurations.Source: §D.1 · p. 23

D.2 Drax-flash Results

Drax-flash is evaluated on the EN benchmark and MLS, with reported runtime improvements over Whisper and Qwen2-Audio.

D.2 Drax-flash Results

  • Drax-flash provides significant runtime improvements, with RTFx compared against Whisper and Qwen2-Audio across transcription sequence lengths.The evaluation covers the EN benchmark and MLS; Table 6 is scoped to the English benchmark.Source: §D.2 · p. 23, Table 6 · p. 23
  • The English benchmark results are reported for Drax-flash in Table 6.Source: Table 6 · p. 23

D.3 Extended English Benchmark Results

The extended English results and MLS analyses examine accuracy-efficiency trade-offs across scoring, sampling temperature, NFE, and use of pmid during inference.

D.3 Extended English Benchmark Results

  • Drax offers control over the accuracy-efficiency trade-off through NFE, candidate set size, and scoring choices on extended English results.The reported results allow users to select an operating point.Source: §D.3 · p. 23

D.4 Sampling Temperature

  • Lower sampling temperatures, such as 0.01–0.1, yield the best MLS WER across languages, while temperature 1.0 noticeably degrades performance.Increasing temperature leads to more diverse generations but also higher error rates.Source: §D.4 · p. 24

D.5 Effect of NFE

  • Drax achieves high-quality sampling with 4–16 steps, and WER shows relatively small improvement for NFE ≥4.The reported plots cover MLS and VoxPopuli per-language WER.Source: §D.5 · p. 24

D.6 Sampling with and without pmid

  • Including pmid during generation consistently hurts performance across datasets, supporting its use as a training-only component.Source: §D.6 · p. 24

D.7 Samples from Audio Distribution

The appendix illustrates samples from the learned audio distribution and Drax’s generation path, including how token sequences refine across steps.

D.7 Samples from Audio Distribution

  • Samples from the learned pmid distribution are acoustically plausible but imperfect.Table 11 presents samples from this learned audio distribution.Source: §D.7 · p. 24, Table 11 · p. 28

D.8 Generation Path

  • Figures 6 and 7 show ensemble prediction through candidate scoring for Tedlium and VoxPopuli (EN), while Figure 8 concerns the effect of sampling steps.Source: Figure 6 · p. 25, Figure 7 · p. 26, Figure 8 · p. 26
  • Drax’s stepwise refinement stabilizes tokens first in shorter blocks and then in longer spans until the full sequence is recovered.Table 10 traces a 16-step example; each column is a token position, and “_” marks positions unchanged from the initial random state.Source: §D.8 · pp. 24–25, Table 10 · p. 27
  • Including pmid during sampling degrades accuracy, supporting its role as a training-only component.Source: Figure 9 · p. 27
Loading 2510.04162v1…