ChannelWeightNetwork

class ChannelWeightNetwork : public madspace::FunctionGenerator

Learned multi-channel weights from a neural network.

The MadNIS approach to multi-channeling: an MLP predicts the per-channel weights \(\alpha_i(x)\) from the MomentumPreprocessing features and a prior, replacing the fixed PropagatorChannelWeights [1, 2]. The network weights are compute-graph globals named after prefix. Call initialize_globals(context) once before use.

batch is the leading batch dimension. c is channel_count.

Arguments

• input – float, shape (batch, MomentumPreprocessing::output_dim()) – the invariant features.

• prior – float, shape (batch, c) – the prior per-channel weights.

Returns

• channel_weights – float, shape (batch, c) – the learned per-channel weights.

References

• [1] T. Heimel, O. Mattelaer, R. Winterhalder, “MadSpace”, https://arxiv.org/abs/2602.06895 (Sec. 2.1.1)

• [2] T. Heimel et al., “MadNIS - Neural multi-channel importance sampling”, https://arxiv.org/abs/2212.06172

Public Functions

ChannelWeightNetwork(std::size_t channel_count, std::size_t particle_count, std::size_t hidden_dim = 32, std::size_t layers = 3, MLP::Activation activation = MLP::leaky_relu, const std::string &prefix = "", bool include_preprocessing = true)

Parameters:

• channel_count – Number of channels.

• particle_count – Number of external particles.

• hidden_dim – Hidden width of the MLP.

• layers – Number of MLP layers.

• activation – MLP activation.

• prefix – Prefix for the trainable global names.

• include_preprocessing – If true, apply MomentumPreprocessing to the input first.

inline const MLP &mlp() const

The underlying network.

inline const MomentumPreprocessing &preprocessing() const

The feature preprocessing.

void initialize_globals(ContextPtr context, std::optional<std::uint64_t> seed = std::nullopt) const

Register the network’s trainable parameters on context.

Parameters:

• context – Context receiving the globals.

• seed – Seed for the random initialization. If unset (the default), the initialization is non-deterministic.

inline const std::string &mask_name() const

Global name of the channel mask.