lloca.backbone.transformer_v2.Transformer

class lloca.backbone.transformer_v2.Transformer(in_channels, attn_reps, out_channels, num_blocks, num_heads, checkpoint_blocks=False, attention_factor=1, mlp_factor=2, dropout_prob=None, preserve_variance=True, elementwise_affine=True, compile=False, compile_kwargs=None)[source]

Bases: Module

LLoCa-Transformer with RMSNorm and GLU.

Combines transformer blocks, each consisting of multi-head self-attention layers, a gated-linear-unit MLP, residual connections, and RMSNorm layers.

Parameters:

  • in_channels (int) – Number of input channels.

  • attn_reps (str) – Representation of each attention head.

  • out_channels (int) – Number of output channels.

  • num_blocks (int) – Number of transformer blocks.

  • num_heads (int) – Number of attention heads.

  • checkpoint_blocks (bool) – Use gradient checkpointing for transformer blocks.

  • attention_factor (int) – Factor by which the key, query, and value size is increased over the default value of hidden_channels / num_heads.

  • mlp_factor (int) – Factor by which the activation size is increased over the default value of hidden_channels.

  • dropout_prob (float) – Dropout probability for output.

  • preserve_variance (bool) – Rescale the frame-to-frame transforms by the invariant Lorentz factor of each particle frame, to prevent the variance blowup from large boosts. Needs the reference momentum p_ref in forward().

  • elementwise_affine (bool) – Whether the RMSNorm layers use learnable per-channel affine weights.

  • compile (bool, optional) – Whether to compile the model with torch.compile, by default False.

  • compile_kwargs (Mapping, optional) – Dict forwarded verbatim to torch.compile() (via lloca.utils.compile.compile_model()) when compile=True (e.g. mode, dynamic, fullgraph). Omitted keys fall back to torch’s own defaults.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(inputs, frames, p_ref=None, ptr=None, **attn_kwargs)[source]

Forward pass.

Parameters:

  • inputs (Tensor) – Input data with shape (…, num_items, in_channels)

  • frames (Frames) – Local frames used for invariant particle attention

  • p_ref (Tensor, optional) – Reference (jet) 4-momentum in the global frame, energy-first: per event (..., 4) for a dense layout or per jet (num_jets, 4) with ptr for a packed layout. Required when a preserve_variance flag is on, ignored otherwise.

  • ptr (Tensor, optional) – Jet boundaries for a packed layout; maps the per-jet p_ref to each token.

  • **attn_kwargs

Returns:

outputs – Outputs with shape (…, num_items, out_channels)

Return type:

Tensor