Feature request: Frontier-conformal FFN — uncertainty maps, adaptive flood-fill, and proofreading queues

[theworker02]

Problem / motivation

At connectomics scale, the bottleneck is rarely “can FFN grow an object?” — it is where humans should spend the next hour of proofreading, and when the network should stop wasting FLOPs on regions it is already sure about.

Today’s OSS FFN already maintains a rich internal state during flood-fill (object probability / seed logits, move thresholds, FOV movement policy; see ffn/inference/inference.py), and the manual already leans on forward–reverse oversegmentation consensus and iterative bootstrapping. What is missing is a first-class way to turn that state into:

1. a calibrated uncertainty signal over the growing frontier,
2. an adaptive compute policy (spend more FOVs where topology is ambiguous; early-stop where it is not), and
3. a proofreading queue that maximises expected topology repair per human minute — optionally with soft human seeds that steer the same flood-fill loop.

That package sits squarely in Google’s current FFN trajectory (JAX path, TPU batching, LICONN / ExaSPIM-scale volumes, TensorStore I/O) and is orthogonal to assembly/export (#137) and to simply open-sourcing skeleton ERL evaluation (#41).

Proposed feature: “Frontier conformal FFN”

Add an optional inference mode — working title FrontierConformal — that treats flood-fill as an uncertainty-aware search rather than a fixed-threshold walk.

1. Frontier entropy / disagreement maps (cheap, already almost there)

During segment_at / FOV moves, record per-voxel (or per-frontier-band) statistics from quantities FFN already computes:

• predictive entropy / margin of the object probability field near the active frontier
• disagreement between a short forward seed continuation and a PolicyInverseOrigins / reverse probe (reuse the consensus idea, but locally and online, not as a full second volume pass)
• optional dropout / EMA-weight disagreement if EMA training (Add the option to track EMA weights in FFN training. #117-era direction) is available in the JAX trainer

Write these as a compact auxiliary volume (uint8 quantised is fine), TensorStore / Neuroglancer-Precomputed friendly.

2. Conformal calibration → proofreading queues

On a held-out validation cut with skeletons or proofread merges/splits, calibrate a conformal threshold so that “flag for human review” has a user-chosen error budget (e.g. ≥95% of true merge-error loci fall inside the review set).

Emit a sorted proofreading queue:

priority	locus (xyz)	suspected failure mode	local entropy	suggested action
1	…	false merge neck	0.91	split / reseed
2	…	aborted branch	0.87	extend / new seed

This is the headline deliverable for labs: FFN tells you where to look next.

3. Adaptive compute policy (the “wow” demo)

Use the same score to drive inference:

• High confidence frontier → larger effective step / fewer FOV revisits / early segment finalisation
• High uncertainty frontier → denser FOV sampling, local reverse probe, or automatic resegmentation decision-point insertion (hook into existing agglomeration / find_decision_points machinery)

A 30-second Neuroglancer / Colab demo of FFN visibly slowing down and lighting up at ambiguous necks would get a lot of attention — and it is scientifically honest, not just theatre.

4. Soft human-seed steering (optional, high leverage)

Allow proofreaders to paint a soft seed / ban region that is injected into the working object map (logit bias), then resume flood-fill. Combined with (1)–(3), this becomes a tight human↔model loop for the bootstrapping workflow the manual already recommends, instead of “segment offline → proofread offline → retrain later.”

Why this would matter (and why it fits this repo)

• Closes the loop on consensus + bootstrapping already documented in doc/manual.md, without waiting for a full second global segmentation pass.
• Compute-aware: directly relevant to TPU/GPU utilisation work already landing in-tree (batch size, multi-subvolume runners).
• Human-time-aware: the scarce resource in every Google-scale connectomics effort.
• Composable: uncertainty layers feed Neuroglancer; queues feed proofreading UIs; decision points feed resegmentation; none of this replaces assembly (Feature request: TensorStore-backed global ID reconciliation + Neuroglancer Precomputed export #137) or skeleton metrics (Evaluation the segmentation #41) — it uses them as calibration targets.
• Publishable artifact: even a solid OSS reference implementation + LICONN notebook would be citable methodology, not just a CLI flag.

Rough acceptance criteria

• Flag --uncertainty_mode=frontier_entropy|frontier_consensus_probe on inference produces an auxiliary map aligned with the segmentation bbox.
• A calibration utility fits a conformal threshold on a validation cut and reports empirical coverage.
• A queue writer emits a stable, documented JSONL/CSV of review loci sorted by priority.
• Adaptive policy A/B on a public cut (e.g. LICONN demo volume) shows ≥X% FOV savings at iso-merge-error, or ≤Y% merge-error at iso-FOV (pick one primary metric; report both).
• Notebook: live uncertainty overlay + top-K proofreading suggestions on the existing JAX LICONN demo path.
• Soft-seed resume API (library-level is enough) with a unit test that a human ban region prevents growth across a synthetic neck.
• Docs: short section in doc/manual.md distinguishing this from full-volume consensus and from post-hoc skeleton eval.

Non-goals (for v1)

• Full distributed ID reconciliation / Precomputed export (tracked in Feature request: TensorStore-backed global ID reconciliation + Neuroglancer Precomputed export #137).
• Open-sourcing the complete historical skeleton-ERL stack (Evaluation the segmentation #41) — only need enough metrics to calibrate the conformal gate.
• Replacing agglomeration / SECGAN; those remain upstream/downstream modules this can call.

Prior art / related

• Manual: bootstrapping, forward–reverse consensus, checkpoint selection via topology metrics.
• Evaluation the segmentation #41: evaluation availability (complementary calibration target).
• Feature request: TensorStore-backed global ID reconciliation + Neuroglancer Precomputed export #137: assembly + Neuroglancer export (complementary I/O surface for uncertainty layers).
• What about consensus and agglomeration steps? #20: consensus/agglomeration visibility — this proposal makes consensus online and actionable for humans and for compute.

Happy to sketch a minimal JAX-side prototype API against the LICONN notebook if that would help maintainers evaluate scope.