✗ Failed on benchmark 2026

Tail-Weighted Optimal Batch Scheduling

Implementation & benchmark of arXiv:2607.27731 — Towards joint scaling laws with optimal batch size schedules

Usefulness7/10

Difficulty5/10

Novelty5/10

Optimization Speedup Sample efficiency Accuracy Optimizer Batching Scheduler Transformer Convex analysis Optimization Probability

Source paper: Towards joint scaling laws with optimal batch size schedules arXiv:2607.27731 ⓘ · analyzed Aug 31, 2026

AI-generated research hypothesis, automatically tested. Not peer-reviewed.

Idea description

Replace a static batch size with a schedule optimized for a prescribed learning-rate schedule and a fixed total number of processed examples. Steps whose stochastic-gradient noise has a large effect on the paper's loss bound receive larger batches, with the weighting determined by the future learning-rate tail rather than by a hand-designed warmup or cooldown rule.

Formulas

$$\mathbb{E}L(\mathbf{w}_{\tau};\eta_t)\leq L_*+\frac{D^2}{2\sum_{t=1}^{\tau}\eta_t}+\frac{1}{2}\sum_{t=1}^{\tau-1}\frac{\eta_t^2\mathbb{E}\|\mathbf{g}_t\|^2}{\sum_{k=t+1}^{\tau}\eta_k}+\frac{1}{2}\eta_\tau\mathbb{E}\|\mathbf{g}_\tau\|^2$$

$$\mathbb{E}\|\mathbf g_t(B_t)\|^2\approx a_t+\frac{s_t}{B_t},\qquad c_t=\begin{cases}\displaystyle\frac{\eta_t^2}{2\sum_{k=t+1}^{\tau}\eta_k},&t<\tau,\\[6pt]\displaystyle\frac{\eta_\tau}{2},&t=\tau,\end{cases}$$

$$\min_{B_t>0,\,\sum_{t=1}^{\tau}B_t=C}\sum_{t=1}^{\tau}\frac{c_ts_t}{B_t}\quad\Longrightarrow\quad B_t^*=\frac{C\sqrt{c_ts_t}}{\sum_{j=1}^{\tau}\sqrt{c_js_j}}$$

Mathematical statement

Theorem 1 bounds the expected loss after $\tau$ SGD steps by $L_* + D^2/(2\sum_{t=1}^{\tau}\eta_t)$ plus a weighted sum of stochastic-gradient second moments: $\frac{1}{2}\sum_{t=1}^{\tau-1}\eta_t^2\mathbb{E}\|\mathbf g_t\|^2/(\sum_{k=t+1}^{\tau}\eta_k) + \frac{1}{2}\eta_\tau\mathbb{E}\|\mathbf g_\tau\|^2$. Here $L_*$ is the optimum loss, $D$ is the initial distance to an optimum, $\eta_t>0$ is the learning rate at step $t$, $\mathbf g_t$ is the stochastic gradient, and $\tau$ is the training horizon. Assume the averaged batch gradient satisfies $\mathbb E\|\mathbf g_t(B_t)\|^2\approx a_t+s_t/B_t$, where $B_t$ is batch size, $a_t$ is the batch-independent squared true-gradient component, and $s_t$ is the gradient-noise scale. Define $c_t=\eta_t^2/[2\sum_{k=t+1}^{\tau}\eta_k]$ for $t<\tau$, and $c_\tau=\eta_\tau/2$. Under the fixed example budget $\sum_t B_t=C$, minimizing the batch-dependent part $\sum_t c_ts_t/B_t$ gives $B_t^*=C\sqrt{c_ts_t}/\sum_j\sqrt{c_js_j}$.

Implementation notes

Integrate this into the data-loader and optimizer-step loop, not into the model architecture. Choose a learning-rate schedule $\eta_1,\ldots,\eta_\tau$, a total example budget $C$, minimum and maximum batch sizes, and a planning horizon such as 256 or 1024 optimizer steps. Precompute the tail coefficients $c_t$ from the theorem. At the beginning of each planning window, estimate gradient noise using temporary microbatches: compute per-microbatch gradients $g_{t,i}$, their mean $\bar g_t$, and $\hat s_t=(q-1)^{-1}\sum_{i=1}^q\|g_{t,i}-\bar g_t\|^2$, with $q$ probe microbatches. Set $r_t=\sqrt{c_t\max(\hat s_t,\epsilon)}$, then assign $B_t=\operatorname{round}_{\mathrm{valid}}(C_w r_t/\sum_j r_j)$ within the current window, where $C_w$ is that window's example budget. Enforce $B_{\min}\leq B_t\leq B_{\max}$, renormalize the remaining budget after clipping, and use gradient accumulation when device memory cannot hold $B_t$. Keep the learning rate indexed by optimizer step rather than by examples. Re-estimate noise every 32--128 steps and smooth $\hat s_t$ with an exponential moving average to prevent schedule jitter. For a cheap first experiment, train a 6-layer decoder-only Transformer on WikiText-103 or OpenWebText with AdamW and cosine decay, comparing static batch, linear batch growth, and this schedule under identical total examples. Record validation loss versus examples, wall-clock time, peak memory, gradient-noise estimates, and loss at equal FLOPs. Success means lower validation loss or faster time-to-target than static batching after data-loader and accumulation overhead; failure means no improvement at equal compute.

Verification

Failed on benchmark

Stage 1 · Toy mechanism gate: Passed ✓

Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗

Stage 2 · Practical benchmark result: Significantly worse ✗

Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works

Stage 1 — Mechanism check agent confidence 8/10

Built a NumPy MVP implementing the theorem’s tail coefficients, continuous optimum, bounded integer allocation, and a noisy-quadratic SGD comparison. The mechanism was numerically confirmed: allocation ratios matched sqrt(c_t s_t) to 2.2e-16, the continuous solution matched the predicted optimum exactly, and scaling noise by 0.25/1/4/16 scaled the minimized objective by exactly the same factors while leaving relative allocations unchanged. In the toy training run, tail weighting improved step-32 loss and the theoretical batch-dependent bound versus static batching, but had slightly worse final stochastic loss (0.00883 vs 0.00873), so this is mechanism-confirmed rather than a demonstrated end-to-end win.

Agent confidence: 8/10

Baseline: Static batch: final loss 0.0087314 ± 0.0006217 SE; step-32 loss 0.0292732; theoretical batch objective 0.0364224.
Idea: Tail-weighted integer batch: final loss 0.0088268 ± 0.0006200 SE; step-32 loss 0.0283912; theoretical batch objective 0.0326175 (10.44% lower than static); integer/continuous objective ratio 1.00043. Allocation was batch 26 at step 1 and 18 at step 64, with range 8–26.

Limitations: This was a small synthetic noisy strongly-convex quadratic, not a neural network or real data-loader experiment. Noise scales were prescribed rather than estimated from probe microbatches; wall-clock overhead, GPU memory, gradient accumulation, AdamW, validation loss, and real sample-efficiency were not tested. The tail-allocation identity is mathematically exact under the assumed second-moment model, while practical gains may depend on estimation error and integer/clipping constraints.

How to run: python3 tail_batch_experiment.py

Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.27731) Sep 2, 2026

Significantly WORSE than the tuned baseline (+42.0%, p=0.0081, wins 0 / 8 paired seeds).

Benchmark: Tabular regression (Friedman)
Model: mlp_tiny
Paired seeds: 8
Baseline mean: 6.9889
Idea mean: 9.9224
Effect (Δ): 2.9335 (+42.0%; negative = idea better)
Wins: 0 / 8 paired seeds
p-value: 0.0081 (permutation test, 20 000 shuffles)
Smallest detectable effect: ±10.2%
Mechanism: Not confirmed ✗
Practical effect: Hurts
Baseline tuning: swept over 6 configs

Limitations:

Only the structurally matched tabular track was tested; vision, sequence, and dynamics were not tested. Wall-clock speed, memory, and a rigorous per-step trained-model noise-to-weight correlation were not measured. The report is saved at bench_report.json.

How to run: python3 stage2_bench.py

Verdict computed by deterministic test code from paired-seed statistics — not by the language model.

Stage 2 — Benchmark Aug 31, 2026 latest Failed ✗

Agent confidence: 10/10

Implemented Tail-Weighted Optimal Batch Scheduling on the structurally matched tabular Friedman#1 regression track with the shared mlp_tiny architecture and a local variable-batch Adam loop. Across 8 paired seeds, the idea was significantly worse than the tuned fixed-batch baseline: paired delta +2.933495 MSE, permutation p=0.0081, and 0/8 wins. The mechanism signature is unconfirmed because the run did not quantitatively correlate per-step trained-model noise estimates with sqrt(c_t s_t).

Baseline: Tuned fixed-batch Adam, lr=0.006 and batch=32: 8-seed test MSE mean 6.988916, SD 0.402721.
Idea: Tail-weighted variable-batch Adam, lr=0.006 and per-epoch example budget 8000: 8-seed test MSE mean 9.922412; batch range 41–128, mean 62.5; paired delta +2.933495, p=0.0081.

Benchmark result