DCB — Differentiable Critical Bandwidth

A PyTorch package that makes Silverman's critical bandwidth test (1981) fully differentiable, enabling end-to-end gradient-based optimisation over the modal structure of continuous distributions.

Overview

h_crit is the minimum KDE bandwidth at which a distribution appears unimodal — a classical nonparametric statistic for modality testing. DCB replaces every non-differentiable step with a smooth surrogate, then uses the Implicit Function Theorem (IFT) to compute exact gradients through the root-finding step at O(1) memory cost.

import torch
from dcb import DCBLayer, TrainingLayer

X = torch.randn(100_000, requires_grad=True)  # 1D samples, any n from 1K to 1B
layer = DCBLayer()
h_crit = layer(X)       # differentiable scalar
h_crit.backward()       # exact IFT gradients

# For repeated training-loop use with warm-start bracket caching:
layer = TrainingLayer(warm_start=True)
for batch in dataloader:
    h = layer(batch)    # 1.8× faster after first call on CPU; ~10× on CUDA with compile=True

Installation

pip install diffcb

Or from source:

git clone https://github.com/ryZhangHason/differentiable-critical-bandwidth
cd differentiable-critical-bandwidth
pip install -e ".[dev]"

Accuracy vs R's bw.crit

Validated against R's multimode::bw.crit(data, mod0=1) (Hall & York 2001). Same-sample protocol (identical data fed to both Python and R, 3–5 seeds per n):

¹ n > 1M uses the sketch API (default max_n_exact=1_000_000); accuracy is measured on the 1M sketch vs R on the same sketch. Bimodal data interleaved so every slice is 50/50.

Distribution note: Bimodal = 0.5·N(−2,1) + 0.5·N(+2,1), h_crit ≈ 1.7. Gaussian = N(0,1), h_crit ~ n^{−1/5} (shrinks with n). Independent-sample error (~0.2–0.5%) reflects sampling variability, not algorithmic error.

Hardware Performance (v0.1.10, CPU Apple M-series)

¹ FFT-normalised cost = t_dcb / t_fft_one, where t_fft_one = median time for one rfft(16384, float32) on the same device (0.062 ms, Apple M). Hardware-agnostic metric for cross-chip comparison.

GPU (Kaggle P100, Round 24): 43–116× speedup vs CPU. Peak 116× at n=50K (direct-KDE parallelism).

API Reference

DCBLayer

DCBLayer(
    target_modes=1,            # target number of modes (default 1 = unimodality test)
    use_fft=True,              # FFT histogram path for n > 50K (default True)
    max_n_exact=1_000_000,     # sketch above this n; None = always exact
    sketch_size=500_000,       # target size for the random sketch subsample
    G_min=16384,               # minimum FFT histogram bins; scales with n^0.4 internally
    use_richardson=True,       # Richardson extrapolation (~30% accuracy gain on CPU)
    forward_path='smooth',     # 'smooth' (default, strictly differentiable, gradcheck passes)
                               # 'auto'   (direct-KDE at n≤25K; surrogate gradient)
                               # 'direct' (force direct-KDE; accuracy benchmarks only)
    fft_dtype=torch.float32,   # histogram FFT dtype for speed; mode counting always uses
                               # float64 internally regardless of this setting
    direct_n_max=25_000,       # direct-KDE threshold (only active in 'auto'/'direct')
    direct_M=2048,             # direct-KDE evaluation grid size
)

forward_path='smooth' (default): Both forward and backward use the same smooth M̃ surrogate at all n. torch.autograd.gradcheck passes. Gradients are exact IFT derivatives of the computed h_crit.

forward_path='auto': Uses direct-KDE for the forward pass at n ≤ 25K (zero binning bias) but the IFT backward always uses M̃. Forward and backward use different implicit functions at n ≤ 25K — gradcheck fails. Use only for forward-only benchmarks.

TrainingLayer (for ML training loops)

from dcb import TrainingLayer

layer = TrainingLayer(
    warm_start=True,    # cache h_prev; narrow bracket to [0.95h, 1.05h] → 1.8× speedup
    compile=False,      # torch.compile opt-in (requires float32, Python ≤ 3.11 on CUDA)
    warm_margin=0.05,
    **dcb_kwargs,       # any DCBLayer parameter
)
layer.reset_cache()     # call when distribution shifts significantly

find_signal_direction (critical-bandwidth signal recovery)

For recovering the direction of a real, population-level critical-bandwidth (multimodal) signal in point-cloud data — as opposed to using DCBLayer as a differentiable loss term inside gradient-based training — use the validated cross-fitted, elitist evolution search:

from dcb.signal_recovery import find_signal_direction

X = ...  # (n, d) float64 tensor
q, diagnostics = find_signal_direction(X, budget=2000, seed=0)
# q: (d,) unit vector -- the recovered signal direction

This is NOT gradient-based (gradient ascent on the critical-bandwidth index was found to be consistently anti-aligned with signal recovery — ascending unimodality is not the same objective as recovering structure). Instead it scores each candidate direction by cross-fitting the scale-free index h~(q) = h_crit(Xq)/sd(Xq) across 4 disjoint data folds (min aggregation), which kills spurious per-sample optima while a real population-level signal survives every fold, then returns the single best-scoring candidate from an 8-member evolution strategy.

Validated in a fully pre-registered confirmatory gate (n=20 fresh, disjoint seeds, decided before any data existed): median recovered cosine 0.86, beating random search (p=0.012, bootstrap CI lower +0.067), SPSA (p<0.0001, CI lower +0.300), and CMA-ES (p=0.0017, CI lower +0.065). Zero delivery failures across all 20 seeds. See the project's Round 48 Phase 9–14 experiment log for the full derivation and gate results.

This construction is validated specifically for data where a real critical-bandwidth signal exists along one direction (sharp, dense components riding on a diffuse background); it is not a general-purpose direction-finder and makes no claim on data without that structure.

Gradient quality check

from torch.autograd import gradcheck
layer = DCBLayer(forward_path='smooth')
X = torch.randn(200, dtype=torch.float64, requires_grad=True)
assert gradcheck(layer, (X,), eps=1e-5, atol=1e-4)  # passes with default settings

Richardson extrapolation

By default (use_richardson=True), DCB runs a second bisection at G/2 bins and applies h̃ = (4·ĥ(G) − ĥ(G/2)) / 3, reducing error ~30% on CPU. Both passes use float64 histogram FFTs for mode counting. On GPU this adds ~38% overhead with minimal gain; use use_richardson=False for GPU training loops.

Numerical Design Notes

Float64 mode counting: The bisection trisection and Richardson loops use float64 histogram FFTs (precompute_fft(..., fft_dtype=float64)) for mode counting, regardless of fft_dtype. Float32 histogram FFT errors (~10⁻⁴ relative) create spurious sign changes in f′ near small h_crit (unimodal distributions, h_crit < ~0.5) that cannot be filtered without suppressing genuine small lobes. Float64 reduces this floor to ~10⁻¹⁶, which the 10⁻¹² relative threshold cleanly removes. Float32 is still used for _refine_hcrit refinement evals and speed-sensitive paths.

n-aware G_min: Effective grid size = min(G_min × (n/100K)^0.4, 262144), keeping h_crit / bin_width stable as h_crit shrinks with n (Gaussian regime: h_crit ~ n^{−1/5}).

Known Limitations

• 1D only. No multivariate extension implemented.
• Gaussian max error ~0.40% for ~5–10% of seeds at n=100K. Arises when the disappearing lobe height approaches the 10⁻¹² noise threshold just below h_crit, causing the bisection to converge slightly early and _refine_hcrit to select a neighbouring lobe. Bimodal data (large h_crit) is unaffected.
• MPS float64. The float64 histogram precompute falls back to CPU on MPS devices; MPS does not support float64 FFT. Handled automatically.
• Independent-sample variability (~0.5%). Two independent random draws from the same distribution have h_crit differing by ~0.2–0.5% at n=100K — irreducible MC variance, not algorithmic error.
• torch.compile blockers. MPS: blocked by float64 fallback. CUDA: requires Python ≤ 3.11 or torch ≥ 2.4.

Confirmed Experimental Results

Version History

Repository Structure

dcb/
  layer.py            DCBLayer nn.Module + DCBFunction autograd
  solver.py           IFT root-finder, trisection bisection, Richardson pass
  fft_kde.py          FFT mode counter, direct_mode_count_batch, precompute_fft
  training.py         TrainingLayer with warm-start and compile support
  kde.py              Direct KDE derivatives (IFT backward path)
  utils.py            Grid, Silverman bandwidth, sg() stabiliser
  signal_recovery.py  find_signal_direction: cross-fitted, elitist evolution search
  _fold_search.py      internal: fold-finding machinery behind signal_recovery
tests/           Unit tests (57 passed, 1 xfailed)

Citation

@software{zhang2026diffcb,
  author = {Zhang, Ruiyu},
  title  = {diffcb: Differentiable Critical Bandwidth for PyTorch},
  year   = {2026},
  url    = {https://github.com/ryZhangHason/differentiable-critical-bandwidth},
  note   = {PyPI: diffcb v0.1.10}
}

License

Apache 2.0 — see LICENSE.