Models & Losses

The network architectures — including the anisotropic backbone and dual-head distance-field design — and the full loss suite with every component, its intent and its default weight.

1. Architectures

Two network definitions live side by side. The first is a conventional single-head U-Net used as a clean reference; the second is the topology-preserving dual-head design that the rest of the pipeline is built around.

2. Why the backbone is asymmetric

This is one of the most important architectural decisions in the repository, and it is easy to overlook. A standard 3D U-Net downscales evenly in all three dimensions from the first stage. Here the downsampling sequence is deliberately shaped instead:

stage 0   stride (1, 1, 1)   resolution preserved — keeps the sheet itself intact
stage 1   stride (1, 2, 2)   pools the two in-plane axes, keeps Z — the thin axis — at full resolution
stage 2   stride (2, 2, 2)   uniform downsampling, deep context
stage 3   stride (2, 2, 2)   uniform downsampling, deepest context

Why this shape. A sheet lies roughly in the XY plane: it is a wide manifold that is only a couple of voxels thick along Z, with a small gap to the next wrap. Pooling along Z would destroy the target almost immediately, so the early stages never pool it. The in-plane axes, by contrast, span hundreds of voxels and carry long, smooth structure, so pooling there costs little information and buys receptive field cheaply. Once the network is deep enough that its features are contextual rather than structural, uniform downsampling is safe again — which is why only the early stages are shaped.

Two further choices reinforce the same goal: residual blocks in the encoder and decoder, which make the deeper stages trainable and help preserve fine structure, and instance normalisation rather than batch normalisation, which is stable at the batch size of a single patch per step. The backbone is also configured to emit a feature map rather than logits — the prediction heads own the output layer, which is what allows two differently-activated heads to share one trunk.

3. The dual-head design

The mask head alone gives the optimiser a problem it cannot solve well: topology is discrete, so “almost connected” and “not connected at all” produce indistinguishable gradients. The second head exists to fix that.

Both heads share the entire encoder and the decoder up to the final stage, then diverge into a small sequence of convolutions with their own normalisation and activation. Only the mask head requires an activation at inference; the distance head is used during training and as a diagnostic signal.

Why the distance head helps topology A distance field is continuous, so it carries gradient information through regions where the binary target is flat — inside a sheet the mask is constant, and outside it is constant too. A network forced to regress distance must build an internal representation of where the sheet is heading, including across a gap or through a near-touching region. That representation is exactly what is needed to avoid splitting a sheet or bridging two wraps, and it makes the topological loss terms meaningful rather than degenerate.

4. The nnU-Net network

The second pipeline uses the network nnU-Net selected from the dataset fingerprint: a residual-encoder 3D U-Net with isotropic spacing, trained on 128 × 128 × 128 patches cut from volumes whose median extent is roughly 320 × 314 × 314, and the framework’s standard loss and augmentation strategy. Unlike the custom backbone it downscales uniformly in all axes, which is precisely the difference the custom design is testing. The exact plan is recorded in docs/plans.json so a reported result can be tied to the configuration that produced it.

5. The loss suite

Three loss classes are implemented. They are not alternatives — the topology-preserving loss composes the thin-manifold terms with distance regression, and the anti-bifurcation loss contributes a further set of physical constraint terms. Each component has a weight, and the weights are the main tuning surface of an experiment.

5.1 Thin-manifold loss

The compact objective, and the one used as a baseline. It balances three overlapping signals so the model is not dominated by the overwhelming background class:

The asymmetry is the point. Missing surface is worse than adding a little extra, because a hole changes topology while a slightly thick sheet does not.

5.2 Topology-preserving loss

The full objective, which subsumes the above and adds terms that target manifold structure directly:

Weights are defaults, not dogma These are the class defaults. The reference experiment overrides them explicitly through the loss-weight schedule, which is why the curriculum — not the class constructor — is the place to look when asking what a particular run actually optimised.

5.3 Anti-bifurcation loss

A focused loss targeting the specific errors this task punishes most. It encodes four direct physical expectations about what a sheet may and may not do, and penalises violations of each:

It is parameterised by a maximum number of segments per scan line and a minimum expected sheet separation, so its assumptions about sheet thickness and inter-sheet spacing are explicit and tunable rather than implicit.

It is disabled by default and enabled per experiment, through the parent loss’s bifurcation_weight. Because it encodes fairly specific physical assumptions, it is most useful once the model already produces a recognisable surface and the remaining errors are topological rather than gross segmentation failures — applying it earlier gives it nothing meaningful to measure.

6. Multi-scale evaluation inside the loss

The topological terms evaluate agreement at a set of neighbourhood sizes rather than one, using anisotropic structuring elements by default (an isotropic variant is available instead), and combine morphological closing and opening fractions into the comparison. The reason is that a gap in a sheet can be a single voxel or a dozen, and a topology term with a fixed scale would only ever see one of those. Aggregating across scales means the loss reacts to structural errors regardless of how large they are, which matters because the scored Betti numbers do not care about the size of the defect either.

7. Composing an objective

The practical recipe for changing the objective:

1. Decide whether the failure is geometric (misplaced boundary, wrong thickness) or topological (split, bridge, hole). Geometric failures are addressed with the overlap and surface components; topological ones with the topology, connectivity and bifurcation components.
2. Change weights in the experiment configuration, not in the loss class. This keeps the class as a reusable implementation and the YAML as the experiment record.
3. Phase the change in with the loss-weight schedule if the new term is only meaningful once the model is producing something surface-like.
4. Watch the competition metrics, not the loss value. The composite score is the target; the loss is only a proxy for it, and the two can diverge when a topological term is doing its job.