vesuvius.train: generate configured auxiliary targets instead of training them against nothing

[CVasilopoulos]

In one sentence: Auxiliary tasks in a training config (distance_transform, surface_normals, ...) now get real targets and train, and trainers that cannot make those targets stop at startup instead of reporting Avg Loss = 0.0000.

One real example: Starting with the Scroll 1 instance-label cube 01744_02256_04048 (256^3, label = instance id > 0) and configuration/aux_tasks/surface_normals_3d.yaml plus the distance_transform entry from #1490, I ran vesuvius.train for 5 train and 2 validation steps, and distance_transform and surface_normals reported train loss 0.2130 and 0.2927 instead of 0.0000.

Before: ZarrDataset only loads primary targets and BaseTrainer does not derive auxiliary ones. The auxiliary heads were built but never had a target, and every epoch printed Avg Loss = 0.0000 for them with no warning. That is what #1490 reports: two runs of a 4-run ablation were plain baselines with dead decoders.

The code that derives these targets already exists. #422 moved it out of the dataset into trainers/auxiliary/, where each trainer adds its target in _prepare_batch. But vesuvius.train never selects those trainers, and each accepts only one task type, so a config with both distance_transform and surface_normals could not use them.

After this PR:

• AuxiliaryTrainer sends each configured auxiliary target to the existing trainer for its task_type.
• vesuvius.train with the default --trainer base uses it when auxiliary targets are configured, so the config from auxiliary_tasks silently train nothing with ZarrDataset: aux targets filtered out, no derived-tensor generation, loss reports 0.0000 #1490 trains all three heads without new flags.
• Any trainer that does not generate the configured auxiliary targets raises at the start of _initialize_training, before datasets are built, with an error naming them.
• The docs no longer say the dataset generates these tensors.

Proof: Upstream main 4b3c72882 against this branch on the same base, same data, config and flags, CPU only. The attached images show the evidence script's terminal output. Its output is filtered to the trainer, batch and loss lines. show_batch.py is a small wrapper that prints the tensor keys of the first batch before and after _prepare_batch, then calls vesuvius.train's main unchanged. The vesuvius code:, EXIT= and -- lines are printed by the scripts. Look at the Avg Loss lines.

run: vesuvius.train --config aux-1490.yaml -i data -o out --max-epoch 1 --max-steps-per-epoch 5 --max-val-steps-per-epoch 2 --no-amp --seed 42 (CPU)

1. Before (main), default --trainer base

vesuvius code: /w/villa-1490-base/vesuvius/src/vesuvius/models/training/train.py
Using Base Trainer for supervised training
first batch from dataset: image(1, 1, 64, 64, 64), padding_mask(1, 1, 64, 64, 64), surface(1, 1, 64, 64, 64), is_unlabeled(1,), surface_skel(1, 1, 64, 64, 64)
first batch after _prepare_batch: image(1, 1, 64, 64, 64), padding_mask(1, 1, 64, 64, 64), surface(1, 1, 64, 64, 64), is_unlabeled(1,), surface_skel(1, 1, 64, 64, 64)
[Train] Epoch 1 completed.
  surface: Avg Loss = -0.6954
  distance_transform: Avg Loss = 0.0000
  surface_normals: Avg Loss = 0.0000
[Validation] Epoch 1 summary:
  Task 'surface': Avg validation loss = -0.9476
  Task 'distance_transform': Avg validation loss = 0.0000
  Task 'surface_normals': Avg validation loss = 0.0000
EXIT=0

2. After (this PR), default --trainer base

vesuvius code: /w/villa-1490/vesuvius/src/vesuvius/models/training/train.py
Using Auxiliary Trainer for supervised training with auxiliary targets
first batch from dataset: image(1, 1, 64, 64, 64), padding_mask(1, 1, 64, 64, 64), surface(1, 1, 64, 64, 64), is_unlabeled(1,), surface_skel(1, 1, 64, 64, 64)
first batch after _prepare_batch: image(1, 1, 64, 64, 64), padding_mask(1, 1, 64, 64, 64), surface(1, 1, 64, 64, 64), is_unlabeled(1,), surface_skel(1, 1, 64, 64, 64), distance_transform(1, 1, 64, 64, 64), surface_normals(1, 3, 64, 64, 64)
  distance_transform: min -5.196 max 49.346
  surface_normals: min -1.000 max 1.000
[Train] Epoch 1 completed.
  surface: Avg Loss = -0.4579
  distance_transform: Avg Loss = 0.2130
  surface_normals: Avg Loss = 0.2927
[Validation] Epoch 1 summary:
  Task 'surface': Avg validation loss = -1.0870
  Task 'distance_transform': Avg validation loss = 1.9482
  Task 'surface_normals': Avg validation loss = 1.0957
EXIT=0

3. After (this PR), same config with --trainer mean_teacher, which does not generate auxiliary targets

Using Regular Mean Teacher Trainer for semi-supervised training
ValueError: TrainMeanTeacher cannot generate auxiliary targets ['distance_transform', 'surface_normals']; their heads would train against nothing and report Avg Loss = 0.0000. Use AuxiliaryTrainer (vesuvius.train selects it for --trainer base) or remove these auxiliary_tasks.
EXIT=1

4. Tests

tests/models/training/trainers/test_aux_trainers.py + tests/models/training/test_base_trainer.py
-- BEFORE  11 passed
-- AFTER   15 passed
tests/models
-- BEFORE  182 passed, 1 skipped, 3 deselected
-- AFTER   186 passed, 1 skipped, 3 deselected

Why / where this is useful:

I picked this up because configs with auxiliary tasks were silently training those heads against nothing, and the only earlier fixes just made the run stop. I ran the before and after on a real Scroll 1 labelled cube and checked that the auxiliary losses go from 0.0000 to real values, and that a trainer that can't make the targets stops at startup.

Auxiliary heads such as signed distance and surface normals are meant to give the surface model extra geometric supervision for tracing papyrus sheets. Anyone running ablations or training with them, as in #1490, now gets heads that learn from targets derived per batch, and the config shipped in configuration/aux_tasks/ works with the default command. If they pick a trainer that cannot derive the targets, the run stops before it spends GPU time.

• I personally verified that the example and proof above were produced by this PR on the stated data.

Details

What changed

• trainers/auxiliary/auxiliary_trainer.py (new): AuxiliaryTrainer(BaseAuxTrainer) maps task_type to the five existing single-type trainers and calls their _compute_aux_tensor. It raises at construction for an unknown type. No generation code is duplicated.
• trainers/auxiliary/base_aux_trainer.py: _generated_auxiliary_targets() returns the auxiliary targets the trainer will fill.
• train.py: BaseTrainer._generated_auxiliary_targets() returns an empty set, and _check_auxiliary_targets() raises when a configured auxiliary target is not in it. _initialize_training calls it first. All trainers that override _initialize_training call super() first, so the check covers them too.
• cli.py: only the elif trainer_name == "base": branch changes. It picks AuxiliaryTrainer when any target has auxiliary_task: true.
• docs/training_flow.md, docs/data_formatting.md: the lines that said the dataset derives these tensors now say which trainer does, and that other trainers stop at startup.
• Tests: AuxiliaryTrainer output equals the single-type trainers' output for a mixed config, _prepare_batch stacks both tasks with the right channel counts, an unknown type raises, and _initialize_training raises before _configure_dataset is called.

Why the trainer and not the dataset, and why automatic selection

The generators already live in trainers, and they build ZarrDataset through _build_dataset_for_mgr. A check inside ZarrDataset.__init__ would also stop those trainers. Putting it in the trainer lets each trainer declare what it generates.

auxiliary_tasks is set in the YAML and the docs never mention a trainer flag, so the default command should train them. A new --trainer value would mean the #1490 config still fails until users find it, and it would add another option to the list. If you prefer an explicit --trainer auxiliary, it is a two-line change.

Tested

• Base: upstream main 4b3c72882. Linux x86_64 only, python:3.14-slim container, torch 2.14.0+cpu, zarr 3.3.0, CPU only, 8 cores. Not tested on GPU, macOS or Windows.
• Data: https://dl.ash2txt.org/full-scrolls/Scroll1/PHercParis4.volpkg/volumetric-instance-labels/instance-labels/01744_02256_04048/ (_volume.nrrd 15 MB, _mask.nrrd 0.8 MB, 9 instance ids). Written unchanged to data/images/01744_02256_04048.zarr, with data/labels/01744_02256_04048_surface.zarr = mask > 0 (23.7% of voxels).
• Config: configuration/aux_tasks/surface_normals_3d.yaml plus the distance_transform entry from auxiliary_tasks silently train nothing with ZarrDataset: aux targets filtered out, no derived-tensor generation, loss reports 0.0000 #1490 (SignedDistanceLoss, distance_type: signed), batch_size: 1, patch_size: [64, 64, 64], deep supervision on.
• The schedule flags are passed on the command line because on main the argparse defaults override tr_config (vesuvius.train: argparse schedule defaults (--max-epoch 1000 etc.) silently override YAML tr_config values #1489, open PR vesuvius.train: stop argparse schedule defaults from overriding tr_config (fixes #1489) #1702).
• Tests: pytest tests/models in the same container.

Limitations

• --seed pins the train/val split but not patch sampling, so the two runs trained on different patches (see patch in evidence.log). The comparison is about whether the auxiliary heads get a target at all, not about the loss values.
• Auxiliary targets are derived in the main process after augmentation, one patch at a time, as in the existing aux trainers. With augment_on_device: true, _prepare_batch runs before the on-device transforms, so vector targets such as normals would be transformed like images. I did not change that.
• vesuvius.train --verbose fails at the first-batch printout on main already, with AttributeError: 'list' object has no attribute 'dtype' on patch_info. It fails the same way with this change. I left it out of this PR.
• CrossFrameZarrDataset with auxiliary tasks goes through the same trainer path, but I did not run it.

Prior work and credit

• @flummoxjr reported auxiliary_tasks silently train nothing with ZarrDataset: aux targets filtered out, no derived-tensor generation, loss reports 0.0000 #1490 with the root cause in the dataset and suggested either generation or a fail-fast error.
• Fail fast when ZarrDataset cannot generate auxiliary targets #1491 (@robertlangdonn, closed as stale) and @ge-al's branch fix/aux-tasks-fail-fast (eb53bef) raise in ZarrDataset (and, in the branch, CrossFrameZarrDataset) when auxiliary tasks are configured. The startup error here follows that idea but lives in the trainer, because the aux trainers build ZarrDataset themselves. Fail fast when ZarrDataset cannot generate auxiliary targets #1491 also changed the docs to say generation does not exist. The generators do exist in trainers/auxiliary/, so the docs here point to them. No code from either is used.
• Dataset refactor #422 (@bruniss) added trainers/auxiliary/ and its tests. This PR only wires them up.
• Open PR vesuvius.train: stop argparse schedule defaults from overriding tr_config (fixes #1489) #1702 also edits cli.py, in the argument parser. The change here is in the trainer dispatch further down.

AI-assisted (Claude Code), human-directed. All runs above were executed on real data, not inferred.
Fixes #1490

[vercel]

@CVasilopoulos is attempting to deploy a commit to the scroll Team on Vercel.

A member of the Team first needs to authorize it.