Update model zoo metadata

README.md

@@ -9,7 +9,7 @@ neuron-segmentation experiments on CREMI. The maintained path is:
 - self-supervised dbMiM / MAE-style pretraining on unlabeled EM volumes;
 - anisotropic 3D UNETR affinity finetuning on CREMI;
 - full-volume CREMI A/B/C evaluation with VOI and adapted Rand error (ARAND);
-- waterz-based post-processing with calibration and threshold sweeps.
+- waterz-based reference post-processing with small calibration and threshold sweeps.
 
 The old private cluster launchers, scratch reports, cached bytecode, legacy
 models, and historical dataloaders have been removed from Git. Large data,
@@ -17,8 +17,7 @@ checkpoints, reports, and local experiment outputs are intentionally ignored.
 
 ## Current Method
 
-The current best method is not the original minimal reproduction. It combines
-the following changes that were stable in ablations:
+The official implementation follows the dbMiM neuron-segmentation pipeline:
 
 1. **Anisotropic UNETR backbone**: `UNETRAnisotropicAffinityNet` with
    `32x160x160` input crops, `patch_size=(4,16,16)`, transformer hidden states
@@ -35,16 +34,70 @@ the following changes that were stable in ablations:
    evaluated with `ignore_label=0`, CREMI-style XY boundary ignore distance `1`,
    z boundary ignore `0`, logit calibration biases, and a waterz threshold sweep.
 
-The most useful new finding is that **fixed mixed edge/random masking on full
-EM data (R33)** improves over scratch, old fullEM dbMiM, pure edge masking, and
-fullEM plain MAE. The best absolute VOI is still the smaller publicEM dbMiM
-model (R17), so the README reports both.
-
-The recommended R33 line does **not** use a reinforcement-learning masking
-policy. It uses fixed mixed edge/random masking. RL-style decision modules are
-retained only as ablations: R16 used the older decision module in the publicEM
-pretraining line, while R34/R35 tested adaptive mixed masking and were negative
-under the current CREMI A/B/C protocol.
+## Implementation Details
+
+**Decision-based MIM.** The pretraining stage treats each 3D ViT patch as an
+agent. Each agent chooses a binary action, `keep` or `mask`, and all agents
+share policy parameters. The target network provides global and local features
+for the decision module. The policy receives a team reward from the change in
+masked-reconstruction loss and is trained with an actor-critic/A2C-style
+objective.
+
+**Reconstruction targets.** The masked-image target reconstructs raw EM voxels
+and structure-aware features. The maintained dbMiM pretraining path uses raw
+reconstruction, membrane-aware reconstruction weighting, and a lightweight
+structure consistency term. Plain MAE controls use the same data, crop size,
+model size, mask ratio, and training schedule with the dbMiM-specific structure
+terms disabled.
+
+**UNETR transfer.** Finetuning loads compatible ViT encoder prefixes from the
+pretrained masked-image model into the anisotropic UNETR affinity network:
+`pos_embed`, `patch_embed`, `encoder_blocks`, and `norm`. The convolutional
+decoder and affinity head are trained for the supervised CREMI task.
+
+**Affinity and instance decoding.** CREMI instance labels are converted to
+nearest-neighbor z/y/x affinities. The finetuned network predicts affinity
+logits; calibrated affinities are converted to instance labels with watershed
+fragments and waterz agglomeration.
+
+**Evaluation protocol.** The maintained validation protocol runs full-volume
+sliding-window inference on public CREMI A/B/C labeled volumes, applies
+CREMI-style boundary ignore (`xy=1`, `z=0`), and reports aggregate `voi_sum`
+and `adapted_rand_error` over A/B/C.
+
+## Post-processing Scope
+
+This repository is now closed around the dbMiM pretraining and UNETR affinity
+segmentation line. waterz remains the reference post-processing backend for
+headline VOI/ARAND results.
+
+Learnable and differentiable post-processing is being developed as a separate
+method line in:
+
+**https://github.com/ydchen0806/nnEM-Seg-diff-postprocess**
+
+Learned RAG, differentiable connected components, GPU watershed, and
+GPU/MWS-style solvers are maintained there so this repository can stay focused
+on the official dbMiM pretraining and affinity-segmentation implementation.
+
+The current best controlled result is **R48**, which uses the publicEM dbMiM
+pretrained encoder from the R16 line, seed-isolated finetuning (`seed=309`),
+and a longer 20k-step MSE+MAWS finetune. On the full CREMI A/B/C validation
+protocol it reaches `VOI=0.962154` and `ARAND=0.178252`, improving over both
+the same-seed plain MAE control and the earlier 12k finetune.
+
+The most useful full-data finding is that **fixed mixed edge/random masking on
+full EM data (R33)** improves over scratch, old fullEM dbMiM, pure edge masking,
+and fullEM plain MAE. It remains the best fullEM recipe, while R48 is the best
+absolute checkpoint so far.
+
+The recommended fullEM R33 line does **not** use a reinforcement-learning
+masking policy. It uses fixed mixed edge/random masking. RL-style decision
+modules are retained as ablations and are being stabilized in the R51/R52 line:
+R51 was an unhealthy policy-collapse run, while R52 uses constrained adaptive
+prior sampling. Its downstream R52/R53 finetunes finished successfully but were
+negative, so the current main line stays with R48-style publicEM dbMiM
+pretraining plus longer MSE+MAWS finetuning.
 
 ## Model Zoo
 
@@ -54,8 +107,9 @@ Weights are hosted at:
 
 | Model | HF path | Intended use |
 |---|---|---|
+| PublicEM dbMiM R48 finetune | `weights/publicem_dbmim_r48_seed309_long20k/finetuned_latest.pt` | Best current segmentation checkpoint |
 | PublicEM dbMiM R17 pretrain | `weights/publicem_dbmim_r17/pretrained_latest.pt` | ViT/dbMiM encoder checkpoint for UNETR initialization |
-| PublicEM dbMiM R17 finetune | `weights/publicem_dbmim_r17/finetuned_latest.pt` | Best current publicEM segmentation checkpoint |
+| PublicEM dbMiM R17 finetune | `weights/publicem_dbmim_r17/finetuned_latest.pt` | Earlier publicEM segmentation checkpoint |
 | FullEM mixed-mask dbMiM R33 pretrain | `weights/fullem_mixedmask_dbmim_r33/pretrained_latest.pt` | Recommended full-data dbMiM pretraining checkpoint |
 | FullEM mixed-mask dbMiM R33 finetune | `weights/fullem_mixedmask_dbmim_r33/finetuned_latest.pt` | Recommended full-data segmentation checkpoint |
 
@@ -103,9 +157,8 @@ The protocol is:
 - sweep calibration biases and waterz thresholds;
 - report aggregate A/B/C `voi_sum` and `adapted_rand_error`.
 
-This split is small but matches the controlled ablation goal: isolate whether
-dbMiM pretraining improves the same anisotropic UNETR finetuning recipe over
-scratch and plain MAE controls.
+This validation setup is used to isolate whether dbMiM pretraining improves the
+same anisotropic UNETR finetuning recipe over scratch and plain MAE controls.
 
 ## Results
 
@@ -118,8 +171,11 @@ post-processing thresholds.
 
 | Arm | VOI | ARAND at best VOI | Best ARAND | Conclusion |
 |---|---:|---:|---:|---|
-| R17 publicEM random-mask dbMiM | **1.002919** | **0.188832** | 0.188832 | Best publicEM VOI |
+| R48 publicEM dbMiM, seed309, 20k finetune | **0.962154** | **0.178252** | **0.178252** | Best current global result |
+| R45 publicEM dbMiM, seed309, 12k finetune | 0.986481 | 0.186187 | 0.186187 | Strong same-seed dbMiM result |
+| R17 publicEM random-mask dbMiM | 1.002919 | 0.188832 | 0.188832 | Earlier publicEM dbMiM result |
 | R23 publicEM random-mask plain MAE | 1.027073 | 0.192763 | 0.189247 | Matched MAE baseline |
+| R47 publicEM plain MAE, seed309 | 1.043065 | 0.190743 | 0.190743 | Same-seed MAE control for R45/R48 |
 | R29 publicEM pure edge-mask dbMiM | 1.033564 | 0.186827 | **0.186827** | Best publicEM ARAND, worse VOI |
 | R32 publicEM fixed mixed-mask dbMiM | 1.046538 | 0.206256 | 0.193183 | Negative vs R17/R23 |
 | R34 publicEM adaptive mixed dbMiM | 1.067471 | 0.205437 | 0.200604 | Negative adaptive result |
@@ -128,8 +184,12 @@ post-processing thresholds.
 
 Key deltas:
 
+- R48 dbMiM beats same-seed publicEM plain MAE R47 by `-0.0809` VOI and
+  `-0.0125` ARAND.
+- Extending the seed309 R16 finetune from the R45 12k schedule to the R48 20k
+  schedule improves VOI by `-0.0243` and ARAND by `-0.0079`.
 - R17 dbMiM beats matched publicEM plain MAE R23 by `-0.0242` VOI and about
-  `-0.0004` best ARAND.
+  `-0.0004` best ARAND under the older comparison.
 - R29 edge-mask dbMiM beats same-mask plain MAE R30 by `-0.0440` VOI and
   `-0.0117` best ARAND, but its VOI is worse than R17/R23.
 
@@ -162,6 +222,15 @@ edge fraction per crop. It did not improve downstream segmentation. After step
 fraction was `0.4456` for R34 and `0.3322` for R35. The current adaptive policy
 is therefore kept as a negative ablation rather than the recommended method.
 
+The newer R51/R52 line changes the policy stabilization instead of simply
+increasing reward strength. R51 still collapsed to an uninformative policy.
+R52 constrains mask-ratio and edge-fraction bins, uses a small edge-proxy reward
+with KL-to-prior regularization, clips and normalizes advantages, and freezes
+the policy after warmup. Its pretraining diagnostics were healthier than R51,
+but downstream finetuning was negative: R52 reached `VOI=1.056275`,
+`ARAND=0.194687`; R53 with `membrane_weight=1.0` reached `VOI=1.080240`,
+`ARAND=0.209107`. These are worse than R48 and are kept as negative ablations.
+
 ## Training Strategy
 
 ### Pretraining
@@ -209,7 +278,7 @@ Main settings:
 | Loss | MSE + MAWS, no BCE/Dice in the current winning recipe |
 | Label handling | synchronized image/label augmentation, 2D border widening radius 1 |
 | Batch size | 2 per GPU |
-| Schedule | 12k optimizer steps, lr `8e-5`, encoder lr `1e-5`, weight decay `0.01`, AMP |
+| Schedule | 12k optimizer steps for standard ablations; R48 uses 20k steps. lr `8e-5`, encoder lr `1e-5`, weight decay `0.01`, AMP |
 | Pretrained prefixes | `pos_embed`, `patch_embed`, `encoder_blocks`, `norm` |
 
 ### Evaluation