"""Self-contained ParticleViT model definition for inference. ParticleViT is a generic transformer over the constituents of a jet, with no physics-specific inductive bias. Each particle is one token; a single prepended class token is read out for classification. There is no positional encoding. This file depends only on PyTorch (no einops / flash-attn), so it runs on CPU or any GPU. It reproduces, parameter-for-parameter, the model trained for the paper "Predict before you train: scaling laws for particle physics foundation models"; the released weights load with strict=True. Usage: from modeling_particlevit import ParticleViT model = ParticleViT.from_pretrained("jaluus/ParticleViT-S").eval() logits = model(X, attn_mask=mask) # X: (B, 150, 9), mask: (B, 150) bool """ from __future__ import annotations import json from pathlib import Path import torch import torch.nn as nn import torch.nn.functional as F def _zero_masked_tokens(x: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor: """Zero out tokens where attn_mask is False (padding positions).""" if attn_mask is None: return x return x.masked_fill(~attn_mask.unsqueeze(-1), 0.0) class ParticleInputEmbedding(nn.Module): """Embed kinematics, PID, and vertex features through separate paths. The 9 input features per particle are: 0:4 continuous kinematics (delta eta, delta phi, log pT, log E) 4 categorical particle-ID code (dense integer id in [0, 8]) 5:9 continuous vertex / tracking features (zero where unavailable) """ def __init__(self, num_features: int, embed_dim: int) -> None: super().__init__() if num_features != 9: raise ValueError(f"ParticleInputEmbedding expects 9 features, got {num_features}") self.kin_embed = nn.Linear(4, embed_dim) self.pid_embed = nn.Embedding(9, embed_dim) self.vertex_embed = nn.Linear(4, embed_dim) self.pid_available_embed = nn.Embedding(2, embed_dim) self.vertex_available_embed = nn.Embedding(2, embed_dim) def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None) -> torch.Tensor: if attn_mask is None: real_mask = X[:, :, 2] != 0 else: real_mask = attn_mask.bool() vertex_available = (X[:, :, 5:9] != 0).any(dim=-1) & real_mask pid_available = vertex_available.any(dim=1, keepdim=True).expand_as(real_mask) x = self.kin_embed(X[:, :, :4]) dtype = x.dtype x = x + self.pid_embed(X[:, :, 4].long()).to(dtype=dtype) x = x + self.vertex_embed(X[:, :, 5:9]) x = x + self.pid_available_embed(pid_available.long()).to(dtype=dtype) x = x + self.vertex_available_embed(vertex_available.long()).to(dtype=dtype) return x class SwiGLU(nn.Module): """Gated-linear-unit feedforward with the 8/3 width convention.""" def __init__(self, dim: int, mlp_ratio: float = 8 / 3) -> None: super().__init__() hidden_dim = int(dim * mlp_ratio) self.w13 = nn.Linear(dim, 2 * hidden_dim, bias=False) self.w2 = nn.Linear(hidden_dim, dim, bias=False) def forward(self, x: torch.Tensor) -> torch.Tensor: x1, x3 = torch.chunk(self.w13(x), 2, dim=-1) return self.w2(F.silu(x1) * x3) class MultiHeadAttention(nn.Module): """Multi-head self-attention with query-key normalization (QK-Norm).""" def __init__(self, embedding_dim: int, n_heads: int) -> None: super().__init__() assert embedding_dim % n_heads == 0, "embedding_dim not divisible by n_heads" self.n_heads = n_heads self.head_dim = embedding_dim // n_heads self.packed_input_projection = nn.Linear(embedding_dim, embedding_dim * 3, bias=False) self.q_norm = nn.RMSNorm(self.head_dim) self.k_norm = nn.RMSNorm(self.head_dim) self.output_projection = nn.Linear(embedding_dim, embedding_dim, bias=False) def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor: B, L, _ = X.shape # (B, L, 3*E) -> (3, B, n_heads, L, head_dim); the packed layout is qkv-major. qkv = self.packed_input_projection(X) qkv = qkv.view(B, L, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4) Q, K, V = qkv.unbind(0) # QK-Norm: normalize Q and K per head before SDPA for softmax stability. Q = self.q_norm(Q) K = self.k_norm(K) key_mask = None query_pad_mask = None if attn_mask is not None: attn_mask = attn_mask.to(dtype=torch.bool) key_mask = attn_mask[:, None, None, :] # (B, 1, 1, L) query_pad_mask = (~attn_mask)[:, None, :, None] # (B, 1, L, 1) attn = F.scaled_dot_product_attention(Q, K, V, attn_mask=key_mask) if query_pad_mask is not None: attn = attn.masked_fill(query_pad_mask, 0.0) attn = attn.transpose(1, 2).reshape(B, L, self.n_heads * self.head_dim) return self.output_projection(attn) class TransformerBlock(nn.Module): """Pre-norm transformer block with reordered (double) RMSNorm outside the residual stream and SwiGLU feedforward.""" def __init__(self, embedding_dim: int, num_heads: int, mlp_ratio: float) -> None: super().__init__() self.multihead_attn = MultiHeadAttention(embedding_dim, num_heads) self.ffn = SwiGLU(embedding_dim, mlp_ratio) self.pre_attn_norm = nn.RMSNorm(embedding_dim) self.post_attn_norm = nn.RMSNorm(embedding_dim) self.pre_ffn_norm = nn.RMSNorm(embedding_dim) self.post_ffn_norm = nn.RMSNorm(embedding_dim) def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor: block = self.pre_attn_norm(X) block = self.multihead_attn(block, attn_mask=attn_mask) block = self.post_attn_norm(block) X = X + block ffn = self.pre_ffn_norm(X) ffn = self.ffn(ffn) ffn = self.post_ffn_norm(ffn) X = X + ffn return _zero_masked_tokens(X, attn_mask) class ParticleViT(nn.Module): """ParticleViT: a generic set-transformer over jet constituents.""" def __init__( self, num_features: int = 9, num_classes: int = 210, embed_dim: int = 512, depth: int = 5, num_heads: int = 8, mlp_ratio: float = 8 / 3, ) -> None: super().__init__() self.config = dict( num_features=num_features, num_classes=num_classes, embed_dim=embed_dim, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio, ) self.token_embed = ParticleInputEmbedding(num_features, embed_dim) self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) self.blocks = nn.ModuleList( [TransformerBlock(embed_dim, num_heads, mlp_ratio) for _ in range(depth)] ) self.norm = nn.RMSNorm(embed_dim) self.head = nn.Linear(embed_dim, num_classes) def forward(self, X: torch.Tensor, attn_mask: torch.Tensor | None = None) -> torch.Tensor: """Args: X: padded (B, L, 9) tensor of preprocessed particle features. attn_mask: bool (B, L), True for real particles (padding is False). Returns: logits of shape (B, num_classes). """ x = self.token_embed(X, attn_mask=attn_mask) x = _zero_masked_tokens(x, attn_mask) bs = x.shape[0] cls = self.cls_token.expand(bs, -1, -1) x = torch.cat([cls, x], dim=1) if attn_mask is not None: cls_mask = torch.ones((bs, 1), dtype=torch.bool, device=attn_mask.device) attn_mask = torch.cat([cls_mask, attn_mask], dim=1) for block in self.blocks: x = block(x, attn_mask=attn_mask) x = self.norm(x) x = _zero_masked_tokens(x, attn_mask) return self.head(x[:, 0]) @classmethod def from_pretrained(cls, repo_or_path: str, device: str = "cpu") -> "ParticleViT": """Load config.json + model.safetensors from a local directory or the Hugging Face Hub (repo id like 'jaluus/ParticleViT-S').""" from safetensors.torch import load_file path = Path(repo_or_path) if not path.exists(): from huggingface_hub import snapshot_download path = Path(snapshot_download(repo_or_path)) with (path / "config.json").open() as f: cfg = json.load(f) model = cls( num_features=cfg["num_features"], num_classes=cfg["num_classes"], embed_dim=cfg["embed_dim"], depth=cfg["depth"], num_heads=cfg["num_heads"], mlp_ratio=cfg["mlp_ratio"], ) state = load_file(str(path / "model.safetensors")) model.load_state_dict(state, strict=True) return model.to(device)