bead.src.models package

Submodules

bead.src.models.flows module

Collection of flow strategies for variational inference and density estimation.

This module implements various normalizing flow architectures that can be used to transform simple probability distributions into more complex ones. These flows are particularly useful for improving the expressiveness of variational autoencoders by allowing more flexible posterior distributions.

Classes:

Planar: Planar flow transformation. Sylvester: Standard Sylvester normalizing flow. TriangularSylvester: Sylvester flow with triangular structure. IAF: Inverse Autoregressive Flow. CNN_Flow: Convolutional neural network based normalizing flow. NSF_AR: Neural Spline Flow with autoregressive structure.

class bead.src.models.flows.CNN_Flow(dim, cnn_layers, kernel_size, test_mode=0, use_revert=True)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.flows.IAF(z_size, num_flows=2, num_hidden=0, h_size=50, forget_bias=1.0, conv2d=False)

Bases: Module

forward(z, h_context)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.flows.NSF_AR(dim=15, K=64, B=3, hidden_dim=8, base_network=<class 'bead.src.models.layers.FCNN'>)

Bases: Module

Neural spline flow, auto-regressive. [Durkan et al. 2019]

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameters()

class bead.src.models.flows.Planar

Bases: Module

der_h(x)

Derivative of tanh

forward(zk, u, w, b)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.flows.Sylvester(num_ortho_vecs)

Bases: Module

Sylvester normalizing flow.

der_h(x)

der_tanh(x)

forward(zk, r1, r2, q_ortho, b, sum_ldj=True)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.flows.TriangularSylvester(z_size)

Bases: Module

Sylvester normalizing flow with Q=P or Q=I.

der_h(x)

der_tanh(x)

forward(zk, r1, r2, q_ortho, b, sum_ldj=True)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

bead.src.models.layers module

Custom layer implementations for neural network architectures.

This module provides specialized neural network layers used in various models across the BEAD framework. These include masked layers for autoregressive models, CNN flow layers, graph convolutional layers, and utility functions for spline-based flows.

Classes:

Identity: Simple identity layer that returns its input unchanged. MaskedLinear: Linear layer with masking for autoregressive architectures. MaskedConv2d: 2D convolutional layer with masking capabilities. CNN_Flow_Layer: Base layer for CNN-based normalizing flows. Dilation_Block: Block of dilated convolutions for CNN flows. GraphConvolution: Graph convolutional network layer. FCNN: Simple fully connected neural network. Log1pScaler: Scaler that applies log(1+x) transformation. L2Normalizer: Scaler that applies L2 normalization. SinCosTransformer: Transforms angles to sin/cos features. ChainedScaler: Chains multiple scalers together.

Functions:

searchsorted: Utility for finding indices where elements should be inserted. unconstrained_RQS: Rational quadratic spline transformation with unconstrained inputs. RQS: Rational quadratic spline transformation.

class bead.src.models.layers.CNN_Flow_Layer(dim, kernel_size, dilation, test_mode=0, rescale=True, skip=True)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.layers.Dilation_Block(dim, kernel_size, test_mode=0)

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.layers.EFPEmbedding(n_efp_features: int, embedding_dim: int = 64, gate_type: str = 'sigmoid', gate_threshold: float = 0.05, dropout_rate: float = 0.1, use_layer_norm: bool = True, monitor_sparsity: bool = True)

Bases: Module

Energy-Flow Polynomial (EFP) Embedding Layer.

Transforms high-dimensional EFP features into compact embedding tokens suitable for transformer architectures and other downstream models. Provides learnable feature selection through gated sparsification and dimensionality reduction.

Scientific Motivation:

EFP features are high-dimensional (140-531 features per jet) and potentially redundant. This embedding layer enables: - Compression: Reduces dimensionality while preserving information - Selection: Learnable gating mechanism for feature importance - Sparsification: Threshold-based pruning to prevent overfitting - Standardization: Consistent output format for downstream architectures

Architecture:

Input: (batch_size, n_jets, n_efp_features) ↓ Linear projection (dimensionality reduction) ↓ Gated sparsification (learnable feature selection) ↓ Layer normalization (training stability) ↓ Dropout (regularization) Output: (batch_size, n_jets, embedding_dim)

Parameters:

• n_efp_features (int) – Number of input EFP features (140 or 531)

• embedding_dim (int) – Output embedding dimension (default: 64)

• gate_type (str) – Gate activation function (‘sigmoid’, ‘relu6’, ‘tanh’)

• gate_threshold (float) – Sparsification threshold (default: 0.05)

• dropout_rate (float) – Dropout rate for regularization (default: 0.1)

• use_layer_norm (bool) – Enable layer normalization (default: True)

• monitor_sparsity (bool) – Track gate activation statistics (default: True)

Example

>>> embedding = EFPEmbedding(n_efp_features=140, embedding_dim=64)
>>> efp_features = torch.randn(32, 3, 140)  # batch_size=32, n_jets=3
>>> embeddings = embedding(efp_features)
>>> print(embeddings.shape)  # torch.Size([32, 3, 64])

extra_repr() → str

Extra representation string for the module.

Returns:

String representation of the layer configuration

Return type:

str

forward(efp_features: Tensor, jet_mask: Tensor | None = None) → Tensor

Forward pass through the EFP embedding layer.

Parameters:

• efp_features (torch.Tensor) – Input EFP features of shape (B, J, N_efp)

• jet_mask (torch.Tensor, optional) – Jet validity mask of shape (B, J) True for valid jets, False for padded jets

Returns:

Embedded EFP tokens of shape (B, J, embedding_dim)

Return type:

torch.Tensor

Raises:

ValueError – If input tensor has incorrect shape or contains invalid values

get_sparsity_stats() → dict

Get current sparsity statistics.

Returns:

Dictionary containing sparsity metrics:

• ’sparsity_ratio’: Fraction of gates below threshold

• ’activation_ratio’: Fraction of gates above threshold

• ’total_gates_seen’: Total number of gates processed

Return type:

dict

reset_sparsity_stats()

Reset sparsity monitoring statistics.

class bead.src.models.layers.FCNN(in_dim, out_dim, hidden_dim)

Bases: Module

Simple fully connected neural network.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.layers.GraphConvolution(in_features, out_features, bias=True)

Bases: Module

Simple GCN layer, similar to https://arxiv.org/abs/1609.02907

forward(input, adj)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameters()

class bead.src.models.layers.Identity

Bases: Module

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.layers.MaskedConv2d(in_features, out_features, size_kernel=(3, 3), diagonal_zeros=False, bias=True)

Bases: Module

build_mask()

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameters()

class bead.src.models.layers.MaskedLinear(in_features, out_features, diagonal_zeros=False, bias=True)

Bases: Module

build_mask()

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameters()

bead.src.models.layers.RQS(inputs, unnormalized_widths, unnormalized_heights, unnormalized_derivatives, inverse=False, left=0.0, right=1.0, bottom=0.0, top=1.0, min_bin_width=0.001, min_bin_height=0.001, min_derivative=0.001)

bead.src.models.layers.searchsorted(bin_locations, inputs, eps=1e-06)

bead.src.models.layers.unconstrained_RQS(inputs, unnormalized_widths, unnormalized_heights, unnormalized_derivatives, inverse=False, tail_bound=1.0, min_bin_width=0.001, min_bin_height=0.001, min_derivative=0.001)

bead.src.models.models module

Neural network model architectures for anomaly detection.

This module provides various autoencoder and variational autoencoder architectures with different latent space configurations, flow transformations, and architectural choices. These models can be used for anomaly detection in particle physics data.

Classes:

AE: Basic autoencoder architecture. AE_Dropout_BN: Autoencoder with dropout and batch normalization. ConvAE: Convolutional autoencoder. ConvVAE: Convolutional variational autoencoder. Dirichlet_ConvVAE: Convolutional Dirichlet variational autoencoder. Planar_ConvVAE: ConvVAE with planar normalizing flows. OrthogonalSylvester_ConvVAE: ConvVAE with orthogonal Sylvester flows. HouseholderSylvester_ConvVAE: ConvVAE with Householder Sylvester flows. TriangularSylvester_ConvVAE: ConvVAE with triangular Sylvester flows. IAF_ConvVAE: ConvVAE with inverse autoregressive flows. ConvFlow_ConvVAE: ConvVAE with convolutional normalizing flows. NSFAR_ConvVAE: ConvVAE with neural spline flows. TransformerAE: Autoencoder with transformer components. FlexibleTransformer: Flexible transformer model that can integrate with any VAE model. VAEWithTransformer: VAE model with integrated transformer.

class bead.src.models.models.AE(in_shape, z_dim, *args, **kwargs)

Bases: Module

decode(z)

detach_hooks(hooks: list) → None

encode(x)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

get_activations() → dict

get_hook(layer_name)

get_layers() → list

store_hooks() → list

class bead.src.models.models.AE_Dropout_BN(in_shape, z_dim, *args, **kwargs)

Bases: AE

dec_bn(z)

enc_bn(x)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.ConvAE(in_shape, z_dim, *args, **kwargs)

Bases: Module

decode(z)

encode(x)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.ConvFlow_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with convolutional flows in the decoder.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvAE

decode(z)

encode(x)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reparameterize(mean, logvar)

class bead.src.models.models.Dirichlet_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvAE

decode(z)

encode(x)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reparameterize(mean, logvar)

class bead.src.models.models.FlexibleTransformer(latent_dim: int, n_efp_features: int, efp_embedding_dim: int, output_dim: int, d_model: int = 256, n_heads: int = 8, n_layers: int = 6, d_ff: int = 2048, dropout: float = 0.1, activation: callable = <built-in function gelu>, norm_first: bool = True, use_class_attention: bool = True, max_jets: int = 3, output_activation: callable | None = None, efp_config: dict | None = None)

Bases: Module

Flexible transformer model that can integrate with any VAE model.

This model is designed to be used with any VAE model, processing both VAE latent vectors and raw EFP features. It uses the transformer utilities from transformer_utils.py to build a flexible and modular architecture.

The model can be used in different modes: 1. VAE-only mode: Only process VAE latent vectors 2. EFP-only mode: Only process EFP features 3. Combined mode: Process both VAE latent vectors and EFP features

Parameters:

• latent_dim – Dimension of the VAE latent space.

• n_efp_features – Number of raw EFP features (e.g., 140 or 531).

• efp_embedding_dim – Dimension of the EFP embeddings after compression.

• output_dim – Dimension of the output space.

• d_model – Dimension of the transformer model.

• n_heads – Number of attention heads.

• n_layers – Number of transformer layers.

• d_ff – Dimension of the feed-forward network.

• dropout – Dropout probability.

• activation – Activation function.

• norm_first – Whether to apply normalization before or after attention and feed-forward.

• use_class_attention – Whether to use class attention pooling for the output.

• max_jets – Maximum number of jets per event.

• output_activation – Activation function for the output layer.

• efp_config – Optional configuration dict for EFPEmbedding layer.

forward(latent_z: Tensor | None = None, efp_features: Tensor | None = None, jet_mask: Tensor | None = None) → Tensor

Forward pass.

Parameters:

• latent_z – Optional VAE latent vector of shape (batch_size, latent_dim)

• efp_features – Optional raw EFP features of shape (batch_size, n_jets, n_efp_features)

• jet_mask – Optional mask tensor of shape (batch_size, n_jets)

Returns:

Output tensor of shape (batch_size, output_dim)

Return type:

output

class bead.src.models.models.HouseholderSylvester_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with householder sylvester flows in the decoder.

batch_construct_orthogonal(q)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.IAF_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with inverse autoregressive flows in the decoder.

encode(x)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.NSFAR_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with auto-regressive neural spline flows in the decoder.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.OrthogonalSylvester_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with orthogonal flows in the decoder.

batch_construct_orthogonal(q)

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.Planar_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with planar flows in the decoder.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.TransformerAE(in_dim, h_dim=256, n_heads=1, latent_size=50, activation=<built-in function gelu>)

Bases: Module

Autoencoder mixed with the Transformer Encoder layer

Parameters:

nn (_type_) – _description_

decoder(z: Tensor)

_summary_

Parameters:

z (_type_) – _description_

Returns:

_description_

Return type:

_type_

encoder(x: Tensor)

_summary_

Parameters:

x (_type_) – _description_

Returns:

_description_

Return type:

_type_

forward(x: Tensor)

_summary_

Parameters:

z (_type_) – _description_

Returns:

_description_

Return type:

_type_

class bead.src.models.models.TriangularSylvester_ConvVAE(in_shape, z_dim, *args, **kwargs)

Bases: ConvVAE

Variational auto-encoder with triangular sylvester flows in the decoder.

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class bead.src.models.models.VAEWithTransformer(vae_model: Module, n_efp_features: int, efp_embedding_dim: int, output_dim: int, transformer_config: dict | None = None, efp_config: dict | None = None)

Bases: Module

VAE model with integrated transformer.

This model combines a VAE with a transformer, allowing for flexible integration of VAE latent vectors and raw EFP features. The VAE can be any VAE model, and the transformer is built using the transformer utilities from transformer_utils.py.

Parameters:

• vae_model – VAE model to use.

• n_efp_features – Number of raw EFP features (e.g., 140 or 531).

• efp_embedding_dim – Dimension of the EFP embeddings after compression.

• output_dim – Dimension of the output space.

• transformer_config – Configuration for the transformer.

• efp_config – Optional configuration dict for EFPEmbedding layer.

decode(z)

Decode latent vector through the VAE decoder.

Parameters:

z – Latent vector

Returns:

Reconstructed input

Return type:

x_recon

encode(x)

Encode input through the VAE encoder.

Parameters:

x – Input tensor

Returns:

Latent vector mu: Mean of the latent distribution (if VAE) logvar: Log variance of the latent distribution (if VAE)

Return type:

z

forward(x: Tensor, efp_features: Tensor | None = None, jet_mask: Tensor | None = None, return_all: bool = False)

Forward pass.

Parameters:

• x – Input tensor

• efp_features – Optional EFP features

• jet_mask – Optional mask tensor for jets

• return_all – Whether to return all outputs (reconstruction, transformer output, mu, logvar)

Returns:

transformer_output: Output of the transformer If return_all=True:

reconstruction: Reconstructed input transformer_output: Output of the transformer mu: Mean of the latent distribution (if VAE) logvar: Log variance of the latent distribution (if VAE)

Return type:

If return_all=False

transform(x: Tensor, efp_features: Tensor | None = None, jet_mask: Tensor | None = None) → Tensor

Transform input through the VAE encoder and transformer.

This is a convenience method for inference.

Parameters:

• x – Input tensor

• efp_features – Optional EFP features

• jet_mask – Optional mask tensor for jets

Returns:

Output of the transformer

Return type:

transformer_output

bead.src.models.models.create_flexible_transformer_from_config(config, latent_dim: int) → FlexibleTransformer

Create FlexibleTransformer model from BEAD Config object.

Parameters:

• config – BEAD Config object containing transformer parameters

• latent_dim – Dimension of the VAE latent space

Returns:

Configured transformer model

Return type:

FlexibleTransformer

Example

>>> from bead.src.utils.ggl import Config
>>> config = Config(...)  # Your config with transformer settings
>>> transformer = create_flexible_transformer_from_config(config, latent_dim=20)

bead.src.models.models.create_transformer_config_from_config(config) → dict

Create transformer configuration dictionary from BEAD Config object.

This utility function extracts transformer-related parameters from the main BEAD configuration and creates a dictionary suitable for initializing FlexibleTransformer or VAEWithTransformer models.

Parameters:

config – BEAD Config object containing transformer parameters

Returns:

Configuration dictionary for transformer initialization

Return type:

dict

Example

>>> from bead.src.utils.ggl import Config
>>> config = Config(...)  # Your config with transformer settings
>>> transformer_config = create_transformer_config_from_config(config)
>>> flexible_transformer = FlexibleTransformer(
...     latent_dim=20,
...     efp_embedding_dim=config.efp_embedding_dim,
...     output_dim=config.transformer_output_dim,
...     **transformer_config
... )

bead.src.models.models.create_vae_with_transformer_from_config(config, vae_model: Module) → VAEWithTransformer

Create VAEWithTransformer model from BEAD Config object.

Parameters:

• config – BEAD Config object containing transformer parameters

• vae_model – Pre-initialized VAE model to integrate with transformer

Returns:

VAE model with integrated transformer

Return type:

VAEWithTransformer

Example

>>> from bead.src.utils.ggl import Config
>>> from bead.src.models.models import ConvVAE
>>> config = Config(...)  # Your config with transformer settings
>>> vae = ConvVAE(in_shape=(1, 28, 28), z_dim=20)
>>> vae_transformer = create_vae_with_transformer_from_config(config, vae)

bead.src.models.models.get_activation_function(activation_str: str)

Map activation function string to callable.

bead.src.models.models.get_output_activation_function(activation_str: str)

Map output activation function string to callable.

Module contents