Add Neptune backbone + masked point modeling (arXiv:2510.01733) - #15
Closed
sevmag wants to merge 2 commits into
Closed
Add Neptune backbone + masked point modeling (arXiv:2510.01733)#15sevmag wants to merge 2 commits into
sevmag wants to merge 2 commits into
Conversation
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>
Owner
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 (theprometheusbranch), not upstreammain(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.encodehonors an optionalbatch["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.clsis now optional (Tensor | None = None). Per-token pretexts (MPM) readtokensand ignorecls; the only consumer that needs it (CURTAIN'sQueryCrossAttnEncoder) now guards forNonewith a clear error. Existing backbones still fillcls, so behavior is unchanged.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):
make_sample/collatedata path runs end-to-end.EncodedEventconstructs withoutcls.Scope / follow-ups (not in this PR)
spine_ft.py --backbone neptune) — the fine-tune harness currently hard-codes a graphnet DeepIce.train_curtain.py backbone=neptune task=mpmpipeline 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