[Hendrik-code]

The bug

_region_moments computes per-region voxel count, mean and variance in one matmul, region_mask @ design, where the first column of design is all ones. So the first output column is a plain voxel count. autocast demotes that matmul to fp16, which maxes at 65504, and the count saturates to inf for any region above about 40³ voxels. At the 128³ training patch size the background region always does.

The failure is silent and total:

• counts = inf with sums = inf gives means = NaN
• every per-region Gaussian contribution in RandomRedistributeSegGPU becomes NaN
• the tmax - tmin > 1e-8 guard is False on NaN, so torch.where returns the input
• retain_stats maps that untouched min-max normalisation back to the original mean and std

The transform is then an exact identity. Nothing raises: the torch.isfinite(x).all() guard cannot see it, because the NaN only ever lived in to_add, and record_applications still reports changed=True, because the fp32 round trip leaves 2.4e-7 of rounding noise and the flag is an exact torch.equal.

The matmul form arrived in 20360d8 ("perf: remove the per-sample loops and host stalls from the augmentations"). Before it, the same quantity was mask_flat.sum(dim=1), a reduction, which autocast keeps in fp32. So redistribution has been dead in every run trained since.

Reproduced

On a real preprocessed Dataset804 128³ patch under autocast(float16), five seeds, in_seg 0.25 and 1.0, data of range 1.74:

	region count	max abs change	voxels moved
main	inf	2.384e-07	0.0 %
this PR	1540096.0	0.92 – 1.83	99.9 – 100 %

The patched numbers are bit-identical to the same calls with autocast switched off, which is the result the fix is supposed to produce.

What it cost

Four single-variable retrains of the paper's Ours configuration, this patch and nothing else changed, 1000 epochs each, evaluated per subject against both the pre-refactor runs and their 20360d8-era siblings. Out-of-domain mean Dice, test split:

arm	pre-20360d8	with the dead transform	this PR
Dataset803 (in-phase)	0.838	0.827	0.844
Dataset804 (CT)	0.760	0.639	0.777
Dataset806 (Dixon fat)	0.833	0.813	0.839
Dataset807 (Dixon water)	0.829	0.828	0.840

Worst single cell, Dataset804 evaluated on SpineGAN Dixon water: 0.713 → 0.301 → 0.769 Dice, and its spinal-canal class 0.824 → 0.110 → 0.863. ASSD agrees, 2.47 → 10.67 → 1.64 mm. In-domain cells move by at most 0.006 in either direction, so this is transfer, not fit, and every difference quoted is Wilcoxon-significant at p<0.001 per subject.

Two caveats worth stating. Out-of-domain cohorts that never regressed also improve by +0.004 to +0.051 Dice, largest on CT seen from an MRI-trained model, so the patch is not a pure restoration of the old behaviour; it also lets 58d8783 ("read segmentation regions from label values, not the channel axis") take effect for the first time, since before that the whole foreground collapsed to one region. And each cell is a single fold-0 run, so none of this separates a seed effect within a cell.

The fix

with torch.autocast(device_type=design.device.type, enabled=False):
    sums = region_mask.to(torch.float32) @ design.to(torch.float32)  # (R, 3)

fp32 represents these counts exactly (128³ is far below 2²⁴) and TF32 keeps the matmul as fast as the fp16 one: _region_moments goes 0.329 → 0.403 ms per call on an A40 at the training shape, the whole transform 9.24 → 9.63 ms, about 0.04 s on a 53.7 s epoch at p=0.4. Median epoch time over the four 1000-epoch retrains came out lower than their siblings' on three of four, so the cost is below card-load noise at this scale.

The moments are returned in fp32 deliberately. Casting back to design.dtype would reintroduce the overflow at the caller.

Tests

unit_tests/test_redistribute_autocast.py pins both halves: that the count survives autocast and agrees exactly with the same call outside it, and that the transform is not an identity under autocast. The four tests fail on main and pass here. They need no GPU, because torch.autocast("cpu", torch.float16) demotes the same op.

They do need a 64³ patch. The identity requires two overflows: with only the count at inf the mean comes out 0 rather than NaN, which is still wrong but still moves the data. The NaN needs sum(x) to overflow too, which in the transform's [0, 1] working space takes roughly 131k voxels, so 48³ does not reproduce it and 64³ does. That is in a comment next to the constant.

🤖 Generated with Claude Code