MadnisTraining

class MadnisTraining

MadNIS training loop for one subprocess’s channels [1, 2, 3].

Alternates sampling batches from the channels’ Integrand / ChannelWeightNetwork (online) or a replay buffer of past samples (buffered) with AdamOptimizer steps on the stratified variance loss, periodically drops underperforming channels, and tracks progress in a status history for reporting. MultiMadnisTraining coordinates one MadnisTraining per subprocess.

References

Public Functions

MadnisTraining(ContextPtr generator_context, ContextPtr optimizer_context, const Config &config, const std::vector<std::shared_ptr<Integrand>> &integrands, const std::optional<ChannelWeightNetwork> &cwnet, std::optional<std::uint64_t> seed = std::nullopt, std::size_t channel_index_offset = 0)
Parameters:

  • generator_context – Context the sampling runtimes run on.

  • optimizer_context – Context the trained globals and optimizer live on.

  • config – Training configuration.

  • integrands – One integrand per channel to train.

  • cwnet – Optional shared channel-weight network.

  • seed – Top-level run seed.

  • channel_index_offset – Offset added to this subprocess’s local channel indices when deriving seeds, so MultiMadnisTraining’s subprocesses (which all share the same top-level seed) get non-overlapping DerivedSeed channel-index ranges.

inline const Config &config() const

The configuration passed to the constructor.

void train_step(std::size_t batch_index)

Run training batch batch_index: sample, optimize, and update the status history.

std::vector<std::size_t> active_channels() const

Indices of the channels not yet dropped.

inline std::size_t active_channel_count() const

Number of channels not yet dropped.

double average_loss() const

Mean loss over the recent training history.

double average_learning_rate() const

Mean learning rate over the recent training history.

double buffered_fraction() const

Fraction of recent steps taken on buffered samples.

inline std::size_t generated_event_count() const

Total number of samples generated so far.

std::size_t buffer_event_count() const

Total number of samples currently held in the replay buffers.

inline const std::vector<std::size_t> &status_batches() const

Batch index of each status-history entry, appended every config.log_interval batches.

inline const std::vector<double> &status_losses() const

Loss of each status-history entry.

inline const std::vector<std::size_t> &status_channel_counts() const

Active channel count of each status-history entry.

inline const std::vector<double> &status_learning_rates() const

Learning rate of each status-history entry.

inline const std::vector<double> &status_buffered_fractions() const

Buffered-step fraction of each status-history entry.

inline const std::vector<std::size_t> &status_generated_events() const

Cumulative generated-event count of each status-history entry.

inline const std::vector<std::size_t> &status_buffer_sizes() const

Buffer size of each status-history entry.

Public Static Functions

static inline void set_abort_check_function(std::function<void(void)> func)

Install func, called periodically during training to check for a requested abort.

struct Config

Tunable parameters of a MadnisTraining run.

Public Members

double learning_rate

Adam learning rate.

std::size_t batches

Number of training batches.

std::size_t log_interval

Batches between two status-history entries.

std::size_t integration_history_length

Number of integration-history entries kept per channel.

std::size_t channel_dropping_interval

Batches between two channel-dropping checks.

double channel_dropping_threshold

Minimum integral fraction below which a channel is dropped.

std::size_t cpu_generator_batch_size

Generator batch size on a CPU device.

std::size_t gpu_generator_batch_size

Generator batch size on a GPU device.

std::size_t gpu_generator_batch_granularity

Granularity generator batches on a GPU device are rounded to.

std::size_t generator_target_size_factor

Multiple of the training batch size kept buffered per channel.

std::size_t batch_size_offset

Constant added to every channel’s requested batch size.

std::size_t batch_size_per_channel

Minimum training-batch events requested per active channel.

double uniform_channel_ratio

Fraction of the training batch drawn uniformly across channels rather than proportionally to their integral.

AdamOptimizer::LRSchedule lr_schedule

Adam learning-rate schedule.

double adam_beta1

Adam first-moment decay rate.

double adam_beta2

Adam second-moment decay rate.

double adam_eps

Adam numerical-stability constant.

double adam_weight_decay

Adam L2 weight-decay coefficient.

double grad_clip_threshold

Maximum gradient norm; 0 disables clipping.

std::size_t buffer_capacity

Maximum number of samples kept in the replay buffer per channel; 0 disables buffering.

std::size_t minimum_buffer_size

Minimum buffered samples before buffered steps are taken.

double buffered_steps_fraction

Fraction of training steps done on buffered samples, reached once the buffers are full.

std::size_t buffer_skip_batches

Number of initial batches during which no samples are buffered.

double buffer_unweighting_quantile

Quantile used to estimate the max weight for buffer unweighting.

double fixed_cwnet_fraction

Fraction of training spent with the channel-weight network frozen, before it starts training too.

double softclip_threshold

Softclip threshold applied to the channel weights; 0 disables it.

std::size_t compressed_channel_weight_count

Number of channel weights kept per event when compressing the multi-channel weight loss.