lloca.backbone.attention.LLoCaAttention

class lloca.backbone.attention.LLoCaAttention(attn_reps, num_heads, preserve_variance=True, variance_eps=0.01)[source]

Bases: Module

Attention with frame-to-frame transformations.

Parameters:

  • attn_reps (TensorReps) – Tensor representation of a single attention head.

  • num_heads (int) – Number of attention heads

  • preserve_variance (bool) – Rescale the pre-attention (local->global) q/k/v and post-attention (global->local) vectors by 1/gamma_i^grade to prevent the variance blowup from large boosts. Needs the reference momentum p_ref in prepare_frames().

  • variance_eps (float) – Small mass floor (energy units) that keeps gamma_i finite for near-lightlike jets.

forward(q_local, k_local, v_local, **attn_kwargs)[source]

Execute LLoCa attention.

Strategy 1) Transform q, k, v into global frame 2) Apply attention in global frame 3) Transform output back into local frame

Comments - Dimensions: … (optional), H (head), N (particles), C (channels). - Extension to cross-attention is trivial but we don’t have this right now for convenience. Strategy: frames_q for queries (in contrast to frames=frames_kv).

Parameters:

  • q_local (torch.tensor) – Local queries of shape (…, H, N, C)

  • k_local (torch.tensor) – Local keys of shape (…, H, N, C)

  • v_local (torch.tensor) – Local values of shape (…, H, N, C)

  • **attn_kwargs – Optional arguments that are passed on to the attention backend

Returns:

out_local – Attention output in local frame of shape (…, H, N, C)

Return type:

torch.tensor

prepare_frames(frames, p_ref=None, ptr=None)[source]

Prepare local frames for LLoCa attention (called once per forward pass).

Parameters:

  • frames (Frames) – Local frames of shape (…, N, 4, 4).

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

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