From ed3bd11aa26062faf64c678f1db1be4811c419bf Mon Sep 17 00:00:00 2001 From: Severin Magel <116261790+sevmag@users.noreply.github.com> Date: Sat, 15 Aug 2026 15:26:48 -0400 Subject: [PATCH 1/2] Add Neptune backbone + masked-point-modeling pretext (arXiv:2510.01733) Reproduce the paper's method as SPINE plugins so it can be compared to CURTAIN on a shared backbone, isolating the pretext from the architecture. - NeptuneBackbone: the per-hit encoder of the paper's reference (prometheus) implementation -- per-hit content MLP (charge only) + learned additive 4D positional MLP + a plain transformer + masked-mean readout, no FPS. Honors an optional batch["pos_mask"], swapping masked hits' coordinates for learned mask embeddings before the positional MLP. - MPMTask/MPMHead: masked point modeling -- mask a fraction of hits' space-time coordinates (spatial/temporal/spatiotemporal), reconstruct them per token with smooth-L1. Token content is position-free so the reconstruction is non-trivial. - configs/backbone/neptune.yaml, configs/task/mpm.yaml. Co-Authored-By: Claude Opus 4.8 --- configs/backbone/neptune.yaml | 10 ++ configs/task/mpm.yaml | 9 ++ src/spine/backbones/neptune.py | 127 +++++++++++++++++++++++ src/spine/pretrain/mpm/__init__.py | 5 + src/spine/pretrain/mpm/task.py | 157 +++++++++++++++++++++++++++++ 5 files changed, 308 insertions(+) create mode 100644 configs/backbone/neptune.yaml create mode 100644 configs/task/mpm.yaml create mode 100644 src/spine/backbones/neptune.py create mode 100644 src/spine/pretrain/mpm/__init__.py create mode 100644 src/spine/pretrain/mpm/task.py diff --git a/configs/backbone/neptune.yaml b/configs/backbone/neptune.yaml new file mode 100644 index 0000000..fe2f5f6 --- /dev/null +++ b/configs/backbone/neptune.yaml @@ -0,0 +1,10 @@ +# Per-hit Neptune backbone (paper / prometheus-branch architecture): per-hit +# content MLP + learned additive 4D positional MLP + a plain transformer, no FPS. +# Reduced size (~ reduced-DeepIce scale) for the pretext comparison; raise +# d_model/depth toward the paper's 768/12 for a full-scale run. +_target_: spine.backbones.neptune.NeptuneBackbone +d_model: 128 +depth: 3 +n_heads: 8 +dim_feedforward: 512 +dropout: 0.0 diff --git a/configs/task/mpm.yaml b/configs/task/mpm.yaml new file mode 100644 index 0000000..da00522 --- /dev/null +++ b/configs/task/mpm.yaml @@ -0,0 +1,9 @@ +# Masked point modeling (Yu, Kamp & Arguelles, arXiv:2510.01733) as a SPINE +# task. Pairs with backbone=neptune. Select with `task=mpm`. +_target_: spine.pretrain.mpm.MPMTask +max_pulses: 768 +center_time: true +mask_ratio: 0.15 +mode: temporal # spatial | temporal | spatiotemporal (paper: pretrain_task) +centroid_loss_weight: 1.0 +time_loss_weight: 1.0 diff --git a/src/spine/backbones/neptune.py b/src/spine/backbones/neptune.py new file mode 100644 index 0000000..f253fa4 --- /dev/null +++ b/src/spine/backbones/neptune.py @@ -0,0 +1,127 @@ +"""Per-hit Neptune backbone: the masked-point-modeling encoder of Yu, Kamp & +Arguelles, "Reducing Simulation Dependence in Neutrino Telescopes with Masked +Point Transformers" (arXiv:2510.01733). + +This mirrors that paper's reference implementation -- the ``prometheus`` branch +of github.com/felixyu7/neptune -- not the latest upstream. The current ``main`` +branch has since moved to a different design (parameter-free 4D RoPE and FPS +point-cloud patchification), so this file is pinned to the published +architecture to keep comparisons faithful to the paper's method rather than the +evolving one. + +Each hit is its own token: the paper's below-``max_tokens`` path, so there is no +FPS patchification here. Content and position are kept separate on purpose -- +the per-hit MLP embeds only non-positional features (charge), while a learned 4D +MLP embeds absolute space-time position and is ADDED to the token. A plain +Transformer encoder mixes tokens and a masked mean over real hits forms the +event embedding (the paper has no CLS token). Holding position out of the token +content is what makes a masked-position pretext non-trivial: with charge alone +visible, the encoder must infer a hit's location -- see ``spine.pretrain.mpm``. + +``encode`` honors an optional ``batch["pos_mask"]`` (bool ``[B, L]`` over pulses) +with ``batch["pos_mask_mode"]`` ("spatial" | "temporal" | "spatiotemporal"), +swapping masked hits' coordinates for learned mask embeddings before the +positional MLP; CURTAIN and fine-tuning leave it unset and run unmasked. +""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn + +from spine.backbones.base import Backbone, EncodedEvent + + +def _mlp(in_dim: int, hidden: tuple[int, ...], out_dim: int) -> nn.Sequential: + """GELU/LayerNorm MLP stack ending in a linear projection to ``out_dim``.""" + layers: list[nn.Module] = [] + last = in_dim + for h in hidden: + layers += [nn.Linear(last, h), nn.GELU(), nn.LayerNorm(h)] + last = h + layers.append(nn.Linear(last, out_dim)) + return nn.Sequential(*layers) + + +class NeptuneBackbone(Backbone): + """Per-hit Neptune encoder (paper / ``prometheus`` architecture). + + ``pos_cols`` index the space-time coordinates in the pulse feature vector as + ``(x, y, z, t)`` (spatial first, time last); ``content_cols`` index the + position-free features the token MLP sees (charge). + """ + + def __init__( + self, + d_model: int = 128, + depth: int = 3, + n_heads: int = 8, + dim_feedforward: int = 512, + dropout: float = 0.0, + pos_cols: tuple[int, ...] = (0, 1, 2, 3), + content_cols: tuple[int, ...] = (4,), + pos_hidden: tuple[int, ...] = (64, 256), + content_hidden: tuple[int, ...] = (256,), + ): + """Build the per-hit content/positional MLPs and the transformer. + + Args: + d_model: Token / embedding width. + depth: Number of transformer encoder layers. + n_heads: Attention heads. + dim_feedforward: Encoder feed-forward width. + dropout: Encoder dropout. + pos_cols: Feature columns for space-time position, spatial then time. + content_cols: Feature columns for the position-free token content. + pos_hidden: Hidden widths of the positional MLP. + content_hidden: Hidden widths of the content MLP. + """ + super().__init__() + self.out_dim = d_model + self.pos_cols = list(pos_cols) + self.content_cols = list(content_cols) + self.content_mlp = _mlp(len(content_cols), content_hidden, d_model) + self.pos_mlp = _mlp(len(pos_cols), pos_hidden, d_model) + layer = nn.TransformerEncoderLayer( + d_model, + n_heads, + dim_feedforward, + dropout, + activation="gelu", + batch_first=True, + norm_first=False, + ) + self.encoder = nn.TransformerEncoder(layer, depth) + self.ln = nn.LayerNorm(d_model) + # "Position unknown" tokens in raw coordinate space, swapped in for + # masked hits before the positional MLP (masked point modeling). + self.spatial_mask_emb = nn.Parameter(torch.randn(3) * 0.02) + self.time_mask_emb = nn.Parameter(torch.randn(1) * 0.02) + + def encode(self, batch: dict) -> EncodedEvent: + """Encode a collated batch of jagged pulses into per-hit tokens.""" + pulses = batch["pulses"] + x0 = pulses.to_padded_tensor(0.0) + lengths = pulses.offsets().diff() + length = x0.shape[1] + mask = torch.arange(length, device=x0.device)[None] < lengths[:, None] + pos = x0[..., self.pos_cols] + content = x0[..., self.content_cols] + + pos_mask = batch.get("pos_mask") + if pos_mask is not None: + mode = batch.get("pos_mask_mode", "spatiotemporal") + m = (pos_mask & mask).unsqueeze(-1) + xyz, t = pos[..., 0:3], pos[..., 3:4] + if mode in ("spatial", "spatiotemporal"): + xyz = torch.where(m, self.spatial_mask_emb, xyz) + if mode in ("temporal", "spatiotemporal"): + t = torch.where(m, self.time_mask_emb, t) + pos = torch.cat([xyz, t], dim=-1) + + tok = self.content_mlp(content) + self.pos_mlp(pos) + tok = self.encoder(tok, src_key_padding_mask=~mask) + tok = self.ln(tok) + denom = mask.sum(1, keepdim=True).clamp(min=1) + cls = (tok * mask.unsqueeze(-1)).sum(1) / denom + return EncodedEvent(tokens=tok, token_mask=mask, cls=cls) diff --git a/src/spine/pretrain/mpm/__init__.py b/src/spine/pretrain/mpm/__init__.py new file mode 100644 index 0000000..96f36fe --- /dev/null +++ b/src/spine/pretrain/mpm/__init__.py @@ -0,0 +1,5 @@ +"""Masked point modeling pretext (arXiv:2510.01733).""" + +from spine.pretrain.mpm.task import MPMHead, MPMTask + +__all__ = ["MPMHead", "MPMTask"] diff --git a/src/spine/pretrain/mpm/task.py b/src/spine/pretrain/mpm/task.py new file mode 100644 index 0000000..3186d2b --- /dev/null +++ b/src/spine/pretrain/mpm/task.py @@ -0,0 +1,157 @@ +"""Masked point modeling (MPM): the self-supervised pretext of Yu, Kamp & +Arguelles, "Reducing Simulation Dependence in Neutrino Telescopes with Masked +Point Transformers" (arXiv:2510.01733), reproduced as a SPINE task. + +A random fraction of hits have their space-time coordinates hidden -- replaced, +inside ``NeptuneBackbone``, by learned mask embeddings -- while their charge +stays visible; the head reconstructs the hidden coordinates from the encoder's +per-hit outputs with a smooth-L1 loss. ``mode`` selects which coordinates are +masked and scored: "spatial" (xyz), "temporal" (t) or "spatiotemporal" (both), +matching the paper's ``pretrain_task``. + +Pairs with ``spine.backbones.neptune.NeptuneBackbone``, whose per-hit token +content is position-free, so recovering a masked hit's coordinate is non-trivial +(it must be inferred from the hit's charge and the rest of the event). The mask +is per-hit, faithful to the paper's below-``max_tokens`` regime. Positions are +reconstructed from the pulses themselves, so no geometry asset is needed. + +Coordinate columns follow the ``NeptuneBackbone`` default position layout +``(x, y, z, t)`` at columns 0-3 of the scaled pulse feature vector. +""" + +from __future__ import annotations + +from typing import Any + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +from spine.pretrain.base import PretrainTask, Sample + +_XYZ = slice(0, 3) +_T = slice(3, 4) + + +class MPMHead(nn.Module): + """Per-hit coordinate reconstruction heads (centroid xyz + time).""" + + def __init__(self, dim: int): + """Build the xyz and time linear predictors over encoder tokens.""" + super().__init__() + self.centroid = nn.Linear(dim, 3) + self.time = nn.Linear(dim, 1) + + def forward(self, _query_pos: Tensor, enc: Any) -> tuple[Tensor, Tensor]: + """Predict per-hit xyz and time from the encoded tokens. + + The query-position argument the engine passes is unused: MPM + reconstructs one prediction per encoder token, not per query. + """ + return self.centroid(enc.tokens), self.time(enc.tokens) + + +class MPMTask(PretrainTask): + """Masked point modeling over per-hit tokens (paper-faithful).""" + + objectives: list = [] + + def __init__( + self, + geo: dict, + scaler: Any, + max_pulses: int = 768, + center_time: bool = True, + mask_ratio: float = 0.15, + mode: str = "temporal", + centroid_loss_weight: float = 1.0, + time_loss_weight: float = 1.0, + ): + """Assemble the MPM task. + + Args: + geo: Geometry asset (unused; positions come from the pulses). + scaler: Detector feature scaling, applied at collate time. + max_pulses: Cap on hits fed to the encoder per event. + center_time: Reference times to the charge-weighted mean. + mask_ratio: Fraction of hits masked per event (1.0 masks all, + the paper's directional setting). + mode: Which coordinates to mask and score -- "spatial", + "temporal" or "spatiotemporal". + centroid_loss_weight: Weight on the xyz reconstruction term. + time_loss_weight: Weight on the time reconstruction term. + + Raises: + ValueError: If mode is not one of the three tasks. + """ + if mode not in ("spatial", "temporal", "spatiotemporal"): + raise ValueError( + f"mode must be spatial|temporal|spatiotemporal, got {mode!r}" + ) + self.geo = geo + self.scaler = scaler + self.max_pulses = max_pulses + self.center_time = center_time + self.mask_ratio = mask_ratio + self.mode = mode + self.centroid_loss_weight = centroid_loss_weight + self.time_loss_weight = time_loss_weight + + def make_sample(self, event: dict[str, np.ndarray], rng: np.random.Generator) -> Sample: + """Cap and time-center one event's hits (no split -- MPM masks in-place).""" + p = event["pulses"] + lay = self.scaler.layout + if len(p) > self.max_pulses: + p = p[rng.choice(len(p), self.max_pulses, replace=False)] + p = p.astype(np.float32).copy() + if self.center_time: + w = np.clip(p[:, lay.charge], 0.0, None) + 1e-6 + p[:, lay.t] -= float((w * p[:, lay.t]).sum() / w.sum()) + return dict(pulses=p) + + def collate(self, samples: list[Sample]) -> dict: + """Pack pulses and draw a per-event random position mask.""" + + def jag(tensors): + return torch.nested.nested_tensor(tensors, layout=torch.jagged) + + scaled = [self.scaler.scale_pulses(torch.from_numpy(s["pulses"])) for s in samples] + lengths = [t.shape[0] for t in scaled] + lmax = max(lengths) + pos_mask = torch.zeros(len(samples), lmax, dtype=torch.bool) + for b, n in enumerate(lengths): + k = max(1, int(round(self.mask_ratio * n))) + pos_mask[b, torch.randperm(n)[:k]] = True + # qpos/label satisfy the engine collate contract; qpos is unused by + # MPMHead and label only sizes the logged batch (total hits). + return dict( + pulses=jag(scaled), + qpos=jag([t[:, _XYZ].clone() for t in scaled]), + label=jag([torch.ones(n) for n in lengths]), + pos_mask=pos_mask, + pos_mask_mode=self.mode, + ) + + def build_head(self, dim: int) -> nn.Module: + """Construct the per-hit reconstruction head.""" + return MPMHead(dim) + + def loss(self, output: Any, batch: dict) -> tuple[Tensor, dict[str, float]]: + """Smooth-L1 reconstruction of masked hits' coordinates.""" + pred_xyz, pred_t = output + x0 = batch["pulses"].to_padded_tensor(0.0) + lengths = batch["pulses"].offsets().diff() + valid = torch.arange(x0.shape[1], device=x0.device)[None] < lengths[:, None] + m = batch["pos_mask"].to(valid.device) & valid + total = pred_xyz.new_zeros(()) + metrics: dict[str, float] = {} + if self.mode in ("spatial", "spatiotemporal") and m.any(): + term = F.smooth_l1_loss(pred_xyz[m], x0[..., _XYZ][m]) + total = total + self.centroid_loss_weight * term + metrics["loss_mpm_centroid"] = float(term.detach()) + if self.mode in ("temporal", "spatiotemporal") and m.any(): + term = F.smooth_l1_loss(pred_t[m], x0[..., _T][m]) + total = total + self.time_loss_weight * term + metrics["loss_mpm_time"] = float(term.detach()) + return total, metrics From 01a037ef8a560638287100f733f3fd2a17e54720 Mon Sep 17 00:00:00 2001 From: Severin Magel <116261790+sevmag@users.noreply.github.com> Date: Sat, 15 Aug 2026 15:26:49 -0400 Subject: [PATCH 2/2] Make EncodedEvent.cls optional Per-token pretexts (masked point modeling) read `tokens` and never need the pooled event embedding, so cls now defaults to None. The one consumer that requires it -- CURTAIN's QueryCrossAttnEncoder -- guards for None with a clear error. Existing backbones still fill cls, so behavior is unchanged. Co-Authored-By: Claude Opus 4.8 --- src/spine/backbones/base.py | 9 +++++++-- src/spine/pretrain/curtain/head.py | 5 +++++ 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/src/spine/backbones/base.py b/src/spine/backbones/base.py index 5a78b03..add0ad5 100644 --- a/src/spine/backbones/base.py +++ b/src/spine/backbones/base.py @@ -13,11 +13,16 @@ @dataclass class EncodedEvent: - """What every backbone returns.""" + """What every backbone returns. + + ``cls`` is optional: a backbone with no event-level readout leaves it + ``None``, and consumers that need an event embedding (the query encoder) + must guard for it. Per-token pretexts read ``tokens`` and ignore ``cls``. + """ tokens: Tensor # [B, L, D] per-pulse token embeddings token_mask: Tensor # [B, L] bool, True = real pulse (not padding) - cls: Tensor # [B, D] pooled event embedding + cls: Tensor | None = None # [B, D] pooled event embedding, or None class Backbone(nn.Module): diff --git a/src/spine/pretrain/curtain/head.py b/src/spine/pretrain/curtain/head.py index ac8f2ad..2265ee9 100644 --- a/src/spine/pretrain/curtain/head.py +++ b/src/spine/pretrain/curtain/head.py @@ -85,6 +85,11 @@ def forward(self, query_pos: Tensor, enc: EncodedEvent) -> Tensor: Returns: [B, Q, D] per-query embeddings. """ + if enc.cls is None: + raise ValueError( + "QueryCrossAttnEncoder requires an event-level `cls` embedding, " + "but the backbone returned cls=None" + ) kv = torch.cat([enc.cls.unsqueeze(1), enc.tokens], dim=1) # [B,1+L,D] ones = torch.ones( enc.token_mask.shape[0], 1, dtype=torch.bool, device=enc.token_mask.device