A paired A/B harness testing Johnson-SU normalization against per-neuron z-scoring for short-context neural activity forecasting on ZAPBench.
Result in one line: Johnson-SU gives a small but reliable MAE improvement at 200- and 2,000-neuron subsets, and becomes unstable by 10,000 neurons, where run-to-run variance exceeds the effect itself. No claim is made at full 71,721-neuron scale.
ΔF/F calcium traces are skewed and heavy-tailed. The ZAPBench trace baselines normalize per neuron with z-scoring, which assumes something closer to symmetry. Johnson-SU fits four parameters per neuron (two of them shape parameters governing skewness and tail weight), so it can absorb that asymmetry before the model sees the data. This repo asks whether doing so actually helps, and answers honestly.
Negative = Johnson-SU beat z-scoring. Three seeds per configuration at 200 and 2,000 neurons, and five at 10,000. The two additional 10,000-neuron seeds were run after the first three disagreed in sign, so the seed count at that scale is a response to the observed instability rather than a number fixed in advance; it is stated here for that reason.
| Neurons | Model | seed 0 | seed 1 | seed 2 | seed 3 | seed 4 | mean | sd |
|---|---|---|---|---|---|---|---|---|
| 200 | TSMixer | −0.33% | −0.91% | −0.66% | — | — | −0.63% | 0.29 |
| 200 | TiDE | −0.17% | −0.24% | −0.13% | — | — | −0.18% | 0.06 |
| 2,000 | TSMixer | −1.54% | −0.90% | −1.73% | — | — | −1.39% | 0.43 |
| 2,000 | TiDE | −0.27% | −0.26% | −0.31% | — | — | −0.28% | 0.03 |
| 10,000 | TSMixer | −0.54% | −0.36% | +1.20% | −2.27% | −1.56% | −0.71% | 1.32 |
16 of 17 paired runs favour Johnson-SU. The single exception is at the largest scale tested — and there individual runs range from −2.27% to +1.20% about a mean of −0.71%, a standard deviation of 1.32 against effect sizes well under 1% at smaller subsets. The correct reading is unreliable at 10k: not a persistent gain, and not a clean reversal either.
Raw per-run outputs are the *.json files in this repo, including z-MAE, Johnson-MAE, CorrH for both arms, window counts and Wilcoxon p-values.
Claimed: at short context (C=4), per-neuron Johnson-SU normalization reduces MAE relative to z-scoring at 200 and 2,000 neuron subsets, across three seeds and two architectures (five seeds at 10,000), under equal training budgets with no tuning for either arm.
Not claimed: anything at 71,721 neurons. Anything about the magnitude at 10,000 neurons — five seeds there span 3.5 percentage points, so no point estimate is meaningful. Anything about long context (C=256 was not run). Any absolute MAE comparable to the ZAPBench leaderboard — 1,500 training steps on a neuron subset with no early stopping is far below the official training budget, so these numbers are only meaningful as a paired delta between two normalizations of the same model.
Known oddity: at 10,000 neurons, seed 2 shows worse MAE (+1.20%) but better CorrH (0.103 vs 0.087) under Johnson. Unexplained; possibly noise, possibly a real accuracy-versus-temporal-structure tradeoff at scale.
Both arms are identical except for the normalization:
- Fit normalizers on the official training split only — z-score parameters, and
scipy.stats.johnsonsu.fitper neuron. Every fit succeeded at every scale tested (100%), so nothing silently fell back to z-scoring. The harness falls back to z-scoring for any neuron whose fit fails or produces non-finite values. - Transform, train, predict.
- Inverse-transform predictions back to ΔF/F before computing any metric, so neither arm is measured in a coordinate system that flatters it.
- Paired Wilcoxon across test windows, plus the seed-to-seed spread — which is the honest measure of reproducibility. The Wilcoxon statistic only measures within-run consistency.
Splits, offsets and horizons are taken verbatim from zapbench/constants.py rather than re-derived: CONDITION_OFFSETS, CONDITION_PADDING = 1, CONDITIONS_TRAIN = (0,1,2,4,5,6,7,8) with TAXIS (index 3) held out, TEST_FRACTION = 0.2, VAL_FRACTION = 0.1, PREDICTION_WINDOW_LENGTH = 32.
Models are PyTorch reimplementations following the official tsmixer.py and tide.py configs: TSMixer with 2 mixing blocks and mlp_dim 256; a TiDE-style channel-independent residual MLP with two 128-unit hidden layers, no covariates. They are not the official implementations, and the two models see different amounts of training data, so cross-model comparison in this repo is not valid — only the within-model Johnson-vs-z comparison is.
pip install torch tensorstore scipy numpy
# args: N_NEURONS STEPS SEEDS "model,context,instance_norm" OUTPUT.json
python run_sweep.py 2000 1500 0,1,2 "tsmixer,4,False" tsmixer_2k.json
python run_sweep.py 2000 1500 0,1,2 "tide,4,False" tide_2k.json
python run_sweep.py 10000 1500 0,1,2 "tsmixer,4,False" tsmixer_10k.jsonThe data loads anonymously from the public zapbench-release GCS bucket — no credentials needed. Runs on CPU; a 2,000-neuron TSMixer arm takes roughly two minutes, a 10,000-neuron arm roughly twelve. Note that CPU and GPU arithmetic differ enough to move the magnitude (an earlier T4 run of the 2,000-neuron TSMixer configuration gave −0.83%), though not the direction.
ZAPBench's Finding #8 reports that MAE clusters anatomically rather than spreading evenly, with higher error in the dorsal-anterior pallium. That gives a concrete hypothesis for the scaling behaviour here: if the normalization benefit is concentrated in particular regions, averaging over larger neuron subsets dilutes it toward zero — which is what 200 → 2,000 → 10,000 looks like.
The data to test it is already public. SEGMENTATION_DATAFRAMES in zapbench/constants.py points to a 52 MB dataframe carrying centroid_x/y/z, area, extent, bbox_* and ellipticity for all 71,721 neurons; POSITION_EMBEDDING_SPECS adds a 71,721 × 192 position embedding. The plan is to retain per-neuron error instead of averaging it, then test whether the Johnson-vs-z delta varies with anatomical position and specifically whether the pallium behaves differently.
The Johnson system of distributions is due to Johnson (1949); the AS-99 fitting algorithm to Hill, Hill & Holder (1976). This work uses scipy.stats.johnsonsu, not a novel implementation of that math. The contribution here is the paired experimental comparison and the scaling result, not the statistics.
ZAPBench and its dataset are the work of Lueckmann, Immer, Chen, Li, Petkova, Iyer, Hesselink, Dev, Ihrke, Park, Petruncio, Weigel, Korff, Engert, Lichtman, Ahrens, Januszewski and Jain — ZAPBench: A Benchmark for Whole-Brain Activity Prediction in Zebrafish, ICLR 2025, arXiv:2503.02618. Please cite them for the benchmark.
Harness code MIT licensed — see LICENSE. Issues and corrections welcome; if you find an error in this comparison I would rather know.
The instability at 10,000 neurons is arguably more valuable than the gain at 2,000. It means a single training run at that scale carries almost no information about whether a normalization choice helps — the run-to-run spread is several times the effect being measured. Worth knowing before anyone reports a one-run comparison on this benchmark, in either direction.
A two-page write-up is included as Echelon_ZAPBench_Johnson-SU_Scaling_Study_SyanKazi_2026-07.pdf.