lloca.backbone.attention.LLoCaAttention

class lloca.backbone.attention.LLoCaAttention(attn_reps, num_heads)

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

forward(q_local, k_local, v_local, **attn_kwargs)

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)

Prepare local frames for processing with LLoCa attention. For a single forward pass through the network, this method is called only once for efficiency.

Parameters:

frames (torch.tensor) – Local frames of reference for each particle of shape (…, N, 4, 4) where N is the number of particles.