neurojepa

A bit-faithful Rust port of NYUMedML/Neuro-JEPA — a V-JEPA-2-derived self-supervised model for volumetric 3-D brain MRI (T1 / T2 / FLAIR) — running on the RLX compiler + runtime.

The forward pass reproduces the reference PyTorch model lane-for-lane: validated to f32 machine epsilon on identical weights for the encoder, predictor, Mixture-of-Experts, attentive pooler, and JEPA loss.

Install

This crate is standalone: it depends only on the published RLX crates (rlx-ir, rlx-runtime, rlx-flow, rlx-opt, rlx-compile, rlx-autodiff) at 0.2.8 from crates.io — no workspace, path, or unpublished dependencies.

[dependencies]
neurojepa = "0.0.1"

CPU is the default. Opt into an accelerator backend with a feature (see Backends), e.g. neurojepa = { version = "0.0.1", features = ["metal"] }.

Numerical parity (vs the reference PyTorch model, identical weights)

Component	max abs diff (CPU)	backend
Encoder (gated RoPE attn, fused qkv, exact GELU)	5.96e-8	native + compiled
Predictor	7.45e-9	native + compiled
MoE block (DeepSeek-V3 bias routing)	1.86e-7	native + compiled
Encoder with a MoE layer	4.84e-8	native + compiled
Attentive pooler (depth-1)	5.22e-8	native
JEPA L1 loss	exact	native
NeuroRunner from .safetensors (public API)	4.47e-8	native

GPU backends reproduce the encoder at cosine ≈ 1.0 (reduction-order / fast-math means they are not bit-identical, which is expected).

References are regenerated from the upstream repo by tests/gen_python_ref.sh (needs python3 + torch + einops + timm) and checked in tests/python_parity.rs.

What it reproduces

This is not a stock ViT — these reference details are matched exactly:

• per-axis 3-D conv patch embed ([pd, ph, pw]);
• fused qkv projection + per-head sigmoid attention gate (proj_attn_gate);
• the V-JEPA-2 rotate_queries_or_keys RoPE with its dual-frequency broadcast, expressed in-graph as q ⊙ COS + rotate_pairs(q) ⊙ SIN;
• per-axis position decomposition (encoder vs predictor use different strides);
• exact (erf) GELU;
• DeepSeek-V3 "bias" MoE: gate (no bias) → softmax → bias buffer for top-k selection only → route_scale, with a single shared-expert MLP. In the compiled / training graph the top-k is realized as an in-graph count_greater(e) < k routing mask, so it equals the native sparse result and stays differentiable.

Usage

Inference (native, bit-exact)

use neurojepa::{FaithfulModel, NeuroJepaConfig};

let model = FaithfulModel::load(NeuroJepaConfig::vit_base(), "model.safetensors")?;
let tokens = model.encode(&volume);          // [num_patches, embed]

Runner (native on CPU, compiled multi-backend otherwise)

use neurojepa::{NeuroRunner, MaskSampler};
use rlx_runtime::Device;

let runner = NeuroRunner::builder()
    .weights("model.safetensors")
    .device(Device::Metal)               // omit for native CPU
    .build()?;

let enc   = runner.encode_volume(&volume)?;
let masks = MaskSampler::from_config(runner.config()).sample(0);
let pred  = runner.predict(&enc, &masks)?;   // if the checkpoint has a predictor
let pool  = runner.pool(&enc)?;              // if it has an attentive pooler

CLI

cargo run --bin neurojepa -- \
    --weights model.safetensors [--config config.json] \
    [--device cpu|metal|mlx|gpu|cuda] [--predict] [--pool] [--dry]

Backends

Device::Cpu runs the native reference (faithful); any other device compiles the faithful_graph and runs it through RLX. Each is behind a feature flag forwarding to rlx-runtime:

Feature	Backend	Status (faithful encoder)
(default)	CPU	✅ bit-exact
metal	Apple Metal	✅ cosine 1.0 (verified, M4 Pro)
mlx	Apple MLX	✅ cosine 1.0 (verified, M4 Pro)
gpu	wgpu (Metal/Vulkan/DX12)	✅ cosine 1.0 (verified, M4 Pro)
coreml	CoreML / ANE	✅ cosine 1.0 (verified; rank-4 attention path)
vulkan	wgpu Vulkan	✅ compiles (needs a Vulkan loader to run)
cuda / rocm / tpu	NVIDIA / AMD / Google	✅ compiles (driver loads at runtime; needs the target hardware)

tests/backends.rs compiles + runs the encoder on every backend enabled at build time and checks it against CPU. Run all the Apple paths in one shot with:

cargo test --features metal,mlx,gpu,coreml --test backends -- --nocapture

Backends are additive feature flags forwarding to rlx-runtime; the convenience aggregates apple-silicon, apple, nvidia-gpu, amd-gpu, and all-backends are also available.

Training

build_faithful_train produces the combined encoder + predictor forward with an L1 latent loss and runs rlx_autodiff::grad_with_loss; gradients flow through the gated attention, the dual-frequency RoPE, and the (differentiable) MoE routing.

let tg = neurojepa::build_faithful_train(&cfg, &weights, &ctx_idx, &tgt_idx, 0, 1)?;
// compile tg.backward, feed ("hidden", "target", "d_output"=1.0), step the optimizer.

Scope / honest limitations

Bit-exact for standard inference at the base configuration (RoPE, GELU MLP, dense or MoE, attentive pooling). Not yet matched:

• masking is structurally faithful (3-scale block sampler) but not bit-parity (different RNG; no boundary-erosion replication);
• the training loop (EMA schedule, foreground-aware loss weighting, multi-scale mask loss) is not a 1:1 port — only the autodiff forward/loss is validated;
• not ported: SwiGLU (use_silu), pooler depth > 1 + classifier head, the non-RoPE predictor, multi-mask predictor, and bf16;
• the compiled / training MoE evaluates all experts densely (correct, but a perf trade-off vs true sparse routing).

License

GPL-3.0-only. See the repository root.