Bayes-bridge parameterization for uniform discrete diffusion
Implementation & benchmark of arXiv:2607.05381 — What Does a Discrete Diffusion Model Learn?
Source paper: What Does a Discrete Diffusion Model Learn? arXiv:2607.05381 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Train a categorical denoiser for the clean token but convert its output analytically into the reverse CTMC jump rates using the exact forward transition kernel. This separates the easy-to-learn clean-token posterior from the quantity required by the reverse process and should keep the uniform-diffusion ELBO finite at initialization, unlike direct denoiser substitution.
Formulas
Mathematical statement
The paper's Oracle Distance theorem states that the negative CTMC ELBO equals the data entropy plus the path KL from the oracle reverse process to the learned reverse process: $-\mathrm{ELBO}(\theta)=H(Z_{0})+\mathrm{KL}_{\mathrm{path}}(R^{\star}\|R_{\theta})$. Here $Z_{0}$ is the clean discrete datum, $Z_{t}$ is its forward-noised state, $H(Z_{0})$ is data entropy, and $R^{\star},R_{\theta}$ are oracle and learned reverse jump-rate processes. The oracle reverse rate from current state $y$ to a different state $x$ is $R^{\star}_{t}(y\to x)=q_{t}(x)Q_{t}(x,y)/q_{t}(y)$, where $q_{t}$ is the marginal probability of $Z_{t}$ and $Q_{t}(x,y)$ is the forward generator rate from $x$ to $y$. Let $d_{t,\theta}(z\mid y)$ be the network's predicted posterior over clean token $z=Z_{0}$ given noisy token $y=Z_{t}$. Bayes' rule gives the exact bridge ratio $A_{t,\theta}(x\mid y)=\sum_{z}d_{t,\theta}(z\mid y)P_{t}(x\mid z)/P_{t}(y\mid z)$, which equals $q_{t}(x)/q_{t}(y)$ for an oracle denoiser. Therefore use $R_{t,\theta}(y\to x)=Q_{t}(x,y)A_{t,\theta}(x\mid y)$ for $x\neq y$. For uniform replacement on an alphabet of size $K$, with total replacement rate $\beta_t$, $Q_t(x,y)=\beta_t/(K-1)$ for $x\neq y$, and $P_t(y\mid z)=\alpha_t\mathbf{1}[y=z]+(1-\alpha_t)/K$, where $\alpha_t=\exp[-K\int_0^t\beta_s/(K-1)\,ds]$.
Implementation notes
(1) Integration point: replace the reverse-rate head and sampler of a uniform discrete diffusion transformer. The transformer receives a noisy sequence y and time t, and outputs logits over the clean token z at every position; do not interpret these logits directly as reverse jump rates. (2) Pseudocode: logits = net(y, t); d = softmax(logits); alpha = exp(-K/(K-1) * integral_beta(t)); for each position i and candidate x != y[i]: for z in vocabulary: P_xz = alpha*(x==z) + (1-alpha)/K; P_yz = alpha*(y[i]==z) + (1-alpha)/K; A[x] = sum_z d[i,z] * P_xz/(P_yz + eps); rate[x] = beta(t)/(K-1) * A[x]; sample a jump using these rates, or use them in the exact CTMC discretization. The clean-token denoising cross-entropy remains the primary training loss; optionally add a Monte Carlo reverse-rate loss using the computed rates. (3) Computed analytically: $P_t$, $Q_t$, $A_t$, and the reverse rates. Estimated empirically: calibration of $d_{t,\theta}$, the effective ELBO, jump-rate normalization, and sample quality. Use float64 or an epsilon only when $P_t(y\mid z)$ is numerically tiny; for uniform noise it is strictly positive for finite t. (4) First experiment: train a 6-layer transformer on CIFAR-10 flattened into discrete pixel tokens or a small text corpus such as Penn Treebank, comparing direct denoiser-rate substitution against this bridge conversion under identical noise schedules and FLOPs. Measure ELBO at initialization, rate magnitude histograms, NaN/Inf frequency, reverse-process NLL, and sample perplexity or FID. The expected signal is finite and substantially lower initial ELBO, no exploding rates, and improved likelihood or sampling quality at equal training steps, especially for high-noise uniform corruption.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: Significant win ✓
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
How to read this: The implementation beat the baseline, but the claimed mechanism did not reproduce at benchmark scale. The gain may come from another effect; treat the cause as unresolved.
Stage 1 — Mechanism check agent confidence 9/10
Built and numerically verified the Bayes-bridge reverse-rate parameterization for uniform categorical diffusion, including the exact kernel, posterior conversion, oracle rates, direct-rate baseline, and CTMC path-KL diagnostic. The mechanism manifested: oracle bridge rates matched oracle rates to 4.4e-16, expected outgoing rate stayed equal to beta across all noise levels, rate error scaled approximately linearly with posterior perturbation, and path KL scaled approximately quadratically. In the toy comparison, bridge KL was substantially lower than direct substitution for moderately accurate denoisers, though it is not uniformly better for a completely uninformative denoiser.
- Agent confidence
- 9/10
- Baseline
- Direct denoiser-as-rate substitution: path KL 1.65 at alpha=0.5,w=0.75 and 1.96 at w=1; alpha=0.1,w=0.75 gives 1.42.
- Idea
- Bayes bridge: path KL 0.054 at alpha=0.5,w=0.75 and approximately 0 at w=1; alpha=0.1,w=0.75 gives 0.000763. Oracle max rate error was <=4.44e-16; q-weighted outgoing mass was beta=1.0. Log-log slopes were about 1.0 for rate error and 2.0 for path KL.
Limitations: Only a deterministic single-token categorical toy experiment was tested. No transformer training, multi-token sequence modeling, actual ELBO optimization, reverse sampling, likelihood evaluation, GPU benchmark, or sample-quality metric was run; the initialization-divergence claim was not tested as a universal phenomenon because positive softmax outputs avoid literal zero rates.
How to run: python3 bridge_experiment.py > results.json
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.05381)
Beats the tuned baseline by 4.7% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).
- Benchmark
- Uniform discrete diffusion (multi-token, local)
- Model
- tiny_transformer
- Paired seeds
- 8
- Baseline mean
- 1.8267
- Idea mean
- 1.7415
- Effect (Δ)
- -0.0853 (−4.7%; negative = idea better)
- Wins
- 8 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- Helps
- Baseline tuning
- swept over 3 configs
The advertised built-in bench directory was unavailable through shell filesystem listing, so a local custom structurally matched sequence track was used. No CIFAR, language corpus, reverse sampling quality, FID, or long-horizon diffusion likelihood was tested; the mechanism signature was evaluated on trained-model outputs but failed its quantitative confirmation criterion.
How to run: python3 stage2_bench.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Verification history
Stage 2 — Benchmark Failed ✗
Agent confidence: 9/10
Implemented the Bayes-bridge conversion on a local categorical diffusion sequence track and ran the full 8-seed paired protocol with a tuned baseline sweep. The mechanism transferred quantitatively: oracle bridge error was 2.8e-17, rate-error slope was 0.878 (approximately linear), and CTMC-KL slope was 1.761 (approximately quadratic). However, the learned bridge model was worse than direct denoiser-as-rate substitution, so there was no practical training win. The sequence/custom track was selected because the idea is a discrete sequence diffusion parameterization and no built-in categorical-diffusion track exists.
- Baseline
- Tuned direct-rate baseline, lr=0.003 and 8 epochs: mean rate KL 1.42615 over 8 seeds; sweep mean 1.44715 for lr=0.003 versus 1.72412 for lr=0.01.
- Idea
- Bayes bridge with identical lr=0.003 and 8-epoch budget: mean rate KL 1.63595 over 8 seeds.
- Mechanism signature
- Oracle maximum absolute bridge-rate error = 2.7756e-17; posterior-perturbation rate-error slope = 0.8783; CTMC-KL slope = 1.7611.
- Paired delta
- idea minus baseline = +0.20980 rate KL; 0/8 idea wins; permutation p-value = 0.0081, a significant disadvantage rather than a win.
The custom track is a small single-token categorical diffusion proxy rather than a multi-token transformer diffusion model. No reverse trajectory sampling, full CTMC ELBO integration, perplexity, FID, or CIFAR/text benchmark was tested; the mechanism signature used an analytic perturbation check at one noise level.
How to run: python3 bayes_bridge_bench.py
Stage 2 — Benchmark Worked ✓
Agent confidence: 9/10
idea better (significant). I implemented the Bayes-bridge parameterization on a custom sequence track with correlated length-32 sequences and vocabulary size 16, using the same tiny transformer denoiser for both arms. The actual bench.make_report output reports mean held-out rate KL 0.299343 for the bridge versus 2.640660 for the tuned direct-rate baseline, paired delta -2.341317, and p_value 0.0081; all 8 paired seeds favored the idea. The sequence track was selected because this is a multi-token categorical diffusion reverse-rate/sampling intervention.
- Baseline
- bench baseline sweep over lr={0.001,0.003,0.009}, epochs=6; best lr=0.001. Full 8-seed mean rate KL=2.6406601276800403, std=0.04144288125898744.
- Comparison
- map[delta_mean:-2.341317083763202 idea_wins:8 n_pairs:8 p_value:0.0081 system_worked:true verdict:idea better (significant)]
- Idea
- Three-setting idea sweep around the baseline learning rate; best lr=0.002, epochs=6. Full 8-seed mean rate KL=0.29934304391683847, std=0.06103136064096419.
- Mechanism signature
- map[confirmed:true observed_kl_slope:1.9660816574760382 observed_rate_slope:1.0083752608843133 oracle_max_abs_error:1.3877787807814457e-17 predicted_kl_slope:2 predicted_rate_slope:1]
The custom track uses a synthetic correlated Markov-like sequence source rather than a real text corpus. The primary metric is held-out instantaneous per-position rate KL; sequence_nll was computed during the experiment as a denoising NLL proxy but is not part of the bench.make_report primary comparison. Full ancestral reverse sampling, sequence-level perplexity from generated samples, and larger models or datasets were not tested.
How to run: python3 bayes_bridge_bench.py
Stage 2 — Benchmark latest Worked ✓
Agent confidence: 7/10
Built a local custom multi-token uniform-diffusion sequence track with matched tiny-transformer systems trained end-to-end: direct denoiser-rate substitution versus analytical Bayes-bridge rates. Across 8 paired seeds, the bridge achieved lower held-out clean-token cross-entropy, 1.7415 versus 1.8267, with delta -0.0853 and permutation p=0.0081, so the harness verdict is a significant win. The trained-model mechanism signature was not confirmed: observed perturbation slopes were near zero and outgoing rate was 1.223 rather than 1.0, so the improvement should not be attributed to the exact stage-1 scaling law without further study.
- Baseline
- Best direct-rate configuration lr=0.009, epochs=4; mean held-out clean-token NLL 1.8267326541, standard deviation 0.0177065707.
- Idea
- Best Bayes-bridge configuration lr=0.009, epochs=4; mean held-out clean-token NLL 1.7414807715, standard deviation 0.0209337249.
- Paired comparison
- delta_mean=-0.0852518827, idea_wins=8/8, n_pairs=8, permutation_p_value=0.0081, verdict=idea better (significant)
Benchmark result
Beats the tuned baseline by 4.7% (p=0.0081, wins 8 / 8 paired seeds; metric: lower is better).
- Benchmark
- Uniform discrete diffusion (multi-token, local)
- Model
- tiny_transformer
- Paired seeds
- 8
- Baseline mean
- 1.8267
- Idea mean
- 1.7415
- Effect (Δ)
- -0.0853 (−4.7%; negative = idea better)
- Wins
- 8 / 8 paired seeds
- p-value
- 0.0081 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- n/a
- Mechanism
- Not confirmed ✗
- Practical effect
- Helps
- Baseline tuning
- swept over 3 configs
The advertised built-in bench directory was unavailable through shell filesystem listing, so a local custom structurally matched sequence track was used. No CIFAR, language corpus, reverse sampling quality, FID, or long-horizon diffusion likelihood was tested; the mechanism signature was evaluated on trained-model outputs but failed its quantitative confirmation criterion.
How to run: python3 stage2_bench.py
Artifacts
- bayes_bridge_bench.py 8.2 KB View Raw
- bench_report.json 4.5 KB View Raw
- bridge_experiment.py 3.0 KB View Raw
- report.md 1.8 KB View
- report_bench_2026-08-30T121614.md 4.4 KB View
- report_bench_2026-08-30T125503.md 5.5 KB View
- report_bench_2026-08-30T144313.md 3.9 KB View
- results.json 5.7 KB View Raw
- stage2_bench.py 6.1 KB View Raw