Add Neptune backbone + masked point modeling (arXiv:2510.01733)

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

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

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):

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)

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

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"]

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