03 / RESEARCH ▸ scroll 000%

RESEARCH ·NeurIPS 2024

Meta-DiffuB

為每個句子排一條專屬的加噪時間表

A contextualized sequence-to-sequence text-diffusion model with meta-exploration: a scheduler network learns a per-sentence noise schedule, trained by the generation quality it produces — state-of-the-art and plug-and-play.

code ↗ interactive — scroll to step through the method
▸

Background

New here? Start with what a diffusion model is and why text is the hard case.

Diffusion models generate by starting from pure noise and removing it a little at a time until something coherent emerges — the idea behind the recent wave of image generators. Applying it to text, where the output is a sequence of discrete words rather than continuous pixels, is newer and trickier.

A diffusion model is steered by a noise schedule: how much noise to add, then remove, at each step. Standard text-diffusion models use one fixed schedule for every sentence — but sentences differ, some easy and some hard, so a single schedule is a compromise. Meta-DiffuB's idea is to learn a per-sentence schedule: a small scheduler network reads the sentence and decides how much noise it needs, and it is trained by the quality of the text the main model then produces — a meta-exploration loop where the generator's success teaches the scheduler.

▸ key terms

Diffusion model — generates data by repeatedly denoising. Noise schedule — the per-step plan of how much noise to add/remove. Seq2seq — maps an input sentence to an output sentence (e.g. translation, summarization). Meta-exploration — learning how to guide/explore — here the scheduler learns from the generator's results.

METHOD · scheduler–exploiter
wˣ Scheduler Bψ · meta-explorer βˣ contextualized noise Exploiter Dθ · S2S diffusion diffusion step t → βₜ noise √-schedule one fixed schedule · every sentence treated alike ι = Bψ(wˣ) — Meta-Instructions T F T T F T F ‘skipping’ · T → step noise up · F → hold level βˣ = skipping(ι, β√) · trained by policy gradient diffusion step t → easy → more noise hard → less noise βˣ adapts per sentence — the key move reverse diffusion with βˣ zT ~ N(0,I) z₀ → ŷ generating rounding → discrete target sentence Dθ → BLEU r before update Dθ′ → BLEU r′ after update R_β = r′ − r meta-reward ∇ψ J(ψ) · policy gradient → update Scheduler STATE OF THE ART 4 / 4 Seq2Seq benchmarks · vs DiffuSeq, SeqDiffuSeq, Dinoiser, PLMs CC QT WA QQP scheduler = plug-and-play · no fine-tuning at inference
DiffuSeq imposes one fixed noise schedule on every sentence — regardless of how hard it is to generate.
STEP 00 · the problem

One schedule for every sentence

Standard S2S-diffusion (DiffuSeq) adds noise on a fixed √-schedule. A trivial paraphrase and a hard open-domain reply are corrupted exactly the same way.

But sentences differ in difficulty — non-contextualized noise leaves performance on the table.

STEP 01 · the scheduler

A network that reads the sentence

The Scheduler Bψ — a small Seq2Seq model — reads the conditioning sentence wˣ and emits a sequence of Meta-Instructions ι, each labelled True or False.

A ‘skipping’ rule turns those labels into noise: True steps the noise up, False holds it.

ι = Bψ(wˣ) → βˣ = skipping(ι, β√)
STEP 02 · contextualized noise

Less noise for hard sentences

The result is a per-sentence schedule βˣ that bends away from the fixed baseline. Harder sentences get less noise to preserve signal; easier ones get more to boost diversity.

This is the move that non-contextualized schedulers can't make.

STEP 03 · the exploiter

Generate with the scheduled noise

The Exploiter Dθ — the S2S-diffusion model — diffuses and denoises using βˣ, recovering z₀ step by step, then rounds it back into a discrete target sentence ŷ.

STEP 04 · meta-reward

Generation quality teaches the scheduler

How much did the exploiter improve? Compare BLEU before and after an update: the meta-reward R_β = r′ − r flows back through a policy gradient to train the scheduler.

The scheduler learns how to noise — never touching the generator's loss directly.

STEP 05 · the result

State of the art, plug-and-play

Meta-DiffuB sets a new bar across four Seq2Seq benchmarks, beating prior diffusion models and fine-tuned PLMs.

Better still, the trained scheduler drops into existing models like DiffuSeq as a plug-and-play module — no fine-tuning required.