Conditional-TC Reveal Scheduler
Implementation & benchmark of arXiv:2608.25505 — Conditional Total Correlation and the Serial Depth of Adaptive Parallel Sampling
Source paper: Conditional Total Correlation and the Serial Depth of Adaptive Parallel Sampling arXiv:2608.25505 ⓘ · analyzed Aug 29, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace confidence-only masked diffusion decoding with an adaptive scheduler that chooses batches whose unrevealed tokens have low conditional total correlation given the already revealed context. The scheduler should preserve large parallel batches when token predictions are conditionally independent, but split highly dependent tokens into separate rounds to reduce forward-KL error.
Formulas
Mathematical statement
For a target distribution P(X_1,\ldots,X_n), let C be the set of already revealed coordinates with realized values x_C, and let A be the coordinates revealed in the next round. Conditional total correlation is \mathrm{TC}(X_A\mid x_C)=\sum_{i\in A}H(X_i\mid x_C)-H(X_A\mid x_C), where H is Shannon entropy in bits. If the decoder samples the next batch from the product of exact conditional marginals q_A(x_A\mid x_C)=\prod_{i\in A}P(x_i\mid x_C), then the one-round forward KL contribution is exactly D_{\mathrm{KL}}(P(X_A\mid x_C)\|q_A(\cdot\mid x_C))=\mathrm{TC}(X_A\mid x_C). For an adaptive reveal policy with rounds A_1,\ldots,A_T, the total error is therefore the target expectation \mathbb{E}[D_{\mathrm{KL}}(P\|Q_\pi)]=\mathbb{E}_{X\sim P}[\sum_{t=1}^{T}\mathrm{TC}(X_{A_t}\mid X_{C_t})], where C_t=\cup_{s<t}A_s. The implementation estimates this cost using the neural model's conditional marginal entropies and Monte Carlo estimates of the joint conditional entropy; exact zero conditional TC indicates a batch that can be sampled in parallel without KL penalty.
Implementation notes
(1) Integrate this into the masked diffusion sampler immediately before each denoising/reveal step. The current mask defines U, the revealed tokens define C, and the scheduler must choose the next batch A subset of U. Retain the model's logits for each masked position as estimates of p_theta(X_i|x_C). (2) Use a two-stage greedy scheduler. First rank U by confidence or entropy and form candidate batches of sizes B, B/2, ..., 1. For each candidate A, compute marginal entropy H_i=-sum_v p_i(v)log2 p_i(v). Estimate the joint entropy with K stochastic completions or masked-token samples x_A^(k) from the model: H_joint=-(1/K)sum_k log2 p_theta(x_A^(k)|x_C). If the architecture cannot score a joint assignment, train or attach a small pairwise dependence head g_theta(i,j,C), and use the computable approximation H_joint≈sum_i H_i-sum_{i<j}g_theta(i,j,C), with g clipped nonnegative. Compute TC_hat=sum_i H_i-H_joint. Choose the largest candidate batch whose TC_hat is below a user budget epsilon_t; otherwise choose the minimum-cost candidate or a singleton. Reveal the chosen tokens using either sampled values or argmax values, update C, and repeat. (3) The paper supplies the identity and the additive round cost; K, the candidate batch sizes, epsilon_t, and the joint-entropy estimator are engineering choices. Estimate uncertainty by repeating the TC estimate with independent completion samples and penalize uncertain estimates using TC_upper=TC_hat+2*std(TC_hat). (4) First test on a small masked diffusion language model such as a 12-layer Transformer on WikiText-103 or a synthetic autoregressive process with known conditional dependencies. Compare fixed-size parallel decoding, confidence-based decoding, and TC scheduling at equal denoising forward passes. Measure generated-token NLL or KL against an autoregressive teacher, quality at a fixed number of rounds, and average batch size at a fixed quality target. Success is a lower NLL/KL than confidence decoding at the same number of rounds, or fewer model evaluations to reach the same perplexity; on synthetic Markov data the scheduler should discover that tokens separated beyond the Markov order can be revealed together.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not tested
Stage 2 · Practical benchmark result: Not run
Methodology: Toy-system gate first; the benchmark stage runs only after a pass. How verification works
Stage 1 — Mechanism check agent confidence 9/10
Built an exact synthetic MVP for conditional-TC scheduling on four independent, strongly correlated binary pairs. The core identity matched numerically to 1.1e-16 bits, and the scheduler correctly measured positive TC for same-pair batches and zero TC across independent pairs. It reduced accumulated forward KL from 2.8544 bits for confidence batching to approximately 0 at the cost of increasing average rounds from 2 to 5, so the accuracy effect is real but this experiment does not demonstrate a speedup.
- Agent confidence
- 9/10
- Baseline
- Confidence batching: 2.854412 forward-KL bits, 2.0 mean rounds, mean batch size 4.0
- Idea
- TC scheduler: 1.01e-15 forward-KL bits, 5.0 mean rounds, mean batch size 1.6; same-pair TC=0.713603 bits and cross-pair TC=0 bits
- Math check
- TC=0.713603042884 bits, direct product-marginal KL=0.713603042884 bits, absolute error=1.11e-16
Limitations: This is an exact enumerable synthetic distribution, not a neural masked-diffusion model. It uses exact conditional entropies rather than noisy model logits or Monte Carlo joint-entropy estimates, does not test GPU performance, and does not establish a speed win at equal quality or equal forward-pass count.
How to run: python3 tc_scheduler_experiment.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Artifacts
- report.md 1.5 KB View
- tc_scheduler_experiment.py 4.5 KB View Raw