Node-JEPA

Node-JEPA learns what "normal" looks like for any graph-structured data — code, molecules, networks, circuits — and flags what deviates. It improves with use: the more you interact with it, the better it gets at distinguishing signal from noise in your specific context.

It does this by parsing your data into a graph, encoding it with a small GNN trained via JEPA (self-supervised, no labels), and comparing the result against a learned baseline. When you tell it which findings are useful, it adjusts. When you show it a new type of data, it detects the domain shift and adapts.

The model is 427K parameters. It fits on a Raspberry Pi 5.

What the graph representation buys you

A linter matches textual patterns. Node-JEPA operates on the actual structure of your code (or data): the AST, control flow, and data flow graph. This means it can find things like:

• A subprocess.Popen call on line 47 that has no error handling — not because the word "subprocess" is on a blocklist, but because the call node has no path to any except clause in the control flow graph.
• A variable assigned on line 12 that is never read — because the data flow graph has no outgoing edges from that assignment node.
• A function nested 5 control structures deep — measured by actual AST depth, not indentation heuristics.

For non-code domains, it finds structurally unusual nodes (hubs, isolates, degree anomalies) relative to what it's learned is typical for that domain.

Dreaming

The encoder includes a dream() method that rolls the predictor forward through latent space — exploring what the model "expects" to come next from a given state.

model = NodeJEPAWorldModel.from_pretrained("EPSAGR/Node-JEPA")
z = model.encode(x, edge_index)

# Free dreaming: where does latent space flow from here?
trajectory = model.dream(z, steps=10)

# Steered dreaming: interpolate toward a target through predictor-space
trajectory = model.dream(z, steps=10, target_embedding=z_target)

# Trust gate: does the dream match reality?
check = model.verify_dream(trajectory[-1], z_actual, threshold=0.15)
# → {"similarity": 0.92, "verified": True, "distance": 0.08}

verify_dream() is the trust gate for downstream decoding — nothing gets acted on unless the dreamed state is close enough to ground truth.

Quick start

from node_jepa.training.model import NodeJEPAWorldModel

model = NodeJEPAWorldModel.from_pretrained("EPSAGR/Node-JEPA")
z = model.encode(x, edge_index)  # [N, 32] node embeddings

Full pipeline

From raw input to anomaly scores in one command. Auto-detects source type.

python -m pipeline.ingest --source ./my-repo          # code
python -m pipeline.ingest --source molecule.sdf        # molecular
python -m pipeline.ingest --source ./papers/           # text

The pipeline extracts entities (tree-sitter for code, RDKit for molecules, LLM for text), infers relations, assembles a graph, encodes it, and scores each node against a learned baseline. LLM backend is configurable — ollama, llama.cpp, HF, or any OpenAI-compatible endpoint. See pipeline/config.yaml.

GitHub repos

from pipeline.adapters.github import GitHubAdapter

adapter = GitHubAdapter("owner/repo", token="...")
pyg_data = adapter.to_pyg()  # files, commits, co-modification edges

Proactive mode

Run as a daemon that watches your codebase and pushes findings when something changes:

python -m node_jepa.deploy.daemon --repo /path/to/code --webhook http://your-agent:8080/findings

It learns from your behavior. Fix an unhandled-exception finding, it keeps surfacing those. Ignore missing-docstring findings for a week, it stops pushing them. State persists across restarts.

API for external tools

Expose analysis as an HTTP endpoint. Any LLM agent (Cursor, Cline, aider) can call it to get graph-grounded findings before deciding what to edit.

python -m node_jepa.deploy.api --repo /path/to/code --port 7685

GET /analyze                     → all findings
GET /analyze?severity=high       → high severity only
GET /analyze?kind=unhandled_call → specific finding type
GET /graph/embedding             → GNN embedding vector

Auto-domain

Feed it any graph. It figures out the domain (by density, size, feature shape), allocates structure tokens, and trains itself.

from node_jepa.deploy.auto_domain import AutoDomainManager

mgr = AutoDomainManager(model)
mgr.ingest(pyg_data)          # accepts PyG, JSON, or NetworkX
mgr.train_step()              # JEPA self-supervised, no labels
findings = mgr.analyze()      # what's unusual for this domain?
mgr.feedback("hub_node", useful=True)

EWC regularization prevents forgetting previously learned domains.

Training

The training loss is: L = L_jepa + lambda_dec * L_decorrelation + mu_rec * L_reconstruction + lambda_smooth * L_smoothness

Smoothness regularization penalizes high cosine distance between adjacent node embeddings, keeping the latent space navigable for dreaming and Phase 3 decoding. Trajectory snapshots (embeddings + node IDs) are saved every 50 epochs for historical latent mapping.

All hyperparameters are configurable in the sweep config. See SCRIPTS.md for which training script to use.

Benchmark results

Task	Result	Context
Cora node classification (linear probe)	82.8%	GraphMAE: 84.2%, BGRL: 82.8%
PubMed node classification (linear probe)	83.2%	Beats GraphMAE (81.1%)
QM9 HOMO regression (linear probe)	R²=0.763	2K molecules, target was 0.5
Code embedding rank	31/32 active dims	No representational collapse
Git transition rollout	1.02x error growth over 10 steps	Bounded, non-diverging

Downstream code intelligence (Devign, POJ-104) is weak. The encoder was only trained on its own codebase — it needs broader pretraining to generalize.

Install

pip install -r requirements.txt          # core model
pip install -r pipeline/requirements.txt  # full pipeline

torch-geometric requires matching PyTorch + CUDA versions. See PyG install guide.

Project layout

node_jepa/              core model
  core/                 JEPA loss, masking, EMA, smoothness regularization
  encoder/              GAT backbone + per-domain structure tokens
  deploy/               daemon, API, auto-domain, case bank, code analyzer
  training/             NodeJEPAWorldModel (from_pretrained / save_pretrained / dream)
  world_model/          transition predictor
  harness/              agent memory, surprise, planner, actions
  learning/             experience buffer, EWC, imagination engine
  perception/           Tree-sitter code extraction, graph builder
  evaluation/           linear probe, benchmarks, eval CLI
  heads/                task heads (regression, classification, forecast)
  curriculum/           multi-domain training phases
  config/               default hyperparameters
  infra/                Pi 5 audit, ONNX export
pipeline/               end-to-end ingestion + scoring
  adapters/             GitHub adapter
  kg_builder/           entity extraction, relation inference, schema discovery
  scorer.py             anomaly scoring (KMeans → learned classifier)
  encoder.py            encoder + predictor wrapper (encode, dream, verify_dream)
  ingest.py             CLI entrypoint
scripts/                training scripts (see SCRIPTS.md)
tests/                  unit tests
pretrained/             checkpoints

Checkpoints

File	What
model.safetensors	Full model, HF format, 427K params
encoder.pt2	Encoder only, 71KB
pretrained/v2_scaled.pt2	Full model, PyTorch format
pretrained/git_transition_2000.pt2	Code transition model
pretrained/md17_transition_v2.pt2	Molecular transition model
node-jepa-phase1	QM9 encoder, R²=0.763

How it learns

World model — JEPA self-supervised: mask graph nodes, predict from context. Per-domain structure tokens adapt to each domain's topology. Smoothness loss keeps the latent space navigable. EWC prevents catastrophic forgetting across domains. Weights persist across restarts.

Dreaming — the predictor rolls forward through latent space, generating trajectories that can be steered toward targets or left to free-associate. verify_dream() gates whether a dreamed state is trustworthy enough to act on.

Case bank — every interaction stores (embedding, action, reward). On the next similar situation, the system retrieves what worked before. Cosine similarity over GNN embeddings.

Finding memory — tracks which findings you act on vs ignore. Computes per-type fix rates. Adjusts priority: high fix rate = surface more, low fix rate = suppress. Persists across restarts.

Citation

@software{node_jepa,
  title = {Node-JEPA},
  author = {EPSAGR},
  year = {2026},
  url = {https://huggingface.co/EPSAGR/Node-JEPA}
}

MIT license.