MoE-JEPA

MoE-JEPA: Mixture-of-Experts on top of V-JEPA 2 for Synthetic Image Forensics

Paper: Unifying Semantic Priors and High-Frequency Traces: Enhancing V-JEPA with Mixture-of-Experts for Robust Synthetic Image Forensics
Authors: Simone Teglia, Irene Amerini — ALCOR Lab, Sapienza University of Rome
Venue: 2026 Workshop on AI for Multimedia Forensics & Disinformation Detection (AI4MFDD) @ ECCV 2026
Code: github.com/ALCOR-Lab-DIAG/MoE-JEPA

---

Model Description

MoE-JEPA is a dual-branch image forensics detector that combines semantic priors from a frozen V-JEPA 2 backbone with high-frequency noise traces from a constrained Bayar convolutional branch. The two streams are fused through a learned per-image gate, and the combined representation is classified into three categories: Real, Fully-Synthetic, and Tampered (locally manipulated).

Architecture at a glance

Input Image
    │
    ├──► [Semantic Branch]
    │        Frozen V-JEPA 2 ViT-L encoder
    │        └─ Last 3 transformer blocks replaced with residual MoE layers
    │               • 1 frozen shared expert (preserves pretrained knowledge)
    │               • Top-2 of 6 trainable routed experts
    │        Patch tokens → Gated Attention MIL pooling → H_semantic
    │
    └──► [Noise Branch]
             BayarConv (constrained pixel predictor)
             └─ Lightweight CNN → H_noise
    
    Adaptive Gated Fusion:
        α  = sigmoid(Linear([H_semantic, H_noise]))
        H_fused = H_semantic + α · H_noise
    
    MLP Head → 3-way classification (Real / Fully-Synthetic / Tampered)

---

How to Load and Use

1. Install dependencies

pip install torch transformers huggingface_hub

2. Reconstruct the model

The model class is provided in model.py in this repository. Download it alongside the weights:

import torch
from huggingface_hub import hf_hub_download

# Download weights and model source
weights_path = hf_hub_download(repo_id="simoneteglia/MoE-JEPA", filename="model_weights.pt")
model_path   = hf_hub_download(repo_id="simoneteglia/MoE-JEPA", filename="model.py")

3. Load the model

import torch
import importlib.util, sys

# --- Dynamically import model.py from the downloaded path ---
spec = importlib.util.spec_from_file_location("moejepa_model", model_path)
mod  = importlib.util.module_from_spec(spec)
sys.modules["moejepa_model"] = mod
spec.loader.exec_module(mod)

MoEJEPADeepfakeDetector3ClassesMOD = mod.MoEJEPADeepfakeDetector3ClassesMOD

# --- Instantiate the model with the same hyperparameters as this checkpoint ---
model = MoEJEPADeepfakeDetector3ClassesMOD(
    model_name="facebook/vjepa2-vitl-fpc64-256",
    num_experts=6,
    lr=1e-5,
    aux_loss_weight=0.1,
    router_temperature=0.5,
    weight_decay=0.05,
    train_last_n_layers=3,
    lambda_energy=0.0,          
    unfreeze_last_n_blocks=0,
    pooling_strategy="mil",
    top_k=2,
    do_noise_stream=True,
    out_classes=3,
)

# --- Load the state_dict ---
checkpoint = torch.load(weights_path, map_location="cpu", weights_only=True)
model.load_state_dict(checkpoint["state_dict"])
model.eval()

print("Model loaded successfully!")
print(f"  Trained for {checkpoint['epoch']} epochs ({checkpoint['global_step']} steps)")

4. Run inference

from transformers import AutoVideoProcessor
from PIL import Image
import torch

# Label map
LABELS = {0: "Real", 1: "Fully-Synthetic", 2: "Tampered"}

# V-JEPA 2 uses a video processor; for a single image we replicate it as frames
processor    = AutoVideoProcessor.from_pretrained("facebook/vjepa2-vitl-fpc64-256", trust_remote_code=True)
image        = Image.open("your_image.jpg").convert("RGB")

# Repeat the single frame to meet the expected clip length (2 frames)
frames       = [image] * 2
inputs       = processor(videos=[frames], return_tensors="pt")
pixel_values = inputs["pixel_values_videos"]   # [1, 2, 3, H, W]
# The model's forward() accepts [B, T, C, H, W] from the video processor

with torch.no_grad():
    logits = model(pixel_values)               # [1, 3]
    pred   = torch.argmax(logits, dim=-1).item()

print(f"Prediction: {LABELS[pred]}  (logits: {logits.squeeze().tolist()})")

---

Citation

@inproceedings{teglia2026moejepa,
  title     = {Unifying Semantic Priors and High-Frequency Traces:
               Enhancing V-JEPA with Mixture-of-Experts for Robust
               Synthetic Image Forensics},
  author    = {Teglia, Simone and Amerini, Irene},
  booktitle = {Workshop on AI for Multimedia Forensics \& Disinformation
               Detection (AI4MFDD) at ECCV},
  year      = {2026},
}

---

Contact

Simone Teglia — [email protected]
ALCOR Lab, Sapienza University of Rome — alcorlab.diag.uniroma1.it

---

License

This model is released under the MIT License.