Add LGATr-slim and LLoCa - #31
Open
hqucms wants to merge 7 commits into
Open
Conversation
Port the L-GATr-slim network (lgatr v2.0.0 nets/slim.py + layers/slim_layers.py) and the tagging-guide LGATrSlimWrapper dense path into a self-contained weaver/nn/model/LGATrSlim.py: - Slim layers (SlimLinear/SlimRMSNorm/SlimGLU/SlimSelfAttention/SlimMLP/SlimBlock) and the LGATrSlim core net, with fp32 precision islands under autocast (minimum_autocast_precision / naive_amp), gradient checkpointing, optional torch.compile, and dead-tail parameter freezing. Bit-exact parity with the reference implementation verified via state_dict transfer. - LGATrSlimTagger: weaver-facing wrapper with four-momentum embedding (E, px, py, pz), symmetry-breaking spurions, global class token, native SDPA key masking, optional masked-mean aggregation, and SequenceTrimmer reuse. Defaults follow the tagging-guide tag_slim config (2.03M params). - movedim replaced by a normalized-permute helper: negative dims produce invalid Transpose nodes in the TorchScript ONNX exporter. - Example network config (test/networks/example_LGATrSlim.py) and unit tests covering forward shapes, padding/permutation invariance, for_inference softmax, backward, and the weaver.train.onnx export path with dynamic axes. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015cGUWDpH3eU9dXgb3DrgpK
Port of the LLoCa (Lorentz Local Canonicalization) transformer from the lloca package (dev branch, commit 8a5bb43) and the tagging-guide wrapper: transformer_v2 backbone with frame-to-frame LLoCa attention (including preserve_variance), LearnedPDFrames frames-net, and the TensorReps transform machinery, wrapped as a weaver-facing LLoCaTransformerTagger. Everything runs on weaver's dense zero-padded layout; the equivectors edge convolution is a masked dense reimplementation of the upstream torch_geometric MessagePassing module, so there is no torch_geometric dependency. Attention is fixed to native scaled_dot_product_attention, and the frames-net edge standardization initializes lazily from the first batch into checkpointed buffers. Tests cover shapes, padding/permutation invariance, Lorentz invariance (without spurions), backward, ONNX export (float32 paths), and cross-checks against the upstream lloca implementation. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_015cGUWDpH3eU9dXgb3DrgpK
Port the sparse-padding flash attention support from tagging-guide (the LGATrWrapper/TransformerWrapper _forward_sparse paths) and the lgatr/lloca attention_backends packages. Both taggers gain an `attention_backend` option: - "native" (default, unchanged): dense zero-padded SDPA; required for ONNX. - "varlen": torch's native flash-attention varlen kernel (torch.nn.attention.varlen.varlen_attn, torch >= 2.10; no extra deps). - "flash": the flash-attn package, preferring the FlashAttention-3 interface and falling back to FlashAttention-2. For the varlen/flash backends the taggers switch to the packed sparse layout: padding is dropped (dense_to_sparse), per-event cu_seqlens drive block-diagonal varlen attention, global class tokens are interleaved at each event's segment start, and readout is per segment (class-token gather or segment mean). The backend is selected from the attention kwargs, as upstream. On CPU, where the varlen kernels are unavailable, the packed layout falls back to a materialized block-diagonal SDPA mask. The LLoCa frames-net stays dense (mathematically equivalent either way). The kernel wrapper zero-pads the head dim to a multiple of 8 and passes the softmax scale explicitly so it keeps the un-padded 1/sqrt(channels) (LGATr-slim head dims, e.g. 28, are not multiples of 8); fp32 inputs are cast to the autocast half dtype and back. LLoCaAttention additionally accepts token-resolved p_ref for the packed layout. Tested on CPU (packed matches dense to ~1e-7) and on an RTX 5090 with both kernels (matches dense fp32 to <=5e-4 in fp16; fwd+bwd speedup at batch 512 / maxlen 128 jagged: ~2.5x for LGATr-slim, ~1.15x for LLoCa). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QZ3Uh4TGsZvcjKMR6Huxyx
Port the xformers memory-efficient attention support from tagging-guide: attention_backend="xformers" runs block-diagonal attention over the packed tokens via xformers.ops.memory_efficient_attention with a BlockDiagonalMask built from the packed sequence lengths (host-side, so this backend incurs a device-to-host sync that the varlen/flash backends avoid). The wrapper casts fp32 inputs to the autocast dtype and zero-pads the head dim to a multiple of 8 with the softmax scale pinned to the un-padded value, as _run_varlen_kernel already does: on recent GPU architectures (e.g. sm_120) xformers has no fp32 kernel and only the fp16/bf16 flash kernels remain. Validated on an RTX 5090 (torch 2.11.0+cu128, xformers 0.0.35): packed xformers vs dense native outputs agree to 4e-6..2.4e-4 for both models, fwd+bwd, with and without bf16 amp; bit-identical to the varlen backend. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01ExsDt5azhwa5gJSMNrCEb8
Port of the tagging-guide "Transformer" baseline (model=tr): the LLoCa transformer_v2 backbone run with identity frames, reusing the existing LLoCaTransformer port. Computes the seven standardized kinematic tagging features internally, supports class-token or masked-mean readout, the packed varlen/flash/xformers attention backends, and ONNX export (with momentum_float64=False). The upstream spurions survive the identity-frames path as featureless valid tokens, reproduced as num_register_tokens. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01X8kqxhVmL1tef9t3pbeFPM
Add an auxiliary_scalars option to LGATrSlimTagger and make the existing ones actually skip the computation instead of computing all seven features and slicing them away: - PlainTransformerTagger: auxiliary_scalars=None now bypasses the momentum path entirely, so pf_vectors is unused and v may be None. - LLoCaTransformerTagger: auxiliary_scalars=None skips the global-frame features fed to the frames-net; the new local_auxiliary_scalars selects the local-frame features fed to the transformer (previously hardcoded to "all"), and None also skips the local-frame four-momentum transforms. - LGATrSlimTagger: gains auxiliary_scalars/momentum_float64 to compute the features internally from pf_vectors instead of taking them from the data config. PlainTransformerTagger and LGATrSlimTagger default to None, i.e. the features come from pf_features via the weaver data config as for ParticleTransformer; note this changes PlainTransformerTagger's default away from the upstream behaviour, which is auxiliary_scalars="all". LLoCaTransformerTagger keeps "all" on both switches. The shared helpers (get_auxiliary_scalars and the four-momentum primitives) move from LLoCaTransformer.py to a new weaver/nn/model/kinematics.py, since LLoCaTransformer imports from LGATrSlim and could not be imported back. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018SaqkBKBrGSDjHD8QPdhWi
`ParticleTransformer`'s default `Embed` applies `BatchNorm1d` + `LayerNorm` + `Linear` + activation. `use_plain_embed=True` reduces it to a bare stack of `Linear` layers, reproducing the single `linear_in` of `PlainTransformerTagger` so the two models can be compared without the input-normalization confound. Exposed on `ParticleTransformer` and both tagger wrappers, so it can be set from the command line with `-o use_plain_embed True`. It overrides `normalize_input`, is rejected together with `use_conv_embed`, and warns when `embed_dims` has more than one entry (a plain stack of linear layers collapses to a single one). Also fixes `Embed.forward`: the (N, C, P) -> (N, P, C) transpose used to sit inside the `input_bn is not None` branch, so `normalize_input=False` applied the embedding across the sequence axis instead of the channel axis. That path was never exercised before, but `use_plain_embed` needs it. `conv_embed` is `Identity` unless `use_conv_embed`, which asserts `normalize_input`, so the default path is unchanged. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018SaqkBKBrGSDjHD8QPdhWi
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
1 participant
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.
No description provided.