Make sampler step count configurable for validation

validate_marginals and collect_samples could already vary flow ODE
steps for inference (giant predict --steps), but training-time
marginal validation and DDIM evaluation were stuck at hardcoded
defaults. Add a validate_steps config/CLI option and forward steps to
sample_ddim consistently.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

giant/analysis.py

@@ -39,8 +39,9 @@ Three tiers of checks, building on the aggregate marginal/KL check in
    generation artifact.
 
 `collect_samples` takes a `steps` argument (forwarded to the flow ODE
-integrator) so a later sampler-step-count ablation can sweep it by calling this
-function repeatedly without any new plumbing.
+integrator or, in ddim mode, the DDIM substep count) so a later
+sampler-step-count ablation can sweep it by calling this function repeatedly
+without any new plumbing.
 """
 
 from __future__ import annotations
@@ -205,11 +206,17 @@ def collect_samples(
     bundle: ModelBundle,
     val_loader: DataLoader,
     n_batches: int | None = None,
-    steps: int = 10,
+    steps: int | None = None,
 ) -> SampleCollection:
-    """Run the sampler over `val_loader`, pairing generations with real targets + conditioning."""
+    """Run the sampler over `val_loader`, pairing generations with real targets + conditioning.
+
+    `steps` overrides the number of sampler steps (flow ODE steps or DDIM
+    substeps); `None` keeps each sampler's own default. Unused in "ddpm"
+    mode, which always runs the full schedule.
+    """
     model = bundle.model
     device = bundle.device
+    steps_kw = {} if steps is None else {"steps": steps}
 
     cond_list, real_list, gen_list = [], [], []
     for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
@@ -219,11 +226,11 @@ def collect_samples(
         cond_cat = cond_cat.to(device)
 
         if bundle.mode == "flow":
-            gen = sample_flow(model, cond_cont, cond_cat, steps=steps)
+            gen = sample_flow(model, cond_cont, cond_cat, **steps_kw)
         elif bundle.mode == "ddpm":
             gen = sample_ddpm(model, cond_cont, cond_cat, bundle.schedule)
         else:
-            gen = sample_ddim(model, cond_cont, cond_cat, bundle.schedule)
+            gen = sample_ddim(model, cond_cont, cond_cat, bundle.schedule, **steps_kw)
 
         cond_full = torch.cat([cond_cont.cpu(), cond_cat.cpu().float()], dim=-1)
         cond_list.append(cond_full.numpy())

giant/cli.py

@@ -78,6 +78,13 @@ def train(
         Optional[int],
         typer.Option(help="Run marginal+KL validation every N epochs (0 disables)"),
     ] = None,
+    validate_steps: Annotated[
+        Optional[int],
+        typer.Option(
+            help="Flow matching ODE steps used during marginal validation "
+            "(ignored in ddpm mode, which always runs the full schedule)"
+        ),
+    ] = None,
     shuffle_buffer: Annotated[
         int, typer.Option(help="Rows held in RAM per worker for shuffling")
     ] = 65536,
@@ -105,6 +112,7 @@ def train(
             "num_workers": num_workers,
             "seed": seed,
             "validate_every": validate_every,
+            "validate_steps": validate_steps,
         }.items()
         if v is not None
     }

giant/config.py

@@ -18,6 +18,7 @@ DEFAULT_CONFIG: dict = {
         "num_workers": 4,
         "seed": 0,
         "validate_every": 10,
+        "validate_steps": 10,
     },
     "model": {
         "hidden_dim": 256,

giant/pipeline.py

@@ -152,4 +152,5 @@ def run_train_job(
         model_config=model_config,
         resume_path=resume,
         validate_every=t["validate_every"],
+        validate_steps=t["validate_steps"],
     )

giant/train.py

@@ -70,6 +70,7 @@ def train(
     model_config: dict | None = None,
     resume_path: str | Path | None = None,
     validate_every: int = 0,
+    validate_steps: int = 10,
 ) -> None:
     out_dir = Path(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -170,7 +171,12 @@ def train(
             if validate_every > 0 and epoch % validate_every == 0:
                 print(f"[epoch {epoch}] marginal validation:")
                 validate_marginals(
-                    model, val_loader, mode=mode, schedule=ddpm_schedule, device=device
+                    model,
+                    val_loader,
+                    mode=mode,
+                    schedule=ddpm_schedule,
+                    device=device,
+                    steps=validate_steps,
                 )
 
             ckpt: dict = {

giant/validate.py

@@ -6,6 +6,10 @@ from giant.constants import LOCAL_TARGET_NAMES
 from giant.sample import sample_flow, sample_ddpm, sample_ddim
 
 
+def _kw(steps: int | None) -> dict[str, int]:
+    return {} if steps is None else {"steps": steps}
+
+
 def _histogram_kl(
     p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
 ) -> float:
@@ -32,12 +36,17 @@ def validate_marginals(
     device: torch.device | None = None,
     n_batches: int | None = None,
     kl_bins: int = 50,
+    steps: int | None = None,
 ) -> dict[str, np.ndarray]:
     """Compare per-dimension marginals of generated vs. real steps.
 
     Returns {"real": (N,9), "generated": (N,9), "kl_divergence": (9,)} in
     normalised space. `kl_divergence[j]` is KL(real || generated) for
     dimension j, estimated from a shared histogram over both samples.
+
+    `steps` overrides the number of sampler steps (flow ODE steps or DDIM
+    substeps); `None` keeps each sampler's own default. Unused in "ddpm"
+    mode, which always runs the full schedule.
     """
     if device is None:
         device = next(model.parameters()).device
@@ -51,11 +60,11 @@ def validate_marginals(
         cond_cat = cond_cat.to(device)
 
         if mode == "flow":
-            gen = sample_flow(model, cond_cont, cond_cat)
+            gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
         elif mode == "ddpm":
             gen = sample_ddpm(model, cond_cont, cond_cat, schedule)
         else:
-            gen = sample_ddim(model, cond_cont, cond_cat, schedule)
+            gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
 
         all_real.append(x1.numpy())
         all_gen.append(gen.cpu().numpy())