Source-linked AI summary
Beyond Token Revision: Investigating Mask-and-Replace Diffusion for Zero-Shot Text-to-Speech
Hounsu Kim, Joonyong Park, Yuki Saito, Satoru Fukayama, Juhan Nam
TL;DR
The paper asks whether mask-and-replace training benefits zero-shot TTS only by enabling inference-time token revision, or whether randomly perturbed training context also helps. DeMaR combines that training with confidence-ranked mask-only sampling, and its WER advantage persists under commit-only sampling. Matched comparisons find benefits from both noisy-context augmentation and replacement supervision without token revision.
Problem
Mask-and-replace training gains are often attributed to self-correction through inference-time revision, but randomly perturbed training context may itself improve generation.
Source: §1 · p. 1Method
DeMaR combines mask-and-replace training with confidence-ranked mask-only sampling while assigning the combined masking and replacement probability to masks at inference.
Source: §3.2 · p. 2Results
Under commit-only sampling, DeMaR beats Mo-DLM at 5.23 versus 5.97% WER for CosyVoice3 and 5.03 versus 5.38% for FlexiCodec.
Source: §5.1 · p. 4Takeaways & Limitations
Noisy-context augmentation and replacement supervision improve WER even when generated tokens are never revised.
Source: §6 · p. 4
Abstract
from arXiv · showhide
Unlike autoregressive models, discrete diffusion-based models for zero-shot text-to-speech generate speech tokens in parallel and can revisit earlier predictions. Mask-and-replace training extends mask-only training by randomly replacing some tokens, and its gains are commonly attributed to self-correction, the ability to revise previously generated tokens. However, exposure to randomly perturbed context during training may itself improve generation, raising the question of whether these gains require inference-time token revision. To investigate this question, we use DeMaR, which combines mask-and-replace training with confidence-ranked mask-only sampling while preserving the total training corruption probability. Trained from scratch on LibriTTS, DeMaR achieves lower word error rates (WER) than autoregressive and mask-only diffusion baselines using the same speech tokenizer. This advantage persists when each token remains unchanged after first being unmasked. Matched training conditions on two heterogeneous speech tokenizers show that both noisy-context augmentation and replacement supervision improve WER under this restriction. These findings demonstrate training-side benefits of replacement beyond enabling inference-time token revision.
1. INTRODUCTION
Discrete diffusion generates speech tokens through repeated parallel predictions, while mask-and-replace training exposes the denoiser to randomly substituted tokens. The introduction asks whether these training-side benefits persist without inference-time token revision and presents DeMaR to test that question.
Background
- Discrete diffusion models recover clean speech-token sequences through repeated parallel predictions, unlike autoregressive models, which generate tokens sequentially.A waveform decoder converts discrete speech tokens to audio; the text supplies linguistic content.Source: §1 · p. 1
Motivation
- Noisy replaced-token context may improve prediction of masked tokens even without replacement supervision or inference-time revision.The paper distinguishes this proposed augmentation benefit from reconstruction loss at replaced positions, which may support self-correction.Source: §1 · p. 1
DeMaR
- DeMaR combines mask-and-replace training with confidence-ranked mask-only sampling, assigning the combined masking and replacement probability to masks at inference.Its commit-only variant tests the same trained model and masking schedule without changing a token after it is first unmasked.Source: §1 · p. 1
Contributions
- DeMaR retains a WER advantage over mask-only diffusion under commit-only sampling, and matched conditions show gains from both noisy-context augmentation and replacement supervision.The comparisons cover CosyVoice3 fixed-rate tokens and FlexiCodec variable-rate content and span sequences.Source: §1 · p. 1
2. RELATED WORK
Related work treats inference-time revision and mask-and-replace training as distinct approaches to handling incorrect visible tokens. DeMaR adapts training–sampling separation to speech-token generation and tests its gains without inference-time revision.
Inference-time token revision
- Prior revision methods remask visible tokens or overwrite them directly, while standard masked decoders and several confidence-ranked generators preserve unmasked tokens, [12].ReMDM and Mask, Sample, Revise apply remasking without retraining [14], while PALLE uses a separate refinement stage; GIDD and SCOPE pair overwriting with training on incorrect visible tokens [16].Source: §2 · p. 1
Mask-and-replace training
- Mask-and-replace training exposes denoisers to incorrect visible tokens as reconstruction targets and as context; prior work applies it across image, sound, text, and TTS settings.GIDD suggests self-correction during denoising, CDLM supervises corrupted visible tokens [20], and SCOPE trains on self-generated predictions matched to its sampler [16].Source: §2 · p. 1
Decoupling training and sampling corruption
- D3RM uses mask-and-replace training with mask-only sampling for piano transcription; DeMaR adapts this separation to speech tokens with confidence-ranked masking.The paper tests whether the resulting gains persist without inference-time revision and whether augmentation and replacement supervision each contribute.Source: §2 · p. 2
3. DEMAR AND MATCHED COMPARISONS
DeMaR trains conditional speech-token denoisers with masking and random replacement, then samples using confidence-ranked masks. The matched comparisons separate training corruption and supervision from inference-time revision.
3. DEMAR AND MATCHED COMPARISONS
- Zero-shot TTS is formulated as text-conditioned speech-token continuation: a denoiser predicts a categorical distribution over the ground-truth token at each target position.The conditioning sequence contains prompt and target text tokens followed by uncorrupted prompt speech tokens.Source: §3 · p. 2
3.1. Training corruption and objective
- At each diffusion time, DeMaR independently retains, masks, or randomly replaces target tokens; the matched mask-only baseline masks with the combined masking-plus-replacement probability.Both models reconstruct original tokens with the same weighted cross-entropy; FlexiCodec applies the loss to independently corrupted content and span streams.Source: §3.1 · p. 2
3.2. Confidence-based sampling
- DeMaR assigns its combined training corruption probability to masks at inference, and confidence rankings select the lowest-confidence positions for masking.Confidence is based either on proposal probability or negative entropy; revisiting allows token revision and demotion of visible tokens to masks.Source: §3.2 · pp. 2–3
3.3. Controlling inference-time revision
- Commit-only sampling uses proposals only at currently masked positions, fills the highest-confidence positions, and disables both token-to-token revision and token-to-mask demotion.The separately tested T2T-only and T2M-only policies permit one revision pathway at a time while keeping the same mask count.Source: §3.3 · p. 3
3.4. Separating augmentation from replacement supervision
- The three matched training conditions isolate no replacement, unsupervised random replacement, and replacement with reconstruction loss at replaced positions.All retain the same sampling mask budget and denominator normalization, enabling comparisons of augmentation and replacement supervision under revisiting and commit-only sampling.Source: §3.4 · p. 3
4. EXPERIMENTAL SETUP
The experiments train speech-token generators from scratch on LibriTTS using two frozen tokenizer and waveform-decoder systems. Decoding is tuned on dev-clean WER, and intelligibility, naturalness, and speaker similarity are evaluated on test-clean.
4.1. Data and model training
- Models are trained from scratch on 585 hours of LibriTTS with frozen CosyVoice3 and FlexiCodec tokenizers and waveform decoders, using a shared 84.95M-parameter backbone.Each model uses 12 layers, hidden dimension 768, 12 attention heads, and AdamW with a 4 × 10^-4 learning rate; three seeds are used for each Section 3.4 condition.Source: §4.1 · p. 3
4.2. Decoding hyperparameter tuning
- Decoding hyperparameters are tuned by dev-clean WER using Gaussian-process Bayesian optimization, with candidates searched on 1,000 utterances and top candidates reranked on a separate 3,000.Search includes temperature, top-p, repetition penalty, and— for AR-LMs—repetition-aware sampling.Source: §4.2 · p. 3
4.3. Evaluation protocol and metrics
- Evaluation uses all 4,835 LibriTTS test-clean utterances and measures intelligibility with WER, naturalness with UTMOS22-strong, and speaker similarity with WavLM-Large.Within each tokenizer, systems share evaluation pairs, text conditioning, reference speech-token prefixes, and frozen waveform decoders.Source: §4.3 · p. 3
5. RESULTS AND ANALYSIS
DeMaR improves WER over autoregressive and mask-only diffusion baselines, and its advantage persists when unmasked tokens cannot be revised. Under that restriction, noisy-context augmentation improves WER and replacement supervision adds further, less uniform gains.
5. RESULTS AND ANALYSIS
- DeMaR lowers Mo-DLM WER from 6.05 to 4.95% for CosyVoice3 and 5.39 to 4.89% for FlexiCodec, while outperforming both Mo-DLM and AR-LM.Table 1 reports generation quality across decoding settings.Source: §5 · p. 3, Table 1 · p. 4
- For CosyVoice3, DeMaR raises UTMOS from 4.127 to 4.276 and SIM from 0.701 to 0.717 relative to AR-LM.For FlexiCodec, both DLMs achieve approximately 3.88 UTMOS and 0.694 SIM.Source: §5 · p. 4
5.1. Contribution of inference-time revision
- Commit-only raises DeMaR WER by just 0.29 percentage points for CosyVoice3 and 0.14 for FlexiCodec, while DeMaR still beats Mo-DLM.Under commit-only, WER is 5.23 versus 5.97% for CosyVoice3 and 5.03 versus 5.38% for FlexiCodec.Source: §5.1 · p. 4, Table 1 · p. 4
5.2. Noisy-context augmentation and replacement supervision
- Without replacement supervision, adding random replacements lowers seed-1 commit-only WER from 6.87 to 5.73% for CosyVoice3 and 5.61 to 5.27% for FlexiCodec.Table 2 compares matched training conditions across seeds; its parentheses report commit-only WER.Source: §5.2 · p. 4, Table 2 · p. 4
- Across seeds, augmentation lowers revisiting mean WER from 6.91 ± 0.24 to 5.61 ± 0.58% for CosyVoice3 and 5.75 ± 0.31 to 5.50 ± 0.31% for FlexiCodec.Replacement supervision further reduces those means to 5.02 ± 0.41% and 4.98 ± 0.30%, respectively, but its gains are less uniform across seeds.Source: §5.2 · p. 4, Table 2 · p. 4
- With replacement supervision, CosyVoice3 WER spans only 5.02–5.08% across four policies, and neither T2T revision nor T2M demotion substantially improves on commit-only.Table 3 controls T2T revision and T2M demotion separately for seed 1’s CosyVoice3 setting.Source: §5.2 · p. 4, Table 3 · p. 4
6. CONCLUSION
The conclusion is that mask-and-replace training benefits do not require inference-time token revision. Noisy-context augmentation and replacement supervision improve WER even when generated tokens are never revised.
6. CONCLUSION
- DeMaR retains its WER advantage over mask-only diffusion when each token remains unchanged after first being unmasked.WER rises by only 0.29 percentage points for CosyVoice3 and 0.14 for FlexiCodec.Source: §6 · p. 4
- Noisy-context augmentation improves WER without replacement supervision in all three seeds for both tokenizers under revisiting and commit-only sampling.Source: §6 · p. 4
- Replacement supervision provides further gains in two of three CosyVoice3 seeds and all three FlexiCodec seeds under both sampling policies.Source: §6 · p. 4