Skip to content

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

Closed
sevmag wants to merge 2 commits into
mainfrom
neptune-mpm-backbone
Closed

Add Neptune backbone + masked point modeling (arXiv:2510.01733)#15
sevmag wants to merge 2 commits into
mainfrom
neptune-mpm-backbone

Conversation

@sevmag

@sevmag sevmag commented Aug 15, 2026

Copy link
Copy Markdown
Owner

What

Adds the method from "Reducing Simulation Dependence in Neutrino Telescopes with Masked Point Transformers" (Yu, Kamp & Argüelles, arXiv:2510.01733) as SPINE plugins, so it can be compared to CURTAIN on a shared backbone — isolating the pretext from the architecture.

  • NeptuneBackbone (src/spine/backbones/neptune.py) — the paper's per-hit encoder: per-hit content MLP (charge only) + learned additive 4D positional MLP + a plain transformer + masked-mean readout, no FPS. It mirrors the paper's reference implementation (the prometheus branch), not upstream main (which has since moved to parameter-free 4D RoPE + FPS patchification) — pinned to the published architecture so comparisons stay faithful to the paper's method. encode honors an optional batch["pos_mask"], swapping masked hits' coordinates for learned mask embeddings before the positional MLP.
  • MPMTask / MPMHead (src/spine/pretrain/mpm/) — masked point modeling, 1-to-1 with the paper: 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. Rides the existing collate/engine contract — no engine or CURTAIN changes.
  • EncodedEvent.cls is now optional (Tensor | None = None). Per-token pretexts (MPM) read tokens and ignore cls; the only consumer that needs it (CURTAIN's QueryCrossAttnEncoder) now guards for None with a clear error. Existing backbones still fill cls, so behavior is unchanged.
  • Configs: configs/backbone/neptune.yaml, configs/task/mpm.yaml.

Why

To answer, cleanly, whether CURTAIN's causal light-front forecasting pretext beats generic masked reconstruction — by running both on the same backbone with the same fine-tune, so the pretext is the only variable. (The two backbones' inductive biases interact with the pretext: DeepIce's learned Minkowski attention bias fights time-masking, whereas Neptune's additive positional MLP + mask embeddings tolerate it — which is exactly why a shared, neutral backbone is the fair place to compare pretexts.)

Validated

Unit smokes (synthetic batches, run on a GPU node):

  • Backbone: correct shapes, masking measurably changes the encoding across all three modes, grads flow to the mask embeddings.
  • MPM: all three modes compute, backward reaches the reconstruction heads and the mask embeddings, and the make_sample / collate data path runs end-to-end.
  • EncodedEvent constructs without cls.

Scope / follow-ups (not in this PR)

  • Fine-tuning a Neptune-pretrained encoder (spine_ft.py --backbone neptune) — the fine-tune harness currently hard-codes a graphnet DeepIce.
  • A full train_curtain.py backbone=neptune task=mpm pipeline run.

Opening as draft because the end-to-end train → fine-tune → compare path isn't exercised yet; the plugins themselves are self-contained and validated.

🤖 Generated with Claude Code

sevmag and others added 2 commits August 15, 2026 15:26
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
@sevmag

sevmag commented Aug 15, 2026

Copy link
Copy Markdown
Owner Author

Re-scoped: #15 bundled the Neptune backbone + MPM pretext with the cls-optional change. Closing in favor of #16 (cls-optional only). The Neptune/MPM work stays on the neptune-mpm-backbone branch for a separate PR if wanted.

@sevmag sevmag closed this Aug 15, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant