Add Neptune point-transformer backbone by felixyu7 · Pull Request #923 · graphnet-team/graphnet

felixyu7 · 2026-09-10T02:15:40Z

Summary

Adds Neptune ("a N Efficient Point Transformer for Ultrarelativistic Neutrino Events", arXiv:2510.01733, reference implementation) as a backbone in graphnet.models.transformer.

Neptune treats an event as an irregular 4D (x, y, z, t) point cloud. It tokenizes the event by farthest point sampling (FPS) into at most num_patches centroids, pools the pulses of each token, runs a rotary-attention transformer over the tokens, and pools to one event vector. It has no edges and no pulse cap, so it drops into a StandardModel with EdgelessGraph over NodesAsPulses.

Builds on #913 (generalised FourierEncoder); this PR will be rebased once #913 merges.

What is added

models/transformer/neptune.py: Neptune(GNN) and PointTransformerEncoder. Columns are passed by name (coordinate_columns, time_column, charge_column, feature_columns) and resolved against input_feature_names. The model takes the Detector's standardized inputs as-is; geometric defaults are in IceCube86 units, and the charge column must be physical (Detector(replace_with_identity=["charge"])). Centroid positions are encoded with Generalise FourierEncoder with a column/frequency schema #913's FourierEncoder.
components/tokenizers.py: FPSTokenizer. Small events keep one token per pulse; larger ones are reduced to centroids and pooled by Voronoi assignment or kNN gather.
components/fps.py: farthest_point_sampling, farthest_point_sampling_with_assign, farthest_point_sampling_with_knn, nearest_assign in pure PyTorch. If the optional torch-fps package is installed, these delegate to it (Triton on CUDA, C++ on CPU).
components/neptune_layers.py: SwiGLU, RoPE4D, AttentionPool, NeptuneTransformerEncoderLayer, NeptuneTransformerEncoder.
models/utils.py: pack / unpack / build_block_mask for variable-length packing with flex_attention at low batch occupancy.
utilities/imports.py: has_torch_fps_package().
DropPath.forward gains optional doc_id / num_docs so stochastic depth stays per-event on the packed path; default behaviour is unchanged.
examples/04_training/09_train_neptune_model.py: direction reconstruction on the Prometheus example data.
tests/models/test_neptune.py: tests for shapes, empty events and batches, gradients, feature-name resolution, time-shift invariance, low precision, config round trip, StandardModel, and FPS/kNN/assign against brute force and against torch-fps.

Dependencies

No new required dependencies. torch-fps is optional; without it Neptune uses the pure-PyTorch FPS, which costs about +36% per training step and +89% at inference on real IceCube events (RTX 5090, batch 64).

Validation

Trained on ~10.9M IceCube NuMu-CC events (direction reconstruction, VonMisesFisher3DLoss) alongside DynEdge under an identical schedule; Neptune reached a better best-validation loss. Tests and pre-commit hooks pass locally on CPU.

🤖 Generated with Claude Code

Adds Neptune (arXiv:2510.01733) as a backbone in graphnet.models.transformer, together with the components it needs: - FPSTokenizer: farthest-point-sampling tokenizer over the 4D (x, y, z, t) point cloud, with Voronoi or kNN pooling and no cap on pulses per event. - Farthest point sampling with fused assign/kNN variants: a pure-PyTorch backend and an optional Triton backend used automatically on CUDA. - FourierPositionEncoder, RMSNorm, SwiGLU, RoPE4D, AttentionPool, and the Neptune transformer encoder layers. - pack/unpack/build_block_mask helpers for variable-length packing with flex_attention, used at low batch occupancy; padded SDPA otherwise. - DropPath gains optional doc_id/num_docs arguments so stochastic depth stays per-event on the packed path. Default behaviour is unchanged. - has_triton_package() import helper. Also adds a training example (09_train_neptune_model.py) and a test module covering the model, tokenizer, and FPS kernels against brute force. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

`FourierEncoder` assumed its input was ordered (x, y, z, time, charge, auxiliary) and applied multipliers fixed to the Kaggle dataset's normalisation, so it index-errors or silently mis-encodes on data laid out differently (issue graphnet-team#753). The multipliers are the more subtle half: they set which separations the sinusoidal ladder can resolve, so carrying them to a detector with a different normalisation changes the resolved physical range without any error. Rename the existing class to `FourierEncoderEPJC`, documented as the EPJ-C Kaggle implementation with those assumptions, and add a `FourierEncoder` taking `schema: Dict[int, float]` -- column index to multiplier. Nothing is assumed about column order or meaning, and the boolean embedding and MLP projection are left to the model, both being problem-specific. Follows the design proposed in the review of graphnet-team#759. The refactor is representation-preserving: the sinusoidal encoding is elementwise, so one embedding per column lays channels out exactly as the strided xyz block did. `tests/models/test_embedding.py` asserts that, and that a schema of the EPJC columns reproduces its pre-projection embedding bit-for-bit. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01F75gXoxTcs4nVRxGp91ENf

`seq_length` is a per-sequence quantity, while `x` is per-step, and the name is also the constructor's embedding-width argument -- so state the `[B]` shape and that each step receives a copy of its sequence's embedded length. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01F75gXoxTcs4nVRxGp91ENf

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01F75gXoxTcs4nVRxGp91ENf

A column's resolvable band has two ends: the multiplier fixes the finest wavelength, `n_freq` fixes how far the ladder reaches beyond it. Only the first was per column, so features wanting different spans could not get them -- and they do differ. On a 1 km neutrino detector the sensor spacing and extent ask for a span of ~180x, event timing for ~3600x and the charge quantisation for ~25x, against the single hardcoded 10000 (a ~9085x ladder) every feature shares today. `schema` now takes either a multiplier or a `(multiplier, n_freq)` pair per column, with a bare multiplier falling back to the encoder's `n_freq`, so existing schemas are unchanged. One `SinusoidalPosEmb` is built per distinct span rather than per column, which keeps the common single-span case at one module and one learnable scale under `scaled`. `SinusoidalPosEmb` also precomputes its ladder into a non-persistent buffer instead of rebuilding it every forward pass, so the extra spans cost nothing and checkpoints are untouched. Its `n_freq` docstring said "Number of frequencies", which is `dim / 2`; it is the span of the ladder, and that mislabel is what hid the missing knob. Verified bit-identical against the previous implementation for the ladder itself, `FourierEncoderEPJC`, `SpacetimeEncoder`, and `FourierEncoder` with bare multipliers; a per-column span changes only its own block. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01F75gXoxTcs4nVRxGp91ENf

Resolve the components/layers.py conflict by adopting upstream's removal of the module and moving the Neptune layer classes (RMSNorm, SwiGLU, RoPE4D, AttentionPool, NeptuneTransformerEncoderLayer, NeptuneTransformerEncoder) to a new components/neptune_layers.py, matching the per-family module layout introduced upstream. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

sevmag · 2026-09-26T20:21:30Z

+                for cross-attention pooling with a learned query.
+            layerscale_init: Initial value of the LayerScale parameters.
+            attn_impl: `"auto"`, `"padded"`, or `"packed"`. Purely a
+                throughput choice; all three are mathematically equivalent.


Make sure to explain what each of the settings means. It should be as concise as possible but still descriptive enough so that people actually understand what you meant by it. The way it is done now is that it is not explained here, but then later on there is a very long comment explaining it, which, imo, is suboptimal

The modes are now documented in the argument's docstring; the long comment is gone.

sevmag

@felixyu7 we are honored that this model gets to live in graphnet. The implementation is mostly carefully crafted by an eager helper (@claude). General stlye comments:

the code is veeeeery verbose, reviewing so many lines of code in one PR is very exhausting. For these kinds of repos like graphnet manual code review is still necessary to some extend. Not for correctness but mainly to make sure it is maintainable
comments in the code should ideally not span multiple lines (more than 2-3)
docstrings should be held to the minimum that is need to understand the code from the context but should also explain the most important parts

More concrete desgin comments:

the data representation stuff like standardization or node sampling should probably be done in the code that lies in graphnet/models/data_representation and not in the model forward
I do not think we should add the triton kernel to. the main branch for now. See my comment about this at the code for alternatives
we should aim for a general version of the FourierEncoder (see #913 ) before we start introducing a new class

Apart from that the PR looks very good! Make sure to adjust the comments or state your point against them and try to deverbose as much as possible!

sevmag · 2026-09-26T20:21:41Z

+
+        self.token_dim = token_dim
+        self.pool_type = pool_type
+        # "auto": packed flex on CUDA, padded SDPA on CPU (where the packed


see comment above

The modes are now documented in the argument's docstring; the long comment is gone.

sevmag · 2026-09-26T20:23:29Z

+
+    def __init__(
+        self,
+        token_dim: int = 768,


why do we not call this d_model as it is called in NeptuneTransformerEncoderLayer, that way it is less confusing

Renamed to d_model / depth.

sevmag · 2026-09-26T20:35:34Z

+            `[B, token_dim]` pooled event representation.
+        """
+        # No cast to `tokens.dtype`: the position encoder computes in float32
+        # anyway, and a bfloat16 round-trip of km-scale coordinates costs up


this comment is a little bit overly specific. I think it confuses more than it helps. While I do agree that we should mention that it is on purpose the we do not cast to a specific numerical type, I think we should reframe this.

E.g. : "# Positions are keept in float32 in purpose to avoid loss of information due to rounding"

or something shorter and less speficic (The rad number is maybe a bit too much)

Now a one-line comment.

sevmag · 2026-09-26T20:42:00Z

+    embeddings. This class embeds *position only*, and is applied to token
+    centroids rather than to individual pulses.
+
+    Time is deliberately excluded: callers are expected to centre time per


I get what this argument is saying, but feeding the time feature which is centered on the charge-weighted mean time is not meaningless at all it gives information on whether a pulse is early or late in the event. I do not think this argument actually holds.

Agreed, that argument was wrong. It's now framed as a design choice: time is encoded relatively via RoPE, and adding column 3 to position_encoding_schema also encodes it absolutely.

sevmag · 2026-09-26T20:45:52Z

+        """
+        # Trig in float32 for accuracy; the MLP follows the ambient autocast
+        # dtype.
+        xyz = coords[..., : self.in_dim].float()


this assumes that the coordinates are the first self.in_dim feature columns, which is not ideal. Let's pass a list of indices that are the indices of the position features in the feature matrix, so that this is more general.

The class is removed in favour of #913's FourierEncoder, whose schema names the columns explicitly and builds the frequencies once.

sevmag · 2026-09-26T22:55:59Z

@@ -0,0 +1,826 @@
+"""Learnable tokenizers turning point clouds into transformer tokens."""


This whole file is really long. Let's de-verbose it; otherwise, reviewing 900 lines in one file is a bit tideous

826 → 519 lines; the duplicated summary-statistics code is now shared.

sevmag · 2026-09-26T22:58:45Z

+from graphnet.models.utils import flex_attention
+
+
+class RMSNorm(LightningModule):


can we not use the torch native RMSNorm https://docs.pytorch.org/docs/2.14/generated/torch.nn.RMSNorm.html ?

Now uses torch.nn.RMSNorm.

sevmag · 2026-09-26T23:03:35Z

I don't think we should introduce triton kernels in graphnet itself as I'm not sure whether people will maintain it and it takes a lot of expertise to do so. I think we should stick with the slower, normal Python code, and if people want, they can implement the Triton kernel in their own branch (you could link a piece of code that does this). The other option is that you write an outside module that we can keep as an optional dependency as we are doing with jammy_flows. then peopel can use it if they want without us having the implementation in the graphnet code

Moved to torch-fps (see summary).

sevmag · 2026-09-26T23:04:18Z

        return False


+def has_triton_package() -> bool:


see comment for triton kernel of fps

Moved to torch-fps (see summary).

sevmag · 2026-09-26T23:04:47Z

    return rw_landing
+
+
+# ------------------------------------------------------------------------


way to verbose comment

…-backbone Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

felixyu7 · 2026-09-30T19:50:42Z

Thanks for the thorough review! I agree with all of it, and the branch now addresses every comment:

Verbosity: the PR is down from ~4.4k to ~2.4k added lines, not counting Generalise FourierEncoder with a column/frequency schema #913's changes (e.g. neptune.py 696 → 392, tokenizers.py 826 → 519). Comments are at most 2–3 lines, and docstrings are kept to the minimum.
Standardization: Neptune now takes the Detector's standardized inputs as-is and needs no special handling.
- Position and time only need a scale, so the geometric defaults are in IceCube86 input units.
- Time centering was unnecessary: time only enters as differences (RoPE, offsets, FPS distances). A test checks that the output is invariant to a time shift.
- Charge must be physical, via the existing Detector(replace_with_identity=["charge"]).
- If charge-weighted centering in the Detector would be useful more broadly, I'm happy to do that in a separate PR.
Node sampling: Neptune doesn't subsample pulses. Farthest point sampling (FPS) is the tokenizer, so it stays in the model.
Triton: moved to the optional torch-fps package (Triton on CUDA, C++ on CPU), guarded like jammy_flows. graphnet only ships the pure-PyTorch path. Here's the cost of not installing it, on 5,000 real IceCube events (RTX 5090, uncapped as Neptune is trained, bs64):

pure PyTorch torch-fps slowdown

training step 20.1 ms 14.8 ms +36%

inference 11.3 ms 6.0 ms +89%

The gap shrinks with batch and model size (+16% for the full-size model). On CPU, FPS is a small part of the forward pass, so torch-fps saves only 3–14% there.
Fourier encoder: removed. The branch builds on Generalise FourierEncoder with a column/frequency schema #913's FourierEncoder, and I'll rebase once Generalise FourierEncoder with a column/frequency schema #913 merges.

@sevmag ready for another look whenever you have time.

sevmag · 2026-10-02T20:47:33Z

Hey @felixyu7, before I rereview this make sure to rebase it on the new main that now as the generalized fourierencoder code from #913.

sevmag · 2026-10-02T20:48:12Z

I also added the numerics guard in the Sinusodial embedding that you were asking about! LMK if you have other things you want to add

# Conflicts: # src/graphnet/models/components/embedding.py # src/graphnet/models/gnn/icemix.py

felixyu7 · 2026-10-03T16:28:47Z

@sevmag should be done, let me know if there are any issues!

sevmag · 2026-10-04T19:27:22Z

Some of these components are general concepts like SwiGLU, 4DRoPE, and attention pool. Lets move them in more general files (not having Neptune in the title) as other people might want to reuse these for their models and then neptune_layers.py is a poor name for the file.

I'm actually also not sure whether we need so many very specific classes like NeptuneTransformerEncoder and PointTransformerEncoder in neptune.py. I think we can make this easier and have the Neptune class in neptune.py do a bit more parsing and work but there being able to remove these medium-level classes. This should also save complexity. I'm also not the biggest fan of the usage of the factory functions here since they complicate the interfaces. Lets stick to args

Done! SwiGLU, AttentionPool and the encoder block (now RotaryTransformerBlock) are in attention_blocks.py, RoPE4D is in embedding.py, and neptune_layers.py is gone. Both encoder wrappers and the layer factory are removed too, so Neptune builds its layers directly from args.

sevmag · 2026-10-04T19:28:42Z

see my comment in neptune_layers. The usage of PointTransformerEncoder might be unnecessary

Removed, see above.

sevmag · 2026-10-04T19:29:47Z

+        if compile_encoder:
+            self.compile_layers()
+
+    def compile_layers(


I think we can get rid of this function and just put the code in the init call saves a lot of lines and makes everything more readable

Inlined into __init__.

sevmag · 2026-10-04T19:31:37Z

+        """
+        x, batch, batch_size = data.x, data.batch, data.num_graphs
+        if self._charge_index is None:
+            charge = x.new_ones(x.shape[0])


what is the case for? Is this maybe a scenario where there are no charge variables so you asign a charge of 1? if yes this should not be done in the model but in the data_representation

Yep, that's the no-charge case. I'd keep it in the model though: the tokenizer uses charge for weighting and token summaries, and for Prometheus every pulse is a single photon hit, so unit charge is the physically correct value rather than a placeholder. Moving it to the data representation would mean a new NodeDefinition just to append a column of ones. Happy to make charge_column required instead if you'd prefer.

sevmag · 2026-10-04T19:35:20Z

+
+# Packing of variable-length batches for block-diagonal `flex_attention`.
+# Compiled on first use, ~50x faster than eager `create_block_mask`.
+_COMPILED_CREATE_BLOCK_MASK: Optional[Callable] = None


wait do you use this global to set compiling??

Nope, it only cached the compiled create_block_mask so it compiles once. It's a plain module-level torch.compile(...) now (lazy, compiles on first call), so the global is gone.

sevmag · 2026-10-04T19:35:31Z

+        doc_id: `[N]` event index of every packed token.
+        n_tokens: Packed sequence length `N`.
+    """
+    global _COMPILED_CREATE_BLOCK_MASK


see comment above

Same, gone.

sevmag · 2026-10-04T19:36:10Z

+    ) -> Tensor:
+        return doc_id[q_idx] == doc_id[kv_idx]
+
+    return _COMPILED_CREATE_BLOCK_MASK(


when would you not use torch.compile?

Fair point. Compile is now on by default (compile_encoder=True), applied per layer, which also cut warmup from ~42s to ~9s at the same step time. Turn it off only for CPU-only or debugging runs (the tests do, to keep CI fast). I also dropped attn_impl: packed flex attention now runs automatically when compiled on GPU.

sevmag · 2026-10-04T19:37:29Z

+
+Ported from the Neptune reference implementation
+(https://github.com/felixyu7/neptune, MIT licence). Calls go to the optional
+`torch_fps` package (Triton kernels on CUDA, C++ on CPU) when it is installed;


make sure to link this repo so people can find it

also create a release so that it is backwards compatible no matter what you do with it after

Tagged and released v0.6.0 (same code as PyPI 0.6.0): https://github.com/felixyu7/torch-fps/releases/tag/v0.6.0. The install hint now pins torch-fps==0.6.*.

sevmag · 2026-10-04T19:40:01Z

+    start_idx: Optional[Tensor],
+    random_start: bool,
+    generator: Optional[torch.Generator],
+    precision: Optional[torch.dtype],


Is there a good reason for why precision should be optional? Silent skips of dtypes are annoying to debug

No good reason. Neptune never set it, and the fallback silently upcast half precision while torch_fps didn't. Removed the arg, so both backends always compute in float32.

Move SwiGLU, AttentionPool and the rotary block (now RotaryTransformerBlock) into attention_blocks.py and RoPE4D into embedding.py; delete neptune_layers.py. Drop the PointTransformerEncoder/NeptuneTransformerEncoder wrappers and the layer factory so Neptune builds its layers from args. Compile per layer by default (compile_encoder=True), dropping attn_impl and compile_layers. Replace the global block-mask cache with a module-level torch.compile. Remove the FPS precision argument; both backends compute in float32. Link and pin torch-fps 0.6. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

felixyu7 · 2026-10-05T02:12:24Z

@sevmag addressed everything in a19935c, ready for another look!

Compiling each layer separately keyed one dynamo cache entry per layer on their shared forward (each layer has its own drop-path rate). With the eval-mode entries that overran the recompile limit, so the deeper layers silently ran eager, where packed flex attention materialises the full N x N mask: training ran 2.6x slower on an H200 (12.3 vs 31.5 batch/s). Compile one module-level function over the whole stack instead, and test that a stack deeper than the limit compiles in train and eval. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

felixyu7 and others added 6 commits September 9, 2026 22:15

Drop the FourierEncoder test module

6869594

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01F75gXoxTcs4nVRxGp91ENf

sevmag reviewed Sep 26, 2026

View reviewed changes

felixyu7 and others added 2 commits September 28, 2026 19:57

Merge graphnet-team#913 (generalised FourierEncoder) into add-neptune…

d9a1cd3

…-backbone Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

address reviewer comments round1

59517a1

felixyu7 mentioned this pull request Sep 30, 2026

Generalise FourierEncoder with a column/frequency schema #913

Merged

Merge remote-tracking branch 'origin/main' into add-neptune-backbone

d62efd5

# Conflicts: # src/graphnet/models/components/embedding.py # src/graphnet/models/gnn/icemix.py

sevmag reviewed Oct 4, 2026

View reviewed changes

		@@ -0,0 +1,826 @@
		"""Learnable tokenizers turning point clouds into transformer tokens."""

		from graphnet.models.utils import flex_attention


		class RMSNorm(LightningModule):

		return rw_landing


		# ------------------------------------------------------------------------

Conversation

felixyu7 commented Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What is added

Dependencies

Validation

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sevmag left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

felixyu7 commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

sevmag commented Oct 2, 2026

Uh oh!

sevmag commented Oct 2, 2026

Uh oh!

felixyu7 commented Oct 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

felixyu7 commented Sep 10, 2026 •

edited

Loading

felixyu7 commented Sep 30, 2026 •

edited

Loading