Feature: adaptive learning rate schedulers

[christopherlovell]

Add Adaptive Learning Rate Schedulers to Training

Implements adaptive learning rate schedulers into the NPE, NLE and NRE methods. All changes maintain backward compatibility - existing code continues to work unchanged.

This feature was implemented with assistance from https://claude.ai/code
Co-Authored-By: Claude [email protected]

Changes Made

Core Implementation

• Enhanced base trainer (sbi/inference/trainers/base.py):
  • Added _create_lr_scheduler() method supporting 6 scheduler types: ReduceLROnPlateau, ExponentialLR,
    CosineAnnealingLR, StepLR, MultiStepLR, and CyclicLR
  • Updated _converged() method with learning rate threshold-based early stopping
  • Added learning rate tracking in training summaries
  • Extended abstract train() signature with scheduler parameters
• Updated NPE trainer (sbi/inference/trainers/npe/npe_base.py):
  • Integrated scheduler creation and stepping logic
  • Added learning rate logging to TensorBoard summaries
  • Maintained full backward compatibility with existing APIs

Testing

• Comprehensive test suite (tests/n*e_scheduler_test.py):
  • Parametrized tests for all scheduler types and configurations
  • Tests for convergence behavior and learning rate tracking
  • Error handling validation for invalid configurations

Documentation

• Practical guide (docs/how_to_guide/20_learning_rate_schedulers.md):
• Advanced tutorial (docs/advanced_tutorials/22_adaptive_learning_rates.ipynb):

API Usage

  # ReduceLROnPlateau with custom parameters
  posterior = inference.train(
      lr_scheduler="reduce_on_plateau",
      lr_scheduler_kwargs={"patience": 10, "factor": 0.5}
  )

  # Advanced configuration
  posterior = inference.train(
      lr_scheduler={
          "type": "cosine_annealing",
          "T_max": 50,
          "eta_min": 1e-6
      }
  )

[janfb]

Thanks a lot @christopherlovell for creating this PR, looks like a great addition. We will do a review of the proposed changes soon!

[janfb]

I did a first high-level pass over the changes - great additions overall! 👏

Before we go into the details, please have a look at my two main comments on the class structure for the LR scheduler options and their kwargs, and for the tests. thanks!

[janfb]

Instead of passing both, lr_scheduler and the corresponding kwargs, I suggest to use dataclasses like we recently introduced for the PosteriorParameters, see #1619. You would need to define a base class LrSchedulerParameters defining the interface and holder all shared parameters, and then a subclass for every LR scheduler type.

We should still offer the basic user to pass a string with the type, then we would just pick the corresponding parameter class and use the default. But as soon as the user wants to specify specific options, they should create their parameter class themselves and pass it here.

This will reduce some of the if-else below and strongly improve the type hints, IDE suggestions and debugging. Let me know if there are any questions and when I am missing something.

[janfb]

The fixture on top is great, but the tests are quite verbose and repetitive at the moment. I suggest to use pytest.mark.parametrize to test the lr schedulers for the different trainers (npe, nle, nre) and scheduler types all in one test file. as mnle is basically nle, I think it's fine to not include it here.

another think to include would be the vector field estimators, e.g., fmpe and npse, because they handle convergence criteria a bit differently, see #1544 .