Johnson-SU normalization vs. z-scoring on ZAPBench traces: consistent gains at small neuron subsets that become unreliable by 10k

[EchelonDevelopment-Syan]

Hi — first, thanks for releasing ZAPBench with the data, the split definitions, and the baseline configs. The fact that zapbench/constants.py pins the condition offsets, the held-out condition and the split fractions is what made this experiment cheap enough for one person to run, and it's why the numbers below can be checked rather than taken on faith.

I ran a small paired experiment on the trace-forecasting task and wanted to report it here, including the part that didn't work out the way I expected.

The question

The standard preprocessing for the trace models is per-neuron z-scoring. ΔF/F traces are visibly non-normal — skewed, heavy-tailed on the positive side — so an obvious thing to try is a transform that fits that shape. I used the Johnson-SU distribution (Johnson 1949; the AS-99 algorithm of Hill, Hill & Holder 1976; here via scipy.stats.johnsonsu.fit), fitting four parameters per neuron on the training split only, forecasting in the transformed space, then inverse-transforming before computing any metric.

Everything is evaluated in ΔF/F space for both arms, so the comparison isn't flattered by measuring error in a stretched coordinate system.

Setup

Identical for both arms, and taken from your published constants rather than re-derived:

• CONDITION_OFFSETS, CONDITION_PADDING = 1, CONDITIONS_TRAIN = (0,1,2,4,5,6,7,8), TAXIS (index 3) held out
• TEST_FRACTION = 0.2, VAL_FRACTION = 0.1, PREDICTION_WINDOW_LENGTH = 32
• Short context, C = 4
• Data: zapbench-release, volumes/20240930/traces/, the (7879 × 71721) matrix
• PyTorch reimplementations of TSMixer (2 blocks, mlp_dim 256) and a TiDE-style channel-independent residual MLP (2 hidden layers of 128), following tsmixer.py / tide.py
• 1,500 Adam steps at lr 2e-3, batch 32, per arm; both arms get the same budget, the same seed, and no tuning
• Normalizers fit on the training split only; every Johnson-SU fit succeeded (100%) at every scale, so nothing silently fell back to z-scoring
• Metric: mean absolute error over all test windows, plus per-window paired Wilcoxon
• Neuron subsets drawn with a fixed RNG; three seeds per configuration (five at 10,000 neurons), matching the 3-seeds-per-model protocol in your Figure 5

Results

Negative delta means Johnson-SU beat z-scoring.

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

At 2,000 neurons the effect is unambiguous: six out of six runs favour Johnson-SU, TiDE extremely tightly (sd 0.03), and per-window Wilcoxon p-values run from 1.7e-18 to 6.2e-53 with roughly 65–78% of test windows improved.

At 10,000 neurons the variance takes over. Five seeds give −0.54%, −0.36%, +1.20%, −2.27% and −1.56%: a mean of −0.71% with a standard deviation of 1.32, against effect sizes well under 1% at the smaller subsets. Four of five still favour Johnson, so on average the benefit may well persist — but the run-to-run spread is now roughly double the effect being measured, and one seed goes the other way by more than the mean itself. Whatever is stable at 2,000 neurons is not stable at 10,000. I would not report a single run at that scale as evidence of anything in either direction, and that is the part I think is worth your attention more than the gain at 2,000.

One detail from seed 2 that I don't have an explanation for: its MAE got worse under Johnson (+1.20%) while its CorrH got better (0.103 vs 0.087). If that's reproducible rather than noise, it would suggest the transform is trading absolute accuracy for temporal structure at scale, which seems worth someone's attention.

What I am and am not claiming

Claiming: on short-context trace forecasting, per-neuron Johnson-SU normalization gives a small but reliable MAE improvement over z-scoring at 200 and 2,000 neuron subsets, across three seeds and two architectures, in a paired comparison with equal training budgets and no tuning for either arm.

Not claiming: anything at full 71,721-neuron scale. Anything about long context — I did not run C=256, and my C=4 and C=256 configurations would have differed in mlp_dim and instance-norm anyway, which would have confounded the comparison. Any absolute MAE that belongs on your leaderboard: 1,500 steps on a neuron subset with no early stopping is far below your training budget, so these numbers are only meaningful as a paired delta.

I'll also say plainly that I got this wrong more than once on the way here. The first two 10k seeds both favoured Johnson and I briefly concluded the effect held at scale; the third flipped hard the other way; the fourth swung further still; the fifth came back. Every intermediate conclusion was drawn from too few seeds — the same lesson your Figure 5 three-seed protocol already encodes, learned the slow way.

Caveats worth stating up front

Everything ran on CPU, which means the arithmetic differs slightly from a GPU run; when I ran the same configuration on a T4 earlier I got −0.83% for TSMixer at 2,000 neurons rather than the −0.90 to −1.73% I see here. The direction was stable across hardware, the magnitude was not, which is another argument for reporting seed distributions rather than point estimates. The Wilcoxon p-values measure within-run consistency across test windows, not reproducibility across training runs — the seed spread is the honest measure of that, and it's why the table shows every seed.

Where I think this goes next

Finding #8 reports that MAE clusters anatomically rather than spreading evenly, notably higher in the dorsal-anterior pallium. That suggests a specific hypothesis for the scaling behaviour above: if the normalization benefit is concentrated in particular regions, then averaging over ever-larger neuron subsets would dilute it toward zero — which is roughly what the 200 → 2,000 → 10,000 progression looks like.

That is testable with what you have already published. SEGMENTATION_DATAFRAMES gives per-neuron centroid_x/y/z, area, extent, bbox and ellipticity for all 71,721 neurons, and POSITION_EMBEDDING_SPECS provides a 71,721 × 192 position embedding. So rather than reporting mean MAE, the next run should retain per-neuron error and regress the Johnson-vs-z delta against anatomical position, then check specifically whether the pallium behaves differently from the rest of the brain. I plan to do that and will report back either way.

One genuine question in the meantime: has anyone on the team already tried alternatives to per-neuron z-scoring for the trace models? Given how non-normal ΔF/F is, I'd assume it came up, and I would much rather hear it was tested and rejected — and why — than spend a week re-deriving a known dead end.

Code to reproduce, including the seed sweep and the raw per-run outputs, is at: https://github.com/EchelonDevelopment-Syan/zapbench-johnson-normalization

Happy to run anything specific if it would be useful, and equally happy to be told this is a known dead end.

— Syan Kazi, Echelon Development & Contract Consulting

[farfromexact]

This is a strong report, especially the distinction between a Wilcoxon result over test windows and reproducibility across training seeds. The 10,000-neuron table is exactly where a single headline number would be misleading.

Before changing the benchmark, I would preserve the raw per-window and per-neuron deltas, record the exact code commit and fitted normalizer parameters, and run a larger seed sweep with the same CPU or GPU class. It would also be useful to compare robust scaling, Yeo-Johnson, and a simple log-like transform under the same budget, then report paired confidence intervals and an effect distribution rather than only p-values.

The anatomical hypothesis is testable, but I would define the region analysis and any multiple-comparison correction before looking at the result. A normalizer interface in the benchmark code could make these comparisons reproducible without rewriting each model. The MAE versus CorrH discrepancy is worth keeping as a separate metric tradeoff question rather than treating one metric as the explanation for the other.

[EchelonDevelopment-Syan]

Thanks, Cong — this is exactly the kind of feedback I was hoping for, and I appreciate the time you put into it. You're right that the 10k table is where a single headline number would mislead, which is why I framed it as a reversal rather than a result. I'd rather this live at "small-subset gain that doesn't survive scale" than overclaim it. Here's what I'll do, in the order you laid out: - Preserve the raw per-window and per-neuron deltas, and pin the exact commit plus the fitted normalizer parameters (γ, δ, ξ, λ per neuron) so the run is fully reproducible. - Run a larger seed sweep on the same hardware class, and report paired confidence intervals and the effect-size distribution rather than leaning on p-values. - Add robust scaling, Yeo-Johnson, and a simple log-like transform as comparators under the same compute budget, so Johnson-SU isn't evaluated in isolation. - For the anatomical hypothesis, I'll pre-specify the region grouping and the multiple-comparison correction before looking at the outcome, so it stays a real test rather than a post-hoc story. - Keep the MAE-vs-CorrH difference as its own metric-tradeoff question rather than using one to explain the other. The normalizer-interface idea is appealing — a thin wrapper that applies a chosen transform at the trace level would let all of these swap in without touching the models. If that's something you'd welcome as a small PR against the benchmark, I'm happy to prototype it; otherwise I'll keep it in the discussion so it stays reproducible. One practical note on scope: the per-neuron Johnson fit is the compute bottleneck at full brain scale, so I'll run the seed sweep at a fixed neuron subset first and be explicit about the scale each number comes from. Thanks again — this gave me a clear path forward. Best, Syan Kazi Founder & CEO | Echelon Development & Contract Consulting, LLC C:(618) 924-2157 Tampa, FL ***@***.*** | echelonbridge.com | syan.earth | neural-recovery.com | AmbientFenceLighting.com | “Every Handshake Builds A Bridge Between You, Ai, & New Reality” CONFIDENTIALITY NOTICE The information contained in this transmission is intended only for the person or entity to which it is addressed and may contain confidential and/or privileged material. If you are not the intended recipient of this information, do not review, retransmit, disclose, disseminate, use, or take any action in reliance upon this information. If you received this transmission in error, please contact the sender and destroy all printed copies and delete the material from all computers

…

On Sun, Jul 26, 2026 at 1:25 AM Cong Ma ***@***.***> wrote: This is a strong report, especially the distinction between a Wilcoxon result over test windows and reproducibility across training seeds. The 10,000-neuron table is exactly where a single headline number would be misleading. Before changing the benchmark, I would preserve the raw per-window and per-neuron deltas, record the exact code commit and fitted normalizer parameters, and run a larger seed sweep with the same CPU or GPU class. It would also be useful to compare robust scaling, Yeo-Johnson, and a simple log-like transform under the same budget, then report paired confidence intervals and an effect distribution rather than only p-values. The anatomical hypothesis is testable, but I would define the region analysis and any multiple-comparison correction before looking at the result. A normalizer interface in the benchmark code could make these comparisons reproducible without rewriting each model. The MAE versus CorrH discrepancy is worth keeping as a separate metric tradeoff question rather than treating one metric as the explanation for the other. — Reply to this email directly, view it on GitHub <#36?email_source=notifications&email_token=BWDOZDJZAI2BBNYYP6DYS6L5GWI4BA5CNFSNUABIM5UWIORPF5TWS5BNNB2WEL2ENFZWG5LTONUW63SDN5WW2ZLOOQXTCNZXHA2DIOBRUZZGKYLTN5XKMYLVORUG64VFMV3GK3TUVRTG633UMVZF6Y3MNFRWW#discussioncomment-17784481>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/BWDOZDJQDDXVG7AYNEKITAT5GWI4BAVCNFSNUABIKJSXA33TNF2G64TZHM4TENBSHA4DONJXHNCGS43DOVZXG2LPNY5TCMBUHE2DSMJSUF3AE> . You are receiving this because you authored the thread.Message ID: ***@***.*** com>

[EchelonDevelopment-Syan]

farfromexact — thank you for the checklist. It took me longer to come back than I meant it to: working through your fourth point properly meant adding a second factor rather than repeating the first experiment, and I pre-registered the whole design before running a cell.

Short version: the result changed my mind about the mechanism.

What I ran. 5 neuron scales (200 / 1k / 2k / 5k / 10k, nested by construction and asserted at runtime) × 2 models × 5 transforms (z, robust, Johnson-SU, Yeo-Johnson, signed-log) × 8 seeds × 2 loss-weighting conditions — native, and equalized to MAD = 1. That's 800 training runs, forming 640 paired comparisons across 80 cells. One GPU class (L4, name recorded in every receipt), one code fingerprint, ZAPBench pinned at 2b0a59eb. The primary family, the multiplicity correction and what would count as a null were fixed in that same document.

What came out. In my earlier post I said Johnson-SU wasn't uniquely responsible, because robust scaling matched it nearly everywhere — and since robust is linear and reshapes nothing, I concluded the mechanism couldn't be reshaping. That was wrong, and the second factor is what showed it. Equalize the loss weighting and robust's advantage goes to zero: it clears the seed-level interval at 9 of 10 model × scale cells under native weighting and 0 of 10 under equalized. Johnson-SU clears 9 of 10 equalized — all five TiDE scales, four of five on TSMixer — and its paired seed-to-seed shift is not distinguishable from zero at 9 of 10 (TiDE at 5,000, for instance: shift +0.39, 95% CI [+0.29, +0.49] for robust, against nothing resolvable for Johnson-SU). Dividing by IQR instead of SD changes each neuron's effective weight in the loss, and on this data that was the whole of robust's contribution.

On MAE versus CorrH — this is the part I'd most like your view on. You asked me to keep the discrepancy as its own tradeoff question rather than letting one metric explain the other, and holding that line is what made the next thing visible.

Under native weighting, signed-log posts the best MAE on TSMixer while its CorrH sits at 0.049–0.076 against roughly 0.09–0.12 for every other arm, and it is outright worse on TiDE, around +4.5%. I treated it as a cautionary negative on that basis. Under equalization both halves of that picture disappear. The TiDE MAE penalty of +4.46 / +4.37 / +4.42 / +4.89 / +4.47 across the five scales becomes +0.04 / −0.01 / −0.01 / +0.01 / +0.00. And the CorrH collapse goes with it: the deficit against z runs −0.034 to −0.055 across all ten native cells, and lands between −0.0002 and +0.0051 across all ten equalized ones — inside the band the other transforms occupy.

So the metric tradeoff I reported looks like an artefact of the loss weighting rather than a property of the transform. I've kept signed-log excluded from every positive claim because my pre-registration commits to that exclusion on the original CorrH grounds regardless of what this run showed, and I'd rather honour that than rewrite the rule after seeing the answer. But the ground it rests on is weaker than when I wrote it, and I think that's worth saying plainly rather than leaving in a footnote.

The other half of that discrepancy was the loose end I left in my original post — the 10,000-neuron seed whose MAE got worse under Johnson while its CorrH improved, which I said I had no explanation for. Eight seeds split it in two. The specific observation doesn't reproduce: at TSMixer 10,000 with Johnson-SU, two seeds of eight regress on MAE, and only one of those also improves on CorrH — the other loses on both (MAE +1.18%, CorrH −0.003). The pattern underneath it does. Pooling all 240 native seed-runs across the non-signed-log arms, per-seed MAE delta and CorrH delta are negatively rank-correlated — Spearman ρ = −0.46, p ≈ 9e−14 — so runs that give up MAE do tend to gain temporal structure. Split by cell it is concentrated exactly where I guessed: TSMixer at 5,000 (ρ = −0.86, p = 0.007) and at 10,000 (robust ρ = −0.88, p = 0.004; Johnson-SU ρ = −0.71, p = 0.047), and absent on TiDE at every scale and on TSMixer at 2,000. This one is exploratory — I did not pre-register it and no claim rests on it, and I'm offering it as the tradeoff question in its own right rather than as either metric explaining the other. But it has the shape I described a week ago, and it looks like a large-TSMixer phenomenon rather than a property of the transform.

On your opening point. The window-versus-seed distinction turned out to be bigger than the one cell I showed you. Of the 31 cells whose seed-level interval spans zero, the window-level bootstrap excludes zero in 17. TSMixer at 10k with robust scaling is the sharpest: window CI [−1.86, −1.17], seed CI [−3.69, +0.67], same eight runs. Reporting the window interval alone would have turned more than half my inconclusive cells into findings. I've printed both side by side everywhere, with the seed-level one governing.

Your list, point by point:

• Raw per-window and per-neuron deltas preserved — both axes in every per-seed .npz, all 160 of them.
• Exact commit and fitted normalizer parameters — full 64-character code digest (the previous run stored a truncated one), the ZAPBench commit, and γ, δ, ξ, λ per neuron in fitted_params.npz for every cell.
• Larger seed sweep on the same hardware class — 8 seeds in every cell, uniform, fixed in advance; a single L4 class with gpu_name recorded and asserted identical across all ten receipt folders.
• Robust scaling, Yeo-Johnson and a log-like transform under the same budget — all four against z at identical budget, with the weighting factor crossed over the whole matrix.
• Paired confidence intervals and an effect distribution rather than p-values — seed-level t-interval and window-level bootstrap reported together for all 80 cells, with every per-seed value printed.
• Region analysis and multiple-comparison correction defined before looking — the correction is Holm–Bonferroni over a pre-specified primary family. One deviation, recorded as one: at n = 8 the exact two-sided Wilcoxon has a minimum attainable p of 0.0078, so Holm over a family of 10 cannot clear 0.05 at any effect size. I applied it to the t-test p instead and reported both. The region grouping is still open — more on that below.
• MAE and CorrH kept separate — no CorrH movement is used to explain an MAE result or the reverse, which is what surfaced the signed-log point above.

What I'm not claiming. At TSMixer 10,000, robust, Johnson-SU and Yeo-Johnson are all unresolved in both weighting arms — robust is −1.51% with a seed interval of [−3.69, +0.67] at eight seeds. From the observed spread it would take roughly 26 seeds to resolve, which I've published alongside the null so nobody has to re-run to find that out. I'd rather leave it inconclusive than lean on the point estimate. (The one cell at that scale that does clear is equalized signed-log, −1.69% [−3.15, −0.23], and it falls under the pre-registered exclusion above.) My evaluation is also a temporal tail-split within the training conditions, with condition 3 / TAXIS excluded entirely rather than used as the held-out condition — a controlled A/B, not your official protocol, and not leaderboard numbers.

What I'd do next. The obvious single-variable experiment is SU versus SB. MIN_MAX_VALUES bounds ΔF/F to roughly (−0.25, 1.5), and on the subset I've tested a moment-based router picks the bounded SB branch for about half of real neurons — yet every arm here forces SU. My anti-dilution guard reports that the SU fit succeeded everywhere, which isn't the same as SU being the right family, so some of this may be a result about a slightly misspecified Johnson-SU. If there's a reason the bound doesn't bite the way I think it does, I'd rather hear it than spend the compute finding out.

Two things I'd defer to you on. Whether this wants re-running under the official held-out-condition protocol before it means anything from your side — I kept the tail-split because it isolates the normalizer cleanly, but that's more your call than mine. And whether there's a region grouping the ZAPBench authors would treat as canonical: you asked me to pre-specify the grouping and the correction before looking, and I'd rather pre-specify yours than invent a parcellation of my own. Per-neuron error is retained for all twenty cells now, including the equalized arm, so the weighting confound is removable either way.

One design lesson for my own next pre-registration, since it bears on how I reported the nulls: ten of the inconclusive cells have their entire interval inside ±0.1%, which is closer to "no effect" than "can't tell" — but I hadn't fixed a TOST margin in advance, so inconclusive is what they get. Next time I'll set one before running rather than reach for it afterwards.

On the normalizer interface — the idea you mentioned — I'm partway there but not as far as I'd like to claim. The five transforms are named in one registry tuple and are fitted and inverted inside a single file, and adding one touches no model code at all. But dispatch is by string rather than through an interface, so what you're describing is a real refactor on my side rather than a rename. If a thin trace-level wrapper would be useful as a small PR against the benchmark, I'd be glad to build it against whatever shape you'd prefer; if you'd rather it stay out of the repo, it lives in the discussion and stays reproducible either way.

Thanks again for taking the time on the first version. Getting a reviewer's checklist on something I'd built on my own is the most useful thing that's happened to this project, and I'd have kept the wrong mechanism story for a lot longer without it. Happy to share the receipts folder, the pre-registration, or any individual cell if it's useful.