⚡ electrical engineering

Token-Based Audio Inpainting via Discrete Diffusion

This paper introduces a novel token-based discrete diffusion framework that effectively restores long missing segments in degraded audio recordings by leveraging pre-trained tokenizers and incorporating derivative-based regularization and span-based absorbing transitions to outperform existing baselines on MusicNet and MAESTRO datasets.

Original authors: Tali Dror, Iftach Shoham, Moshe Buchris, Oren Gal, Haim Permuter, Gilad Katz, Eliya Nachmani

Published 2026-02-18

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tali Dror, Iftach Shoham, Moshe Buchris, Oren Gal, Haim Permuter, Gilad Katz, Eliya Nachmani

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have an old, cherished vinyl record of your favorite symphony. Over the years, the needle has skipped, or perhaps a spill of coffee has ruined a few seconds of the music. You can hear the notes before the gap and the notes after it, but the middle is just silence or static. Your goal is to fill in that missing piece so perfectly that no one can tell it was ever damaged.

This is the problem of Audio Inpainting.

For a long time, computers tried to fix this by looking at the raw sound waves (the squiggly lines) or the frequency charts (like a visual map of the sound). But when the missing piece is long—say, a whole musical phrase—these methods often get confused. They might fill the gap with a weird, robotic noise or a note that doesn't fit the melody. It's like trying to finish a sentence by guessing random words; the grammar might be okay, but the meaning is lost.

The New Approach: Turning Music into LEGO

The paper introduces a new method called AIDD (Audio Inpainting via Discrete Diffusion). Instead of trying to fix the raw sound waves directly, the authors decided to translate the music into a different language first: Discrete Tokens.

Think of the audio signal not as a continuous river of water, but as a long train of LEGO bricks.

The Old Way: Trying to smooth out a broken bridge by pouring more water over it.
The New Way: Taking the bridge apart, turning the broken section into a pile of specific LEGO bricks, and then asking an expert builder to figure out exactly which bricks fit in the gap to rebuild the bridge perfectly.

Here is how their "LEGO builder" works, broken down into three simple steps:

1. The Translator (Tokenization)

First, the computer uses a tool called WavTokenizer. Imagine this as a translator that listens to the music and converts every sound into a short code, like a sequence of numbers or symbols (e.g., "Note C," "Drum hit," "Silence").

Why do this? It's much easier for a computer to predict the next word in a sentence (or the next LEGO brick in a sequence) than to predict the exact shape of a water droplet. This allows the computer to understand the meaning and structure of the music, not just the raw noise.

2. The "Blindfolded" Builder (Discrete Diffusion)

This is the magic part. The computer takes the sequence of LEGO bricks and deliberately breaks some of them, replacing them with a generic "BLANK" brick (a mask).

The Process: The computer then plays a game of "guess the missing piece." It looks at the bricks before and after the gap and tries to guess what the blank bricks should be.
The Twist: It doesn't guess perfectly on the first try. It makes a rough guess, then refines it, then refines it again, over and over (like a sculptor chipping away at stone). This is called Diffusion. Because it's working with LEGO bricks (discrete tokens) instead of water (continuous waves), it can handle huge gaps without getting lost. It can "imagine" a whole new melody that fits perfectly with the rest of the song.

3. The Two Secret Weapons

To make sure the rebuilt music sounds natural and not robotic, the authors added two special rules to the training:

The "Smoothness" Rule (Derivative-Based Loss):
Imagine you are drawing a line on a piece of paper. If you suddenly jerk your hand, the line looks jagged and unnatural. The computer needs to learn that music flows smoothly. The authors added a rule that punishes the computer if its guesses jump around too wildly. It forces the "LEGO bricks" to transition smoothly, ensuring the music doesn't sound like a stuttering robot.
The "Chunk" Rule (Span-Based Masking):
In the past, computers might have been trained by hiding just one brick at a time. But in real life, gaps are often big chunks of missing audio. The authors trained the computer by hiding entire stretches of bricks at once. This teaches the model to understand the big picture and the context of the whole phrase, not just the immediate neighbor.

The Results: Why It Matters

The researchers tested this on two big datasets of classical music (MusicNet and MAESTRO). They created gaps ranging from tiny skips to huge holes (up to 750 milliseconds, which is a long time in music).

The Outcome: Their method filled in the gaps better than any previous method.
The Analogy: If other methods were like a child trying to finish a puzzle by gluing in random pieces that sort of fit, AIDD is like a master puzzle solver who looks at the picture on the box and knows exactly what the missing piece looks like, even if the hole is huge.
Efficiency: Not only was the sound better, but the model was also smaller and faster to train than its competitors.

In Summary

This paper is about teaching computers to "listen" to music in a smarter way. By turning sound into a sequence of discrete tokens (like words in a sentence or bricks in a wall) and using a smart guessing game (diffusion) to fill in the blanks, they can restore damaged recordings with a level of musical intelligence that was previously impossible. It's not just fixing the noise; it's re-imagining the music.

1. Problem Statement

Audio Inpainting is the task of reconstructing missing or corrupted segments of an audio signal. While traditional methods (autoregressive models, sparse representations) and recent deep generative models (VAEs, continuous diffusion models) have shown promise, they face significant limitations when dealing with long gaps (e.g., >150 ms).

Waveform-level diffusion models struggle to capture long-range structure due to the high computational cost of large receptive fields.
Spectrogram/CQT-based methods often lose semantic coherence and struggle with phase reconstruction over extended regions.
VAE-based models maintain local continuity but fail to preserve high-level semantic consistency as gap lengths increase.

The core challenge is reconciling fine-grained temporal details with high-level semantic structure in continuous audio representations.

2. Methodology: AIDD (Audio Inpainting via Discrete Diffusion)

The authors propose AIDD, the first framework to apply discrete diffusion to tokenized music representations. The architecture consists of three main components:

A. Tokenization (WavTokenizer)

Instead of operating on raw waveforms or spectrograms, AIDD converts audio into a compact sequence of discrete tokens using a pre-trained WavTokenizer.

This transforms the inpainting problem from a continuous regression task into a discrete sequence completion task.
The tokenizer compresses high-resolution audio while preserving semantic content and reconstruction fidelity.

B. Discrete Diffusion Model (DiT)

The core generative engine is a Diffusion Transformer (DiT) operating in the categorical token space.

Forward Process: Uses an absorbing transition mechanism where tokens are progressively masked (replaced by a [MASK] token) based on a noise schedule.
Reverse Process: The model learns to predict the "concrete score" (the ratio of probabilities between token states) to denoise the sequence and recover the original tokens.
Architecture: An encoder-only Transformer with rotary positional encoding and time conditioning.

C. Novel Training Strategies

To address the specific challenges of audio inpainting, the authors introduce two key innovations:

Span-Based Masking: Unlike standard discrete diffusion which masks tokens independently, AIDD employs structured span masking. It samples contiguous spans of tokens to mask, simulating realistic gaps. The span length follows a geometric distribution, allowing the model to learn from both fine-grained local corruption and broader semantic perturbations.
Derivative-Based Regularization Loss: To ensure temporal smoothness in the reconstructed audio, a regularization loss ( $\mathcal{L}_{deriv}$ ) is added. This loss penalizes irregular fluctuations in the predicted token embeddings by aligning the first- and second-order derivatives of the predicted embeddings with the ground truth. This enforces smooth temporal dynamics across token boundaries.

3. Key Contributions

First Discrete Diffusion for Audio Inpainting: AIDD is the pioneering work applying discrete diffusion to tokenized audio, enabling stable generation and long-range semantic coherence that continuous models struggle to achieve.
Unified Corruption-Reconstruction Framework: The combination of span-based masking (for structured corruption) and derivative-regularized reconstruction (for smoothness) creates a robust training paradigm specifically tailored for audio gaps.
Efficiency and Scalability: The model achieves state-of-the-art results with significantly fewer parameters and training time compared to strong baselines like CQT-Diff+.

4. Experimental Results

The method was evaluated on two benchmark datasets: MusicNet (classical music) and MAESTRO (piano performances), with gap lengths ranging from 150 ms to 750 ms.

Objective Metrics:
- MusicNet: AIDD outperformed baselines (LPC, A-SPAIN-L, CQT-Diff+) across most gap lengths. Notably, it achieved a ~25% reduction in Fréchet Audio Distance (FAD) at 300 ms gaps compared to CQT-Diff+, indicating superior perceptual realism.
- MAESTRO: AIDD achieved the best Objective Difference Grade (ODG) scores for 375 ms and 750 ms gaps, outperforming GAN-based methods (GACELA, bin2bin) and their MIDI-extended variants.
Subjective Evaluation:
- In Mean Opinion Score (MOS) tests on MAESTRO, AIDD scored 3.64, outperforming GACELA (3.51) and CQT-Diff+ (3.51), and approaching the quality of original audio (4.12).
Efficiency:
- AIDD is approximately half the size of CQT-Diff+ (90M vs. 242M parameters) and required only one-quarter of the training time.
- It supports variable-length inputs, unlike many baselines that require fixed-size inputs.

5. Significance and Future Directions

Paradigm Shift: This work demonstrates that discrete diffusion in a tokenized latent space is a superior approach for long-gap audio restoration compared to continuous waveform or spectrogram diffusion. It effectively bridges the gap between local signal fidelity and global semantic structure.
Generalizability: The proposed training strategies (span masking and derivative regularization) offer a new direction for training discrete diffusion models not just for audio, but potentially for other token-based domains like language.
Limitations: The performance is bounded by the quality and bandwidth (24 kHz) of the underlying tokenizer. Additionally, there is a potential training-inference mismatch regarding how gaps are handled during tokenization.

In conclusion, AIDD establishes a new state-of-the-art for audio inpainting, particularly for challenging, long-duration gaps, by leveraging the structural advantages of discrete token representations and tailored diffusion training objectives.