LICONN ExPID82_1 affinity models

Voxel affinity prediction for the LICONN ExPID82_1 volume — expansion light microscopy of mouse brain tissue, not EM — for use with PyTorch Connectomics (tutorials/neuron_liconn_ist).

Two checkpoints, same architecture and same specimen, differing in the image resolution they consume:

file	image grid (ZYX)	trained on	whole-val VOI ↓
affinity_expid82_18nm_128x128x128.ckpt	[24, 18, 18] nm	1 GPU, 200 k	0.9129
affinity_expid82_9nm_128x128x128.ckpt	[12, 9, 9] nm	4 GPUs, 200 k	0.9701 — worse, see below

If you just want a working model, use the 18 nm one. It is the checkpoint with a published, reproduced end-to-end number. The 9 nm model is a resolution experiment whose motivation and limits are documented below; it is not established as an improvement.

Both share: MedNeXt-L kernel 3 (61,779,399 parameters), 1 channel image → 6 channel banis affinity (channels 0–2 are the r1 nearest-neighbour affinity used downstream), [128, 128, 128] ZYX training patch, 200 k steps from scratch, AdamW lr 1e-3, cosine to 0, scale_sigmoid output activation.

---

9 nm model (affinity_expid82_9nm_128x128x128.ckpt)

The ExPID82_1 source image has a finer 9 × 9 × 12 nm XYZ level than the 18 nm grid the other model uses. This checkpoint consumes it directly.

Files	affinity_expid82_9nm_128x128x128.ckpt (247 MB) · train_config_9nm.yaml
Grid	[12, 9, 9] nm ZYX = 9×9×12 nm XYZ — the native source resolution, not an upsample
Training data	final_proofread_mip0 train split, 540 × 8580 × 6690; validation 290 × 8580 × 6690 = 16.65 G voxels
Label	the same 18 nm FFN-proofread GT, nearest-upsampled 2× on every axis
Schedule	200 k steps from scratch, batch 2 per GPU on 4 GPUs
Label erosion	4, not 2 — see below

Why erosion 4 and not 2

Kisuk erosion is XY-only (a (1, t, t) window on a ZYX array), and nearest-2×-upsample followed by erosion(2r) is bit-identical to native erosion(r). So erosion 4 here reproduces the 18 nm model's erosion 2 separation. Using erosion 2 would have silently given the equivalent of native erosion 1 — weaker separation than the baseline — and made the two models differ in two ways at once instead of one.

The label ceiling — read this before interpreting the model

The proofread segmentation exists only at 18 nm. This model's GT is that segmentation nearest-upsampled 2×, and is bit-exact to it, so it carries no additional boundary information:

• half the r1 affinity edges along each axis are intra-parent and trivially "same" (measured: ~62% of mip0 r1 edges are "same" against ~20% at 18 nm, the excess being exactly the trivial intra-parent ones);
• boundary position stays quantized to the 18 nm grid.

So this model tests whether a sharper input helps despite a coarse label. It cannot test finer boundary supervision, and no data on this volume can.

How much extra signal is actually in the 9 nm image

Measured on matched crops, not assumed. The 18 nm level is a true mean-downsample parent of the 9 nm level (corr 0.984, residual consistent with JPEG noise). Trilinear-upsampling 18 nm back to the 9 nm grid explains r² = 0.766, so 16.4% of in-plane power sits above the 18 nm Nyquist — but that tail is nearly flat, and flat is the signature of noise.

The discriminating test is z-coherence, since real anatomy must persist across a 12 nm z-step while per-slice noise and a spatially-fixed JPEG block grid cannot. Above-Nyquist correlation is +0.308 at lag 1 and falls to the +0.099 artifact floor by lag 2, while the below-Nyquist band still holds +0.54 at a 96 nm separation. Scaling the lag-1 excess against the low band puts genuine structure at roughly a quarter of that tail: order 4% of total image power, at 25–45 nm periods.

That band is where thin processes are dropped, so the idea is motivated — but it is a ~4% lever, not a 16% one, and it is spent against a label that cannot localise better than 18 nm.

Evaluation status

Boundary discrimination (measured). Decoder-free, on the 18 nm grid, against the exact (not upsampled) GT: for each pair of voxels adjacent at 18 nm, how well does the affinity predict that they share a proofread object? An 18 nm edge is the min of the two 9 nm edges it spans. AUC over 24 z-planes, 4 M same-parent edges per axis:

axis	18 nm model	9 nm model	delta
z	0.8991	0.8947	−0.0045
y	0.9000	0.8945	−0.0055
x	0.9013	0.8960	−0.0053

The 9 nm model is not better at separating 18 nm boundaries, and is slightly worse, consistently on all three axes. Caveat: these are two different training runs, and run-to-run spread on this project is of the same order (two 18 nm runs differing only in GPU count differ by 0.006 whole-val VOI), so this gap is not by itself attributable to the input resolution.

Whole-volume VOI (measured). Scored by downsampling the 9 nm segmentation onto the 18 nm grid and comparing against the same GT array behind the 0.9129 above — VOI is not comparable across voxel grids, so a 16.65 G-voxel segmentation scored against an upsampled GT would be a different quantity. Decoder and settings are identical to the 18 nm result (ABISS max in compressed space, ws_high/ws_low = 94th/20th percentile), so only the model differs:

mt	VOI ↓	split	merge	AR err	pred segs
0.38	3.4012	0.5740	2.8272	0.8201	119,125
0.41	1.4984	0.6179	0.8805	0.5129	173,361
0.44	1.0025	0.6600	0.3426	0.3951	244,359
0.47	0.9701	0.7143	0.2558	0.3821	362,265
0.50	1.0389	0.7967	0.2422	0.3818	565,795
0.53	1.2100	0.9130	0.2970	0.3945	873,724

This model is worse than the 18 nm one: 0.9701 against 0.9129, i.e. +0.057, five to six times the ~0.01 VOI noise floor. It is worse on both terms — split +0.041, merge +0.016 — so it is not a split/merge trade-off, and it produces 4.6× more segments. The optimum is interior (0.44 and 0.50 are both worse), so the sweep range was not the problem.

Two controls, because a number like this is easy to get wrong:

• The downsample is not causing it. Scoring the same decode at the opposite corner of every 2×2×2 block gives 0.9736 against 0.9701 — a spread of 0.0035, a third of the noise floor and 6% of the gap.
• Run-to-run variance is not causing it. Two 18 nm runs differing only in GPU count differ by 0.006 whole-val VOI, an order below this gap.

What this does not establish is that the input resolution is the cause. This run differs from the 18 nm baseline in three ways at once: half the physical field of view per axis at the same 128³ patch, half the physical long-range affinity radius (r10 = 120 nm Z / 90 nm XY here against 240/180 there), and the label ceiling above. Separating them needs runs that vary one factor at a time.

Do not use this model expecting an improvement. The practical reading is narrower than the scientific one: for this pipeline, on this volume, against a label quantized to 18 nm, the 9 nm input does not pay for 8× the voxels. That is not a claim that the ~4% of genuine above-Nyquist signal is unusable in principle.

Decoding it needs a 64-bit watershed

At 9 nm the validation volume is 290 × 8580 × 6690 = 16.65 G voxels. ABISS ws uses a uint32 internal watershed index, so a single invocation is capped at 2,147,483,648 voxels (watershed_traits<uint32_t>::high_bit) — this volume is 7.8× over and the binary aborts on an assertion. For reference the 18 nm validation decode sits at 98.3% of that same cap. Decoding at 9 nm therefore requires either ABISS's chunked pipeline or a build with a 64-bit internal index. Note the assert is compiled out under -DNDEBUG, in which case the index silently overflows instead of failing loudly — check your voxel count rather than relying on the assertion.

The EMA weights are not in this file

EMA was enabled (decay 0.999, validate_with_ema: true), so the logged val_loss_total describes EMA weights. Unlike the 18 nm run, this run's training checkpoint does retain them (callbacks → EMAWeightsCallback → ema_state, 517 tensors). They are not included here: this file holds the raw training weights, which are the weights every number above was measured on, because PyTC's test mode does not instantiate the EMA callback and therefore never swaps them in. The EMA weights are unevaluated; ask if you want them published.

---

18 nm model (affinity_expid82_18nm_128x128x128.ckpt)

Files	affinity_expid82_18nm_128x128x128.ckpt (247 MB) · train_config.yaml (the frozen run config)
Grid	[24, 18, 18] nm ZYX = 18×18×24 nm XYZ
Training data	ExPID82_1 train split, 270 × 4290 × 3345, FFN-proofread GT
Schedule	200 k steps from scratch, AdamW lr 1e-3, batch 2 on 1 GPU, cosine to 0
Result	0.9129 VOI on the full held-out validation volume (ABISS decode)

18 nm is the finest scale at which the proofread segmentation exists. The source image has a finer 9×9×12 nm level (used by the other model), but there is no GT there, so this model is trained at the finest paired image+label resolution available.

⚠️ This checkpoint holds raw training weights, not EMA weights

EMA was enabled during training (decay 0.999, validate_with_ema: true), which means the logged val_loss_total values were computed on EMA weights, while the checkpoint's state_dict holds the raw training weights — the framework undoes the EMA swap before writing the checkpoint. EMA-state persistence was added to PyTC three weeks after this run, so this model's EMA weights were never saved and are unrecoverable.

Practical consequences:

• Do not quote the training-time validation loss as this checkpoint's loss; they describe different weights.
• The 0.9129 above is unaffected and reproducible: in test mode the EMA state is never initialised, so that inference ran on exactly the weights in this file.

Results

Full held-out validation volume, 145 × 4290 × 3345 = 2.08 G voxels, ABISS max-affinity agglomeration at ws_high/ws_low = 94th/20th percentile, ws_merge_threshold 0.47:

	VOI ↓	split	merge	Adapted-Rand err ↓	pred segs	GT segs
ABISS max, mt 0.47	0.9129	0.6732	0.2397	0.3195	79,056	35,815

The decoder is split-dominated at this setting, as configured. 0.47 has since been confirmed as the whole-volume optimum by a bracketed sweep (0.44 → 0.9341, 0.47 → 0.9130, 0.50 → 0.9680), superseding the earlier note that it was slab-tuned and probably too high.

---

Shared usage notes

These apply to both checkpoints.

Inference window must be [128, 128, 128]

MedNeXt normalizes with GroupNorm — per-sample, per-channel, no running statistics — so normalization is computed over the sliding window's spatial extent at every block and the forward pass is window-size dependent. Both published results were produced with sliding-window ROI [128, 128, 128], matching the training patch. Change the window and the numbers no longer apply.

Usage

# 18 nm
hf download pytc/liconn affinity_expid82_18nm_128x128x128.ckpt --local-dir ckpt/
python scripts/main.py --config tutorials/neuron_liconn_ist/1_affinity.yaml \
  --mode test --checkpoint ckpt/affinity_expid82_18nm_128x128x128.ckpt

# 9 nm — same recipe, but point the data at the 9 nm export and set
# params.data.resolution_zyx: [12, 9, 9]
hf download pytc/liconn affinity_expid82_9nm_128x128x128.ckpt --local-dir ckpt/

Output is float16 CZYX. Arrays are ZYX, so channel c is the edge along array axis c: ch0 = Z, ch1 = Y, ch2 = X.

Affinities are scale_sigmoid, not probabilities

The pipeline stores sigmoid(0.2 · logit), not a calibrated edge probability. Measured on the 18 nm checkpoint the stored affinity spans about [0.01, 0.80] and never reaches 0.88; on the 9 nm checkpoint, [0.04, 0.81]. ABISS thresholds copied from other PyTC tutorials (e.g. Pinky's ws_high_threshold: 0.88) therefore sit above the data maximum and seed nothing.

tutorials/neuron_liconn_ist/2_abiss.yaml gives ws_high/ws_low as percentiles and keeps ws_merge_function: max, which is monotone-invariant so a sweep in the compressed space covers the same family of segmentations. mean is not monotone-invariant and must not be substituted without uncompressing first — note that ABISS's chunked pipeline agglomerates with mean only.

For ABISS also set channels: [2, 1, 0] (its ws reads XYZC with channel 0 = X edge) and edge_storage: source (this repo writes edge (i, i+1) at voxel i; ws reads the value at i as edge (i−1, i) — without this the boundary map is one voxel off on every axis).

Training data

Image and GT both come from the public ExPID82_1 release: image image_230130b, segmentation 231030_agg_240123 (FFN, proofread). The 18 nm model uses crop [140,240,240] → [555,4530,3585] ZYX split at z = 270; the 9 nm model uses [280,480,480] → [1110,9060,7170] split at z = 540 — exactly 2× on every axis, i.e. the identical physical region.

Limitations

• One acquisition. All numbers are within-volume development evidence on ExPID82_1, not an independent test score, and thresholds are not known to transfer to other specimens.
• Trained on proofread labels from the target volume — supervision on the target domain.
• The GT is 42.5% background (proofread FFN leaves much unlabelled); VOI here ignores the background label.
• The 9 nm model additionally inherits the label ceiling described above, and scores worse than the 18 nm model (0.9701 vs 0.9129); it is published as a documented negative result, not as an improvement.