EP‑01: Pluggable Training Infrastructure for sbi

[janfb]

Summary

This Enhancement Proposal (EP) proposes refactoring sbi’s training infrastructure to reduce duplication,
improve type safety, and enable pluggable logging and early stopping. It introduces
typed configurations, a unified training loop shared across NPE/NLE/NRE, and clean
interfaces for logging backends (TensorBoard, WandB, MLflow, stdout) and early stopping
strategies—while maintaining full backward compatibility.

Why

• Training logic is duplicated across methods (>3k LOC), making maintenance and
  contributions hard.
• Unstructured kwargs lead to runtime errors and poor IDE support.
• Logging is limited; early stopping is missing.
• Community contributions (e.g., PR Feature: adaptive learning rate schedulers #1629) are difficult to integrate under the current architecture.

Proposal

• Typed TrainConfig and method‑specific LossArgs for safer, clearer training
  configuration.
• Unified training loop backend used by NPE, NLE, NRE, VFPE.
• Pluggable logging backends via a simple configuration interface.
• Multiple early stopping strategies (patience, plateau detection, LR threshold).
• Backward compatible: existing .train(...) kwargs continue to work.
• Next step: extract the training loop to a standalone run_training(config, model, loss_fn, loaders, …) -> (model, summary) to avoid further bloating NeuralInference
  and improve testability/extensibility.

API sketch

from sbi.training import TrainConfig, LossArgsNPE, LoggingConfig, EarlyStopping
from sbi.inference import NPE

config = TrainConfig(batch_size=100, learning_rate=1e-3, max_epochs=1000)
logging = LoggingConfig(backend="wandb", project="sbi-experiment")
early_stop = EarlyStopping.validation_loss(patience=20)

inf = NPE(prior=prior)
neural_net = inf.train(config, LossArgsNPE(exclude_invalid_x=True),
                       logging=logging, early_stopping=early_stop)

# Next step (proposed extraction)
# run_training(config, model, loss_fn, train_loader, val_loader, optimizer, scheduler, callbacks)

What we’d like feedback on

1. Which logging backends matter most to you (TensorBoard, WandB, MLflow, stdout, other)?
2. Which early stopping strategies are essential (patience/min_delta, plateau w/ cooldown, LR threshold, custom callbacks)?
3. Any feedback on the idea of extracting the training loop entirely from the base.py
   class?

Links

• PR introducing the dataclasses for TrainConfig and LossArgs: Training dataclasses #1668
• PR #1629: Community early stopping implementation
• NUMFOCUS SDG Proposal: Related funding proposal

[janfb]

some discussion on this EP happened in the corresponding PR: #1674