# Tail-Weighted Optimal Batch Scheduling

- ID: 1504
- Canonical URL: https://synthcore.org/idea/1504/tail-weighted-optimal-batch-scheduling
- API JSON: https://synthcore.org/api/idea/1504.json
- API Markdown: https://synthcore.org/api/idea/1504.md
- Verification status: failed_benchmark
- Source: [arXiv:2607.27731](https://arxiv.org/abs/2607.27731)
- Category: optimization
- Solves: speedup, sample-efficiency, accuracy
- ML areas: optimizer, batching, scheduler, transformer
- Math tags: convex-analysis, optimization, probability
- Ratings: usefulness 7/10; difficulty 5/10; novelty 5/10

## 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.

## 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}\).

## Key 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}}$$

## 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

- Status: failed_benchmark
- Mechanism evidence: yes
- Mechanism confirmed: no
- Practical verdict: harms
- Verdict: 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.

### Mechanism check

- Verdict: 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.
- Confidence: 8/10
- 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.

### Practical benchmark

- Paired seeds: 8
- Baseline mean: 0
- Idea mean: 0
- p-value: 0.0081
- Paired wins: 0/8
- Benchmark verdict: idea worse (significant)

## Artifacts

- [bench_report.json](https://synthcore.org/code/587/bench_report.json)
- [report.md](https://synthcore.org/code/587/report.md)
- [report_bench_2026-08-31T192201.md](https://synthcore.org/code/587/report_bench_2026-08-31T192201.md)
- [results.json](https://synthcore.org/code/587/results.json)
- [stage2_bench.py](https://synthcore.org/code/587/stage2_bench.py)
- [tail_batch_experiment.py](https://synthcore.org/code/587/tail_batch_experiment.py)
- [Download all files as ZIP](https://synthcore.org/download/587)

## Disclaimer

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