diff --git a/flaxdiff/data/sources/images.py b/flaxdiff/data/sources/images.py
index 569373c..5616471 100644
--- a/flaxdiff/data/sources/images.py
+++ b/flaxdiff/data/sources/images.py
@@ -187,6 +187,8 @@ def map(self, element) -> Dict[str, jnp.array]:
                 results = self.tokenize(caption)
                 return {
                     "image": image,
+                    # the class index, which the JEPA linear/kNN probes score against
+                    "label": np.int32(element['label']),
                     "text": {
                         "input_ids": results['input_ids'][0],
                         "attention_mask": results['attention_mask'][0],
diff --git a/flaxdiff/jepa/__init__.py b/flaxdiff/jepa/__init__.py
new file mode 100644
index 0000000..1abef58
--- /dev/null
+++ b/flaxdiff/jepa/__init__.py
@@ -0,0 +1,4 @@
+from .masking import MultiBlockMask, multi_block_mask
+from .models import JepaEncoder, JepaVideoEncoder, JepaPredictor, TokenStack, FactorizedTokenStack
+from .objective import JepaObjective, representation_health, normalize_targets
+from .probes import get_linear_probe_metric, get_knn_probe_metric, linear_probe, knn_probe
diff --git a/flaxdiff/jepa/masking.py b/flaxdiff/jepa/masking.py
new file mode 100644
index 0000000..042a048
--- /dev/null
+++ b/flaxdiff/jepa/masking.py
@@ -0,0 +1,120 @@
+"""I-JEPA multi-block masking over a patch grid.
+
+The paper samples several target blocks with a random scale and aspect ratio
+and takes the context to be what is left. Under jit every mask has to come out
+the same shape on every step, so the geometry is resolved once, at
+construction, against the static patch grid:
+
+  - the block area is fixed to the value inside the configured scale range that
+    admits the most (height, width) factorizations within the aspect ratio
+    range, so each sample still draws a genuinely different block shape and
+    position, and every target block has exactly the same token count;
+  - the context is a uniformly random subset of the complement of the target
+    union, sized S - num_targets * area. That size is always available (the
+    union is at most num_targets disjoint blocks), and when the blocks happen
+    to overlap the context simply drops the surplus rather than changing shape.
+
+Indices refer to positions in the scan-ordered token sequence that
+PatchSequenceEmbed produces, and come out sorted so that an SSM mixer scans
+them in a meaningful order.
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Tuple
+
+import jax
+import jax.numpy as jnp
+
+
+def _factorizations(area: int, grid: Tuple[int, int], aspect: Tuple[float, float]):
+    """(h, w) pairs of the given area that fit the grid and the aspect range."""
+    H_P, W_P = grid
+    pairs = []
+    for h in range(1, min(H_P, area) + 1):
+        if area % h:
+            continue
+        w = area // h
+        if w <= W_P and aspect[0] <= h / w <= aspect[1]:
+            pairs.append((h, w))
+    return pairs
+
+
+@dataclass(frozen=True)
+class MultiBlockMask:
+    """Static mask geometry for one patch grid, plus the sampler over it."""
+    grid: Tuple[int, int]
+    num_targets: int
+    block_shapes: Tuple[Tuple[int, int], ...]
+    num_context: int
+
+    @property
+    def num_patches(self) -> int:
+        return self.grid[0] * self.grid[1]
+
+    @property
+    def block_area(self) -> int:
+        h, w = self.block_shapes[0]
+        return h * w
+
+    def sample(self, rng: jax.Array, batch_size: int):
+        """Draw (context_idx [B, num_context], target_idx [B, M, block_area])."""
+        H_P, W_P = self.grid
+        S = self.num_patches
+        M = self.num_targets
+
+        heights = jnp.asarray([h for h, _ in self.block_shapes])
+        widths = jnp.asarray([w for _, w in self.block_shapes])
+        # flat offsets of a block's tokens from its top-left corner, per shape
+        offsets = jnp.asarray([[(i // w) * W_P + (i % w) for i in range(h * w)]
+                               for h, w in self.block_shapes])
+
+        shape_key, pos_key, context_key = jax.random.split(rng, 3)
+        choice = jax.random.randint(shape_key, (batch_size, M), 0, len(self.block_shapes))
+        corner = jax.random.uniform(pos_key, (batch_size, M, 2))
+        top = jnp.floor(corner[..., 0] * (H_P - heights[choice] + 1)).astype(jnp.int32)
+        left = jnp.floor(corner[..., 1] * (W_P - widths[choice] + 1)).astype(jnp.int32)
+        target_idx = (top * W_P + left)[..., None] + offsets[choice]
+
+        is_target = jnp.zeros((batch_size, S), dtype=bool).at[
+            jnp.arange(batch_size)[:, None], target_idx.reshape(batch_size, -1)].set(True)
+        # target tokens sort last, so the first num_context entries of a random
+        # ordering are a uniform sample of the complement
+        order = jax.random.uniform(context_key, (batch_size, S)) + is_target
+        context_idx = jnp.sort(jnp.argsort(order, axis=1)[:, :self.num_context], axis=1)
+        return context_idx, target_idx
+
+
+def multi_block_mask(
+    grid: Tuple[int, int],
+    num_targets: int = 4,
+    scale: Tuple[float, float] = (0.15, 0.2),
+    aspect: Tuple[float, float] = (0.75, 1.5),
+) -> MultiBlockMask:
+    """Resolve the I-JEPA mask geometry for a patch grid."""
+    S = grid[0] * grid[1]
+    candidates = [
+        (area, _factorizations(area, grid, aspect))
+        for area in range(max(1, math.ceil(scale[0] * S)), math.floor(scale[1] * S) + 1)
+    ]
+    candidates = [(area, shapes) for area, shapes in candidates if shapes]
+    if not candidates:
+        raise ValueError(
+            f"No block of scale {scale} on a {grid} grid has an aspect ratio in "
+            f"{aspect}; widen one of the ranges or use a finer patch grid.")
+
+    midpoint = 0.5 * (scale[0] + scale[1]) * S
+    area, shapes = min(candidates, key=lambda c: (-len(c[1]), abs(c[0] - midpoint)))
+
+    num_context = S - num_targets * area
+    if num_context <= 0:
+        raise ValueError(
+            f"{num_targets} target blocks of {area} tokens leave no context on a "
+            f"{grid} grid ({S} tokens).")
+
+    return MultiBlockMask(
+        grid=grid,
+        num_targets=num_targets,
+        block_shapes=tuple(shapes),
+        num_context=num_context,
+    )
diff --git a/flaxdiff/jepa/models.py b/flaxdiff/jepa/models.py
new file mode 100644
index 0000000..42e4e64
--- /dev/null
+++ b/flaxdiff/jepa/models.py
@@ -0,0 +1,316 @@
+"""JEPA encoders and predictors, arranged from the shared DiT machinery.
+
+A JEPA encoder is the DiT sandwich without the diffusion parts: patchify with
+the 2D sincos signal, run unmodulated ModulatedBlocks over the tokens, norm.
+There is no timestep to condition on, so the blocks run in their unmodulated
+(plain pre-norm) mode, and the mixer is still pluggable - mixer patterns with
+'ssm' give a linear-time S5 encoder.
+
+Position never comes from RoPE here. Both the encoder and the predictor work
+on a masked subset of the sequence, where a token's index in the sequence is
+not its position on the grid, so RoPE is left at identity and position is
+carried entirely by the 2D sincos embedding that travels with each token.
+"""
+
+import jax.numpy as jnp
+from flax import linen as nn
+from typing import Optional, Sequence, Tuple
+from flax.typing import Dtype, PrecisionLike
+
+from ..models.dit_common import (
+    PatchSequenceEmbed, ModulatedBlock, build_block_pattern,
+    identity_rope_freqs, scan_ordered_pos_embed,
+)
+from ..models.vit_common import RotaryEmbedding
+
+
+def gather_tokens(tokens, idx):
+    """Select tokens per sample: [B, S, F] and [B, N] -> [B, N, F]."""
+    return jnp.take_along_axis(tokens, idx[..., None], axis=-2)
+
+
+class TokenStack(nn.Module):
+    """A stack of unmodulated blocks over a token sequence."""
+    features: int
+    num_layers: int
+    num_heads: int
+    mlp_ratio: int = 4
+    ssm_attention_ratio: str = "all-attn"
+    block_pattern: Optional[Sequence[str]] = None
+    ssm_state_dim: int = 64
+    bidirectional_ssm: bool = True
+    dropout_rate: float = 0.0
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    force_fp32_for_softmax: bool = True
+    norm_epsilon: float = 1e-5
+    qk_norm: bool = False
+    attention_impl: Optional[str] = None
+
+    def setup(self):
+        pattern = build_block_pattern(
+            self.num_layers, self.ssm_attention_ratio, self.block_pattern)
+        self.blocks = [
+            ModulatedBlock(
+                features=self.features,
+                num_heads=self.num_heads,
+                mixer='ssm' if kind == 'ssm' else 'attention',
+                modulated=False,
+                mlp_ratio=self.mlp_ratio,
+                dropout_rate=self.dropout_rate,
+                dtype=self.dtype,
+                precision=self.precision,
+                force_fp32_for_softmax=self.force_fp32_for_softmax,
+                norm_epsilon=self.norm_epsilon,
+                qk_norm=self.qk_norm,
+                attention_impl=self.attention_impl,
+                ssm_state_dim=self.ssm_state_dim,
+                bidirectional_ssm=self.bidirectional_ssm,
+                name=f"block_{i}",
+            ) for i, kind in enumerate(pattern)
+        ]
+
+    def __call__(self, tokens, freqs_cis=None, train: bool = False):
+        if freqs_cis is None:
+            freqs_cis = identity_rope_freqs(tokens.shape[-2], self.features // self.num_heads)
+        for block in self.blocks:
+            tokens = block(tokens, conditioning=None, freqs_cis=freqs_cis, train=train)
+        return tokens
+
+
+class FactorizedTokenStack(nn.Module):
+    """Spatial then temporal blocks over [B, T, N, F], as in VideoDiT.
+
+    Time is a real 1D axis that masking never touches, so the temporal half
+    keeps genuine RoPE while the spatial half runs at identity.
+    """
+    features: int
+    num_layers: int
+    num_heads: int
+    max_frames: int = 1024
+    mlp_ratio: int = 4
+    ssm_attention_ratio: str = "all-attn"
+    block_pattern: Optional[Sequence[str]] = None
+    ssm_state_dim: int = 64
+    bidirectional_ssm: bool = True
+    dropout_rate: float = 0.0
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    force_fp32_for_softmax: bool = True
+    norm_epsilon: float = 1e-5
+    qk_norm: bool = False
+    attention_impl: Optional[str] = None
+
+    def setup(self):
+        def stack(name, num_layers, ratio):
+            return TokenStack(
+                features=self.features, num_layers=num_layers, num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio, ssm_attention_ratio=ratio,
+                ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm,
+                dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision,
+                force_fp32_for_softmax=self.force_fp32_for_softmax,
+                norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm,
+                attention_impl=self.attention_impl, name=name)
+
+        # one spatial and one temporal block per layer, built as single-block
+        # stacks so the two halves can be interleaved
+        pattern = build_block_pattern(
+            self.num_layers, self.ssm_attention_ratio, self.block_pattern)
+        self.spatial = [stack(f"spatial_{i}", 1, "all-ssm" if kind == 'ssm' else "all-attn")
+                        for i, kind in enumerate(pattern)]
+        self.temporal = [stack(f"temporal_{i}", 1, "all-attn") for i in range(self.num_layers)]
+        self.temporal_rope = RotaryEmbedding(
+            dim=self.features // self.num_heads, max_seq_len=self.max_frames,
+            name="temporal_rope")
+
+    def __call__(self, tokens, train: bool = False):
+        B, T, N, F = tokens.shape
+        freqs_temporal = self.temporal_rope(seq_len=T)
+
+        tokens = tokens.reshape(B * T, N, F)
+        for spatial, temporal in zip(self.spatial, self.temporal):
+            tokens = spatial(tokens, train=train)
+            tokens = tokens.reshape(B, T, N, F).transpose(0, 2, 1, 3).reshape(B * N, T, F)
+            tokens = temporal(tokens, freqs_cis=freqs_temporal, train=train)
+            tokens = tokens.reshape(B, N, T, F).transpose(0, 2, 1, 3).reshape(B * T, N, F)
+        return tokens.reshape(B, T, N, F)
+
+
+class JepaEncoder(nn.Module):
+    """ViT over an image, optionally restricted to a subset of its patches."""
+    patch_size: int = 16
+    emb_features: int = 384
+    num_layers: int = 12
+    num_heads: int = 6
+    mlp_ratio: int = 4
+    ssm_attention_ratio: str = "all-attn"
+    ssm_state_dim: int = 64
+    bidirectional_ssm: bool = True
+    dropout_rate: float = 0.0
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    force_fp32_for_softmax: bool = True
+    norm_epsilon: float = 1e-5
+    qk_norm: bool = False
+    attention_impl: Optional[str] = None
+    use_hilbert: bool = False
+    use_zigzag: bool = False
+
+    @property
+    def scan_order(self):
+        assert not (self.use_hilbert and self.use_zigzag), \
+            "use_hilbert and use_zigzag are mutually exclusive"
+        return 'hilbert' if self.use_hilbert else 'zigzag' if self.use_zigzag else 'raster'
+
+    def setup(self):
+        self.embed = PatchSequenceEmbed(
+            patch_size=self.patch_size,
+            emb_features=self.emb_features,
+            scan_order=self.scan_order,
+            dtype=self.dtype,
+            precision=self.precision,
+        )
+        self.stack = TokenStack(
+            features=self.emb_features, num_layers=self.num_layers,
+            num_heads=self.num_heads, mlp_ratio=self.mlp_ratio,
+            ssm_attention_ratio=self.ssm_attention_ratio,
+            ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm,
+            dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision,
+            force_fp32_for_softmax=self.force_fp32_for_softmax,
+            norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm,
+            attention_impl=self.attention_impl,
+        )
+        self.norm = nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="norm")
+
+    def __call__(self, x, token_idx=None, train: bool = False):
+        tokens, _ = self.embed(x)
+        if token_idx is not None:
+            tokens = gather_tokens(tokens, token_idx)
+        return self.norm(self.stack(tokens, train=train))
+
+
+class JepaVideoEncoder(nn.Module):
+    """Factorized spatial-temporal encoder over (B, T, H, W, C).
+
+    token_idx selects a tubelet: the same patch positions in every frame, so
+    the factorized layout survives masking untouched.
+    """
+    patch_size: int = 16
+    emb_features: int = 384
+    num_layers: int = 12
+    num_heads: int = 6
+    mlp_ratio: int = 4
+    ssm_attention_ratio: str = "all-attn"
+    ssm_state_dim: int = 64
+    bidirectional_ssm: bool = True
+    dropout_rate: float = 0.0
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    force_fp32_for_softmax: bool = True
+    norm_epsilon: float = 1e-5
+    qk_norm: bool = False
+    attention_impl: Optional[str] = None
+    use_hilbert: bool = False
+    use_zigzag: bool = False
+
+    @property
+    def scan_order(self):
+        assert not (self.use_hilbert and self.use_zigzag), \
+            "use_hilbert and use_zigzag are mutually exclusive"
+        return 'hilbert' if self.use_hilbert else 'zigzag' if self.use_zigzag else 'raster'
+
+    def setup(self):
+        self.embed = PatchSequenceEmbed(
+            patch_size=self.patch_size,
+            emb_features=self.emb_features,
+            scan_order=self.scan_order,
+            dtype=self.dtype,
+            precision=self.precision,
+        )
+        self.stack = FactorizedTokenStack(
+            features=self.emb_features, num_layers=self.num_layers,
+            num_heads=self.num_heads, mlp_ratio=self.mlp_ratio,
+            ssm_attention_ratio=self.ssm_attention_ratio,
+            ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm,
+            dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision,
+            force_fp32_for_softmax=self.force_fp32_for_softmax,
+            norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm,
+            attention_impl=self.attention_impl,
+        )
+        self.norm = nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="norm")
+
+    def __call__(self, x, token_idx=None, train: bool = False):
+        B, T, H, W, C = x.shape
+        tokens, _ = self.embed(x.reshape(B * T, H, W, C))
+        if token_idx is not None:
+            tokens = gather_tokens(tokens, jnp.repeat(token_idx, T, axis=0))
+        tokens = tokens.reshape(B, T, tokens.shape[1], self.emb_features)
+        return self.norm(self.stack(tokens, train=train))
+
+
+class JepaPredictor(nn.Module):
+    """Narrow transformer from context embeddings to target embeddings.
+
+    Context tokens are projected down, mask tokens stand in for the targets,
+    and both carry the sincos signal for the grid position they belong to.
+    """
+    grid: Tuple[int, int] = (14, 14)
+    emb_features: int = 384      # encoder width, in and out
+    predictor_features: int = 192
+    num_layers: int = 6
+    num_heads: int = 6
+    mlp_ratio: int = 4
+    ssm_attention_ratio: str = "all-attn"
+    ssm_state_dim: int = 64
+    bidirectional_ssm: bool = True
+    dropout_rate: float = 0.0
+    dtype: Optional[Dtype] = None
+    precision: PrecisionLike = None
+    force_fp32_for_softmax: bool = True
+    norm_epsilon: float = 1e-5
+    qk_norm: bool = False
+    attention_impl: Optional[str] = None
+    scan_order: str = 'raster'
+    factorized: bool = False     # space-time blocks, for video
+
+    def setup(self):
+        self.proj_in = nn.Dense(features=self.predictor_features, dtype=self.dtype,
+                                precision=self.precision, name="proj_in")
+        self.mask_token = self.param(
+            "mask_token", nn.initializers.normal(0.02), (1, 1, self.predictor_features))
+        stack_kwargs = dict(
+            features=self.predictor_features, num_layers=self.num_layers,
+            num_heads=self.num_heads, mlp_ratio=self.mlp_ratio,
+            ssm_attention_ratio=self.ssm_attention_ratio,
+            ssm_state_dim=self.ssm_state_dim, bidirectional_ssm=self.bidirectional_ssm,
+            dropout_rate=self.dropout_rate, dtype=self.dtype, precision=self.precision,
+            force_fp32_for_softmax=self.force_fp32_for_softmax,
+            norm_epsilon=self.norm_epsilon, qk_norm=self.qk_norm,
+            attention_impl=self.attention_impl,
+        )
+        self.stack = (FactorizedTokenStack(**stack_kwargs) if self.factorized
+                      else TokenStack(**stack_kwargs))
+        self.norm = nn.LayerNorm(epsilon=self.norm_epsilon, dtype=self.dtype, name="norm")
+        self.proj_out = nn.Dense(features=self.emb_features, dtype=self.dtype,
+                                 precision=self.precision, name="proj_out")
+
+    def __call__(self, context, context_idx, target_idx, train: bool = False):
+        """context: [B, (T,) N_ctx, F] -> predictions [B, (T,) N_tgt, F]."""
+        pos_embed = jnp.asarray(
+            scan_ordered_pos_embed(self.predictor_features, *self.grid, self.scan_order),
+            dtype=self.dtype or jnp.float32)
+
+        def positions(idx):
+            pos = pos_embed[idx]                       # [B, N, P]
+            return pos[:, None] if self.factorized else pos
+
+        num_target_tokens = target_idx.shape[-1]
+        context = self.proj_in(context) + positions(context_idx)
+        targets = self.mask_token + positions(target_idx)
+        if self.factorized:
+            targets = jnp.broadcast_to(
+                targets, (*context.shape[:-2], num_target_tokens, self.predictor_features))
+
+        tokens = jnp.concatenate([context, targets], axis=-2)
+        tokens = self.stack(tokens, train=train)
+        return self.proj_out(self.norm(tokens[..., -num_target_tokens:, :]))
diff --git a/flaxdiff/jepa/objective.py b/flaxdiff/jepa/objective.py
new file mode 100644
index 0000000..0c37fde
--- /dev/null
+++ b/flaxdiff/jepa/objective.py
@@ -0,0 +1,152 @@
+"""The I-JEPA / V-JEPA objective.
+
+Predict the representation of masked target blocks from the representation of
+the visible context, in latent space. Three moving parts:
+
+  - the context encoder sees only the context tokens and is trained;
+  - the target encoder sees the whole image and is not: it is the EMA of the
+    context encoder, so it lives in the trainer's ema_params rather than in a
+    parameter subtree of its own, and its branch is stop_gradient'd;
+  - the predictor maps context embeddings plus target positions to the target
+    representations.
+
+Targets are layer-normalized (no learned affine) before the L2 loss, which is
+what keeps the scale of the prediction problem fixed as the encoder drifts.
+The paper's L2 is used rather than the reference implementation's smooth-L1:
+the LN already bounds the target scale, so there are no outliers for smooth-L1
+to protect against, and L2 keeps the loss directly comparable to the paper.
+
+The characteristic failure is silent: both encoders quietly agree on a
+constant and the loss goes to zero. representation_health is reported on every
+step so that collapse is visible in the training curves rather than at the end
+of a probe run.
+"""
+
+from typing import Dict, Tuple
+
+import jax
+import jax.numpy as jnp
+import optax
+
+from ..trainer.objectives import Objective, EMASpec
+from .masking import MultiBlockMask
+
+CONTEXT_ENCODER = "context_encoder"
+PREDICTOR = "predictor"
+
+
+def representation_health(z) -> Dict[str, jax.Array]:
+    """Collapse telemetry for pooled embeddings [B, D].
+
+    repr_std is the per-dimension standard deviation across the batch: it goes
+    to zero exactly when the encoder stops distinguishing inputs. repr_cov_offdiag
+    is the RMS magnitude of the off-diagonal covariance, which rises when the
+    dimensions become redundant (dimensional collapse) even while repr_std holds.
+    """
+    batch_size, dim = z.shape
+    centered = z - jnp.mean(z, axis=0, keepdims=True)
+    cov = (centered.T @ centered) / max(batch_size - 1, 1)
+    off_diagonal = cov * (1.0 - jnp.eye(dim, dtype=cov.dtype))
+    return {
+        "repr_std": jnp.mean(jnp.std(z, axis=0)),
+        "repr_cov_offdiag": jnp.sqrt(jnp.sum(off_diagonal ** 2) / max(dim * (dim - 1), 1)),
+    }
+
+
+def normalize_targets(x, epsilon: float = 1e-6):
+    """Feature-wise layer norm with no learned affine.
+
+    Applied to the target encoder's output so the prediction problem keeps a
+    fixed scale as the encoder drifts, and so shrinking the representation is
+    not a way to lower the loss.
+    """
+    mean = jnp.mean(x, axis=-1, keepdims=True)
+    variance = jnp.var(x, axis=-1, keepdims=True)
+    return (x - mean) * jax.lax.rsqrt(variance + epsilon)
+
+
+class JepaObjective(Objective):
+    """Joint-embedding prediction over images (B,H,W,C) or video (B,T,H,W,C)."""
+
+    tag = "jepa"
+
+    def __init__(
+        self,
+        encoder,
+        predictor,
+        mask: MultiBlockMask,
+        sample_data_key: str,
+        sample_data_shape: Tuple[int, ...],
+        momentum: Tuple[float, float] = (0.996, 1.0),
+        momentum_steps: int = 100_000,
+    ):
+        self.encoder = encoder
+        self.predictor = predictor
+        self.mask = mask
+        self.sample_data_key = sample_data_key
+        self.sample_data_shape = tuple(sample_data_shape)
+        self.is_video = len(self.sample_data_shape) == 4
+        self.ema = EMASpec(
+            decay=optax.linear_schedule(momentum[0], momentum[1], momentum_steps),
+            path=("params", CONTEXT_ENCODER),
+        )
+
+    def init_params(self, rng):
+        encoder_rng, predictor_rng = jax.random.split(rng)
+        sample = jnp.ones((1, *self.sample_data_shape))
+        context_idx = jnp.arange(self.mask.num_context, dtype=jnp.int32)[None]
+        target_idx = jnp.arange(self.mask.block_area, dtype=jnp.int32)[None]
+
+        encoder = self.encoder.init(encoder_rng, sample, context_idx)
+        context = self.encoder.apply(encoder, sample, context_idx)
+        predictor = self.predictor.init(predictor_rng, context, context_idx, target_idx)
+        return {"params": {CONTEXT_ENCODER: encoder["params"],
+                           PREDICTOR: predictor["params"]}}
+
+    def encode(self, encoder_params, data, token_idx=None, train=False, rngs=None):
+        return self.encoder.apply({"params": encoder_params}, data, token_idx,
+                                  train=train, rngs=rngs)
+
+    def loss(self, params, ema_params, batch, rng, step):
+        data = (jnp.asarray(batch[self.sample_data_key], dtype=jnp.float32) - 127.5) / 127.5
+        batch_size = data.shape[0]
+        mask_rng, dropout_rng = jax.random.split(rng)
+        context_idx, target_idx = self.mask.sample(mask_rng, batch_size)
+        num_targets = self.mask.num_targets
+
+        # Target branch: the whole view through the EMA encoder, no gradient
+        full = normalize_targets(self.encode(ema_params["params"][CONTEXT_ENCODER], data))
+        # [B, (T,) S, F] -> [B, M, (T,) n_tgt, F]
+        frame_axis = (1,) if self.is_video else ()
+        gather_idx = target_idx.reshape(batch_size, num_targets, *frame_axis, -1, 1)
+        targets = jax.lax.stop_gradient(
+            jnp.take_along_axis(full[:, None], gather_idx, axis=-2))
+
+        context = self.encode(
+            params["params"][CONTEXT_ENCODER], data, context_idx,
+            train=True, rngs={"dropout": dropout_rng})
+
+        # Each target block is predicted from the same context: fold the block
+        # axis into the batch so one predictor call covers all M of them
+        repeated = jnp.repeat(context, num_targets, axis=0)
+        predictions = self.predictor.apply(
+            {"params": params["params"][PREDICTOR]},
+            repeated,
+            jnp.repeat(context_idx, num_targets, axis=0),
+            target_idx.reshape(batch_size * num_targets, -1),
+            train=True, rngs={"dropout": dropout_rng},
+        ).reshape(targets.shape)
+
+        loss = jnp.mean((predictions - targets) ** 2)
+        pooled = jnp.mean(full, axis=tuple(range(1, full.ndim - 1)))
+        return loss, representation_health(pooled)
+
+    def make_validation_step(self, **_):
+        """Frozen target-encoder embeddings, which the probes score."""
+        def embed(val_state, batch):
+            data = (jnp.asarray(batch[self.sample_data_key], dtype=jnp.float32) - 127.5) / 127.5
+            features = self.encode(
+                val_state.ema_params["params"][CONTEXT_ENCODER], data)
+            return jnp.mean(features, axis=tuple(range(1, features.ndim - 1)))
+
+        return embed
diff --git a/flaxdiff/jepa/probes.py b/flaxdiff/jepa/probes.py
new file mode 100644
index 0000000..1adb196
--- /dev/null
+++ b/flaxdiff/jepa/probes.py
@@ -0,0 +1,79 @@
+"""Frozen-encoder probes, as evaluation metrics the validation loop can run.
+
+A JEPA run has no samples to look at, so representation quality is the only
+signal that the training curve cannot give you. Both probes score the pooled
+embeddings the objective's validation step produces against the batch labels,
+fitting on the first half of the batch and scoring on the second - cheap
+enough to run every epoch, and honest, since nothing is scored on data it was
+fit on. They measure trend, not absolute transfer accuracy; a full-dataset
+probe is a separate offline job.
+"""
+
+import jax
+import jax.numpy as jnp
+import optax
+
+from ..metrics.common import EvaluationMetric
+
+
+def _split(embeddings, labels):
+    fit = len(embeddings) // 2
+    if fit == 0:
+        raise ValueError("probes need at least two samples per validation batch")
+    return embeddings[:fit], labels[:fit], embeddings[fit:], labels[fit:]
+
+
+def linear_probe(embeddings, labels, num_classes: int, steps: int = 100,
+                 learning_rate: float = 1e-2, weight_decay: float = 1e-4):
+    """Accuracy of a logistic regression fit on half the batch, scored on the rest."""
+    fit_x, fit_y, test_x, test_y = _split(embeddings, labels)
+    mean, std = jnp.mean(fit_x, axis=0), jnp.std(fit_x, axis=0) + 1e-6
+    fit_x, test_x = (fit_x - mean) / std, (test_x - mean) / std
+
+    params = {"w": jnp.zeros((embeddings.shape[-1], num_classes)),
+              "b": jnp.zeros((num_classes,))}
+    optimizer = optax.adamw(learning_rate, weight_decay=weight_decay)
+
+    def objective(p):
+        logits = fit_x @ p["w"] + p["b"]
+        return jnp.mean(optax.softmax_cross_entropy_with_integer_labels(logits, fit_y))
+
+    def step(carry, _):
+        p, opt_state = carry
+        grads = jax.grad(objective)(p)
+        updates, opt_state = optimizer.update(grads, opt_state, p)
+        return (optax.apply_updates(p, updates), opt_state), None
+
+    (params, _), _ = jax.lax.scan(step, (params, optimizer.init(params)), None, length=steps)
+    predicted = jnp.argmax(test_x @ params["w"] + params["b"], axis=-1)
+    return jnp.mean(predicted == test_y)
+
+
+def knn_probe(embeddings, labels, num_classes: int, k: int = 20):
+    """Cosine k-NN accuracy, fit half against scored half."""
+    fit_x, fit_y, test_x, test_y = _split(embeddings, labels)
+    fit_x = fit_x / (jnp.linalg.norm(fit_x, axis=-1, keepdims=True) + 1e-8)
+    test_x = test_x / (jnp.linalg.norm(test_x, axis=-1, keepdims=True) + 1e-8)
+
+    similarity = test_x @ fit_x.T
+    neighbours = jnp.argsort(-similarity, axis=-1)[:, :min(k, fit_x.shape[0])]
+    votes = jnp.sum(jax.nn.one_hot(fit_y[neighbours], num_classes), axis=1)
+    return jnp.mean(jnp.argmax(votes, axis=-1) == test_y)
+
+
+def get_linear_probe_metric(num_classes: int, label_key: str = "label", **kwargs):
+    return EvaluationMetric(
+        function=lambda embeddings, batch: linear_probe(
+            embeddings, jnp.asarray(batch[label_key]), num_classes, **kwargs),
+        name="linear_probe_accuracy",
+        higher_is_better=True,
+    )
+
+
+def get_knn_probe_metric(num_classes: int, label_key: str = "label", k: int = 20):
+    return EvaluationMetric(
+        function=lambda embeddings, batch: knn_probe(
+            embeddings, jnp.asarray(batch[label_key]), num_classes, k=k),
+        name="knn_probe_accuracy",
+        higher_is_better=True,
+    )
diff --git a/flaxdiff/metrics/fid.py b/flaxdiff/metrics/fid.py
new file mode 100644
index 0000000..d11d13b
--- /dev/null
+++ b/flaxdiff/metrics/fid.py
@@ -0,0 +1,75 @@
+from .common import EvaluationMetric
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+# The FID InceptionV3 is ~90MB of weights; one copy is enough for every metric
+# built from this module.
+_inception_cache: dict = {}
+
+
+def _get_inception():
+    """Cached (model, params) for the pool3 feature extractor."""
+    if 'inception' not in _inception_cache:
+        from .inception import InceptionV3
+        print("[metrics] Loading InceptionV3 FID weights (cached for reuse)...")
+        model = InceptionV3(pretrained=True)
+        params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 299, 299, 3)))
+        _inception_cache['inception'] = (model, params)
+    return _inception_cache['inception']
+
+
+def frechet_distance(mu_a, sigma_a, mu_b, sigma_b, eps=1e-6) -> float:
+    """Frechet distance between two multivariate gaussians.
+
+    Runs on the host through scipy: the matrix square root of the covariance
+    product has no jax equivalent, and FID is computed once per validation
+    batch so the transfer is irrelevant.
+    """
+    from scipy import linalg
+
+    mu_a, mu_b = np.atleast_1d(mu_a), np.atleast_1d(mu_b)
+    sigma_a, sigma_b = np.atleast_2d(sigma_a), np.atleast_2d(sigma_b)
+
+    covmean, _ = linalg.sqrtm(sigma_a.dot(sigma_b), disp=False)
+    if not np.isfinite(covmean).all():
+        # Singular product covariance, nudge the diagonal as in the reference
+        # implementations rather than returning a nan
+        offset = np.eye(sigma_a.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma_a + offset).dot(sigma_b + offset))
+    if np.iscomplexobj(covmean):
+        covmean = covmean.real
+
+    diff = mu_a - mu_b
+    return float(diff.dot(diff) + np.trace(sigma_a) + np.trace(sigma_b) - 2 * np.trace(covmean))
+
+
+def _gaussian_stats(activations):
+    activations = np.asarray(activations, dtype=np.float64)
+    return activations.mean(axis=0), np.cov(activations, rowvar=False)
+
+
+def get_fid_metric():
+    """FID between the generated batch and the real batch, lower is better.
+
+    Per-batch FID is noisy at typical validation batch sizes and is only
+    meaningful as a relative trend across checkpoints, not as a headline number
+    comparable to published FID-50k.
+    """
+    model, params = _get_inception()
+
+    @jax.jit
+    def activations(images):
+        # Inception wants [-1, 1] at 299x299; pool3 output is [B, 1, 1, 2048]
+        resized = jax.image.resize(images, (images.shape[0], 299, 299, 3), method='bilinear')
+        features = model.apply(params, resized, train=False)
+        return features.reshape(features.shape[0], -1)
+
+    def fid_metric(generated: jnp.ndarray, batch):
+        original = (jnp.asarray(batch['image'], dtype=jnp.float32) - 127.5) / 127.5
+        mu_gen, sigma_gen = _gaussian_stats(activations(generated))
+        mu_real, sigma_real = _gaussian_stats(activations(original))
+        return frechet_distance(mu_gen, sigma_gen, mu_real, sigma_real)
+
+    return EvaluationMetric(function=fid_metric, name='fid')
diff --git a/flaxdiff/models/autoencoder/autoencoder.py b/flaxdiff/models/autoencoder/autoencoder.py
index a307da1..233eaec 100644
--- a/flaxdiff/models/autoencoder/autoencoder.py
+++ b/flaxdiff/models/autoencoder/autoencoder.py
@@ -10,11 +10,19 @@
 @dataclass
 class AutoEncoder(ABC):
     """Base class for autoencoder models with video support.
-    
+
     This class defines the interface for autoencoders and provides
     video handling functionality, allowing child classes to focus
     on implementing the core encoding/decoding for individual frames.
+
+    Latents are normalized as (z - latent_shift) * latent_scale on the way out
+    and inverted on the way in, the SD3 convention. The defaults are the
+    identity; set them to the dataset's own latent mean and 1/std so the
+    diffusion model sees roughly unit-variance, zero-mean inputs.
     """
+    latent_shift: float = 0.0
+    latent_scale: float = 1.0
+
     @abstractmethod
     def __encode__(self, x: jnp.ndarray, **kwargs) -> jnp.ndarray:
         """Abstract method for encoding a batch of images.
@@ -72,13 +80,15 @@ def encode(self, x: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwa
             
             # Encode all frames
             latent = self.__encode__(x_reshaped, key=key, **kwargs)
-            
+
             # Reshape back to include temporal dimension [B, T, h, w, c]
             latent_shape = latent.shape
-            return latent.reshape(batch_size, seq_len, *latent_shape[1:])
+            latent = latent.reshape(batch_size, seq_len, *latent_shape[1:])
         else:
             # Standard image processing
-            return self.__encode__(x, key=key, **kwargs)
+            latent = self.__encode__(x, key=key, **kwargs)
+
+        return (latent - self.latent_shift) * self.latent_scale
     
     def decode(self, z: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwargs) -> jnp.ndarray:
         """Decode latent representations, with special handling for video data.
@@ -95,9 +105,11 @@ def decode(self, z: jnp.ndarray, key: Optional[jax.random.PRNGKey] = None, **kwa
         Returns:
             Decoded output with the same batch and temporal dimensions as input
         """
+        z = z / self.latent_scale + self.latent_shift
+
         # Check for video data (5D tensor)
         is_video = len(z.shape) == 5
-        
+
         if is_video:
             # Extract dimensions for reshaping
             batch_size, seq_len, height, width, channels = z.shape
diff --git a/flaxdiff/models/autoencoder/diffusers.py b/flaxdiff/models/autoencoder/diffusers.py
index 43ee03c..ab8015b 100644
--- a/flaxdiff/models/autoencoder/diffusers.py
+++ b/flaxdiff/models/autoencoder/diffusers.py
@@ -12,7 +12,8 @@
 """
 
 class StableDiffusionVAE(AutoEncoder):
-    def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16):
+    def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16", dtype=jnp.bfloat16,
+                 latent_shift=None, latent_scale=None):
 
         pretrained = load_pretrained_vae(modelname)
         config = pretrained["config"]
@@ -61,10 +62,14 @@ def __init__(self, modelname = "CompVis/stable-diffusion-v1-4", revision="bf16",
             dtype=dtype,
         )
 
-        # Older VAE configs predate the scaling_factor key; 0.18215 is the SD default
-        scaling_factor = config.get("scaling_factor", 0.18215)
-        print(f"Scaling factor: {scaling_factor}")
-        
+        # The VAE's own latent normalization rides on the AutoEncoder seam, so a
+        # caller can override it with per-dataset statistics without a second
+        # scaling path. Older configs predate these keys; 0.0 and 0.18215 are
+        # the SD defaults.
+        self.latent_shift = config.get("shift_factor", 0.0) if latent_shift is None else latent_shift
+        self.latent_scale = config.get("scaling_factor", 0.18215) if latent_scale is None else latent_scale
+        print(f"Latent shift: {self.latent_shift}, latent scale: {self.latent_scale}")
+
         def encode_single_frame(images, rngkey: jax.random.PRNGKey = None):
             latents = enc.apply({"params": params['encoder']}, images, deterministic=True)
             latents = quant_conv.apply({"params": params['quant_conv']}, latents)
@@ -75,11 +80,9 @@ def encode_single_frame(images, rngkey: jax.random.PRNGKey = None):
                 latents = mean + std * jax.random.normal(rngkey, mean.shape, dtype=mean.dtype)
             else:
                 latents, _ = jnp.split(latents, 2, axis=-1)
-            latents *= scaling_factor
             return latents
-        
+
         def decode_single_frame(latents):
-            latents = (1.0 / scaling_factor) * latents
             latents = post_quant_conv.apply({"params": params['post_quant_conv']}, latents)
             return dec.apply({"params": params['decoder']}, latents)
         
diff --git a/flaxdiff/models/dit_common.py b/flaxdiff/models/dit_common.py
index d18903d..923a8d2 100644
--- a/flaxdiff/models/dit_common.py
+++ b/flaxdiff/models/dit_common.py
@@ -14,7 +14,7 @@
 import jax
 import jax.numpy as jnp
 from flax import linen as nn
-from typing import Optional
+from typing import Optional, Sequence
 from flax.typing import Dtype, PrecisionLike
 
 from .common import FourierEmbedding, TimeProjection
@@ -38,6 +38,28 @@ def scan_indices(scan_order: str, H_P: int, W_P: int):
     return None
 
 
+def scan_ordered_pos_embed(emb_dim: int, H_P: int, W_P: int, scan_order: str):
+    """2D sincos position embedding permuted into the scan order, so token i of
+    the sequence carries the signal for the 2D position it actually came from."""
+    pos_embed = build_2d_sincos_pos_embed(emb_dim, H_P, W_P)
+    idx = scan_indices(scan_order, H_P, W_P)
+    return pos_embed if idx is None else pos_embed[idx]
+
+
+def build_block_pattern(num_layers: int, ssm_attention_ratio: str = "3:1",
+                        block_pattern: Optional[Sequence[str]] = None):
+    """Per-layer mixer choice from a ratio string like '3:1', 'all-ssm', 'all-attn'."""
+    if block_pattern is not None:
+        return list(block_pattern)
+    if ssm_attention_ratio == "all-ssm":
+        return ['ssm'] * num_layers
+    if ssm_attention_ratio == "all-attn":
+        return ['attn'] * num_layers
+    n_ssm, n_attn = (int(part) for part in ssm_attention_ratio.split(':'))
+    unit = ['ssm'] * n_ssm + ['attn'] * n_attn
+    return (unit * (num_layers // len(unit) + 1))[:num_layers]
+
+
 class PatchSequenceEmbed(nn.Module):
     """Patchify in raster/hilbert/zigzag order and add the 2D sincos signal.
 
@@ -85,15 +107,10 @@ def __call__(self, x):
         else:
             tokens = self.patch_embed(x)
 
-        # Fixed 2D sincos position embedding (order-invariant spatial signal).
-        # For hilbert/zigzag, reorder the row-major embedding to match the scan
-        # so each token gets the sincos for its real 2D position.
-        pos_embed = build_2d_sincos_pos_embed(self.emb_features, H_P, W_P)
-        pos_embed = jnp.asarray(pos_embed, dtype=tokens.dtype)
-        idx = scan_indices(self.scan_order, H_P, W_P)
-        if idx is not None:
-            pos_embed = pos_embed[idx]
-        tokens = tokens + pos_embed[None, :, :]
+        # Fixed 2D sincos position embedding (order-invariant spatial signal)
+        pos_embed = scan_ordered_pos_embed(
+            self.emb_features, H_P, W_P, self.scan_order)
+        tokens = tokens + jnp.asarray(pos_embed, dtype=tokens.dtype)[None, :, :]
         return tokens, inv_idx
 
 
@@ -203,11 +220,16 @@ class ModulatedBlock(nn.Module):
     mixer='ssm' replaces attention with a bidirectional S5 scan, optionally
     followed by Spatial-Mamba style 2D state fusion. freqs_cis is unused by
     the SSM mixer but kept in the interface so blocks are interchangeable.
+
+    modulated=False drops the adaLN-Zero conditioning path entirely, leaving a
+    plain pre-norm residual block with learned affine norms - the ViT block a
+    JEPA encoder needs, where there is no timestep to condition on.
     """
     features: int
     num_heads: int
     rope_emb: RotaryEmbedding = None
     mixer: str = 'attention'
+    modulated: bool = True
     mlp_ratio: int = 4
     dropout_rate: float = 0.0
     dtype: Optional[Dtype] = None
@@ -227,13 +249,17 @@ def setup(self):
         assert self.mixer in ('attention', 'ssm'), f"Unknown mixer {self.mixer}"
         hidden_features = int(self.features * self.mlp_ratio)
 
-        self.ada_params_module = AdaLNParams(
-            self.features, dtype=self.dtype, precision=self.precision)
+        if self.modulated:
+            self.ada_params_module = AdaLNParams(
+                self.features, dtype=self.dtype, precision=self.precision)
+        # Without modulation the norms carry their own affine, since there is
+        # no conditioning vector left to supply the shift and scale
+        affine = not self.modulated
         self.norm1 = nn.LayerNorm(
-            epsilon=self.norm_epsilon, use_scale=False, use_bias=False,
+            epsilon=self.norm_epsilon, use_scale=affine, use_bias=affine,
             dtype=self.dtype, name="norm1")
         self.norm2 = nn.LayerNorm(
-            epsilon=self.norm_epsilon, use_scale=False, use_bias=False,
+            epsilon=self.norm_epsilon, use_scale=affine, use_bias=affine,
             dtype=self.dtype, name="norm2")
 
         if self.mixer == 'attention':
@@ -307,9 +333,16 @@ def _apply_2d_fusion(self, ssm_output):
 
     @nn.compact
     def __call__(self, x, conditioning, freqs_cis, train: bool = False):
-        scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split(
-            self.ada_params_module(conditioning), 6, axis=-1
-        )
+        if self.modulated:
+            scale_mlp, shift_mlp, gate_mlp, scale_attn, shift_attn, gate_attn = jnp.split(
+                self.ada_params_module(conditioning), 6, axis=-1
+            )
+        else:
+            assert conditioning is None, "an unmodulated block takes no conditioning"
+            # Identity modulation - the block body below collapses to a plain
+            # pre-norm residual block without branching on the mode
+            scale_mlp = shift_mlp = scale_attn = shift_attn = 0.0
+            gate_mlp = gate_attn = 1.0
 
         # --- Mixer path (attention or SSM) ---
         residual = x
@@ -340,12 +373,17 @@ def __call__(self, x, conditioning, freqs_cis, train: bool = False):
         return x
 
 
+def identity_rope_freqs(seq_len: int, dim: int, dtype: Dtype = jnp.float32):
+    """RoPE frequencies that rotate by nothing, for sequences whose 1D index is
+    not a position (scan orders other than raster, masked token subsets)."""
+    shape = (seq_len, dim // 2)
+    return jnp.ones(shape, dtype=dtype), jnp.zeros(shape, dtype=dtype)
+
+
 def neutralized_rope_freqs(rope: RotaryEmbedding, seq_len: int, scan_order: str):
     """RoPE frequencies for the sequence, neutralized to identity for
     hilbert/zigzag scans where the 1D index is not a 2D position (the 2D
     sincos embedding already carries position there)."""
-    freqs_cos, freqs_sin = rope(seq_len=seq_len)
     if scan_order != 'raster':
-        freqs_cos = jnp.ones_like(freqs_cos)
-        freqs_sin = jnp.zeros_like(freqs_sin)
-    return freqs_cos, freqs_sin
+        return identity_rope_freqs(seq_len, rope.dim)
+    return rope(seq_len=seq_len)
diff --git a/flaxdiff/models/registry.py b/flaxdiff/models/registry.py
index 2923b9a..c691712 100644
--- a/flaxdiff/models/registry.py
+++ b/flaxdiff/models/registry.py
@@ -18,6 +18,7 @@
 from .ssm_dit import HybridSSMAttentionDiT
 from .video_dit import VideoDiT
 from .unet_3d import UNet3D
+from ..jepa.models import JepaEncoder, JepaVideoEncoder, JepaPredictor
 
 MODEL_REGISTRY = {
     'unet': Unet,
@@ -29,6 +30,9 @@
     'hybrid_dit': HybridSSMAttentionDiT,
     'video_dit': VideoDiT,
     'unet_3d': UNet3D,
+    'jepa_encoder': JepaEncoder,
+    'jepa_video_encoder': JepaVideoEncoder,
+    'jepa_predictor': JepaPredictor,
 }
 
 ARCHITECTURE_SUFFIX_FLAGS = {
diff --git a/flaxdiff/models/ssm_dit.py b/flaxdiff/models/ssm_dit.py
index 63f3033..5e9f559 100644
--- a/flaxdiff/models/ssm_dit.py
+++ b/flaxdiff/models/ssm_dit.py
@@ -11,7 +11,7 @@
 
 from .dit_common import (
     PatchSequenceEmbed, ConditioningEmbed, PatchSequenceOutput,
-    ModulatedBlock, remat_block, neutralized_rope_freqs,
+    ModulatedBlock, remat_block, neutralized_rope_freqs, build_block_pattern,
 )
 from .vit_common import RotaryEmbedding
 
@@ -49,21 +49,6 @@ def scan_order(self):
             "use_hilbert and use_zigzag are mutually exclusive"
         return 'hilbert' if self.use_hilbert else 'zigzag' if self.use_zigzag else 'raster'
 
-    def _build_block_pattern(self):
-        """Generate block pattern from ratio string."""
-        if self.block_pattern is not None:
-            pattern = list(self.block_pattern)
-        elif self.ssm_attention_ratio == "all-ssm":
-            pattern = ['ssm'] * self.num_layers
-        elif self.ssm_attention_ratio == "all-attn":
-            pattern = ['attn'] * self.num_layers
-        else:
-            parts = self.ssm_attention_ratio.split(':')
-            n_ssm, n_attn = int(parts[0]), int(parts[1])
-            unit = ['ssm'] * n_ssm + ['attn'] * n_attn
-            pattern = (unit * (self.num_layers // len(unit) + 1))[:self.num_layers]
-        return pattern
-
     def setup(self):
         self.embed = PatchSequenceEmbed(
             patch_size=self.patch_size,
@@ -82,7 +67,8 @@ def setup(self):
             dim=self.emb_features // self.num_heads,
             max_seq_len=4096, dtype=self.dtype)
 
-        pattern = self._build_block_pattern()
+        pattern = build_block_pattern(
+            self.num_layers, self.ssm_attention_ratio, self.block_pattern)
         blocks = []
         for i, block_type in enumerate(pattern):
             if block_type == 'ssm':
diff --git a/flaxdiff/predictors/__init__.py b/flaxdiff/predictors/__init__.py
index 68390f2..befce3b 100644
--- a/flaxdiff/predictors/__init__.py
+++ b/flaxdiff/predictors/__init__.py
@@ -7,6 +7,7 @@
     CosineNoiseScheduler,
     KarrasVENoiseScheduler,
     EDMNoiseScheduler,
+    FlowMatchingScheduler,
 )
 
 ############################################################################################################
@@ -39,6 +40,14 @@ def get_target(self, x_0, epsilon, rates) ->jnp.ndarray:
     def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray]) -> jnp.ndarray:
         return 1
 
+    def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray:
+        """||target error||^2 / ||x_0 error||^2 at the given SNR.
+
+        Loss weights (min-SNR-gamma and friends) are defined on the x_0 loss;
+        dividing by this converts them into the space the model is trained in.
+        """
+        return 1.0
+
 class EpsilonPredictionTransform(DiffusionPredictionTransform):
     def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
         # preds is the predicted noise
@@ -46,10 +55,13 @@ def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray])
         signal_rates, noise_rates = rates
         x_0 = (x_t - epsilon * noise_rates) / signal_rates
         return x_0, epsilon
-    
+
     def get_target(self, x_0, epsilon, rates) ->jnp.ndarray:
         return epsilon
 
+    def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray:
+        return snr
+
 class DirectPredictionTransform(DiffusionPredictionTransform):
     def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
         # Here the model predicts x_0 directly
@@ -76,7 +88,27 @@ def get_target(self, x_0, epsilon, rates) ->jnp.ndarray:
         v = signal_rate * epsilon - noise_rate * x_0
         variance = signal_rate**2 + noise_rate**2
         return v / jnp.sqrt(variance)
-    
+
+    def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray:
+        return snr + 1
+
+class FlowMatchPredictionTransform(DiffusionPredictionTransform):
+    """Rectified flow velocity: the model predicts u = epsilon - x_0, the
+    constant velocity of the linear path, so both endpoints are one step away.
+    """
+    def backward_diffusion(self, x_t, preds, rates: tuple[jnp.ndarray, jnp.ndarray]) -> Union[jnp.ndarray, jnp.ndarray]:
+        signal_rate, noise_rate = rates
+        x_0 = x_t - noise_rate * preds
+        epsilon = x_t + signal_rate * preds
+        return x_0, epsilon
+
+    def get_target(self, x_0, epsilon, rates) ->jnp.ndarray:
+        return epsilon - x_0
+
+    def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray:
+        # x_0 error is t times the velocity error, and t = 1 / (1 + sqrt(SNR))
+        return (1 + jnp.sqrt(snr)) ** 2
+
 class KarrasPredictionTransform(DiffusionPredictionTransform):
     def __init__(self, sigma_data=0.5) -> None:
         super().__init__()
@@ -102,6 +134,11 @@ def get_input_scale(self, rates: tuple[jnp.ndarray, jnp.ndarray], epsilon=1e-8)
         c_in = 1 / (jnp.sqrt(self.sigma_data ** 2 + sigma ** 2) + epsilon)
         return c_in
 
+    def target_error_scale(self, snr: jnp.ndarray) -> jnp.ndarray:
+        # x_0 error is c_out times the raw error, and alpha = 1 here so
+        # sigma^2 = 1 / SNR
+        return 1 / self.sigma_data ** 2 + snr
+
 ############################################################################################################
 # Noise schedule presets
 ############################################################################################################
@@ -113,6 +150,10 @@ def get_diffusion_preset(
     sigma_max: float = 80.0,
     rho: float = 7.0,
     sigma_data: float = 0.5,
+    P_mean: float = -0.4,
+    P_std: float = 1.0,
+    shift: float = 1.0,
+    min_snr_gamma: float = None,
 ) -> tuple[NoiseScheduler, NoiseScheduler, DiffusionPredictionTransform]:
     """Named (train schedule, sampling schedule, prediction transform) presets.
 
@@ -120,18 +161,29 @@ def get_diffusion_preset(
     parameterization. Both training and inference build from here, so a model
     is always sampled with the same convention it was trained with.
     """
+    # Only the training schedule ever produces loss weights, so only it needs
+    # to know the parameterization
     if name == 'edm':
-        train = EDMNoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
-        sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
         transform = KarrasPredictionTransform(sigma_data=sigma_data)
-    elif name == 'karras':
-        train = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
+        train = EDMNoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data,
+                                  P_mean=P_mean, P_std=P_std,
+                                  prediction_transform=transform, min_snr_gamma=min_snr_gamma)
         sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
+    elif name == 'karras':
         transform = KarrasPredictionTransform(sigma_data=sigma_data)
+        train = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data,
+                                       prediction_transform=transform, min_snr_gamma=min_snr_gamma)
+        sample = KarrasVENoiseScheduler(1, sigma_min=sigma_min, sigma_max=sigma_max, rho=rho, sigma_data=sigma_data)
     elif name == 'cosine':
-        train = CosineNoiseScheduler(timesteps, beta_end=1)
-        sample = CosineNoiseScheduler(timesteps, beta_end=1)
         transform = VPredictionTransform()
+        train = CosineNoiseScheduler(timesteps, beta_end=1,
+                                     prediction_transform=transform, min_snr_gamma=min_snr_gamma)
+        sample = CosineNoiseScheduler(timesteps, beta_end=1)
+    elif name in ('flow', 'flow_matching'):
+        transform = FlowMatchPredictionTransform()
+        train = FlowMatchingScheduler(shift=shift,
+                                      prediction_transform=transform, min_snr_gamma=min_snr_gamma)
+        sample = FlowMatchingScheduler(shift=shift)
     else:
         raise ValueError(f"Unknown noise schedule preset: {name}")
     return train, sample, transform
diff --git a/flaxdiff/samplers/common.py b/flaxdiff/samplers/common.py
index 910202c..7caab39 100644
--- a/flaxdiff/samplers/common.py
+++ b/flaxdiff/samplers/common.py
@@ -22,31 +22,49 @@ def __init__(
         model_output_transform: DiffusionPredictionTransform,
         input_config: DiffusionInputConfig,
         guidance_scale: float = 0.0,
+        guidance_start: float = 0.0,
+        guidance_stop: float = 1.0,
         autoencoder: AutoEncoder = None,
     ):
         """Initialize the diffusion sampler.
-        
+
         Args:
             model: Neural network model
             params: Model parameters
             noise_schedule: Noise scheduler
             model_output_transform: Transform for model predictions
             guidance_scale: Scale for classifier-free guidance (0.0 means disabled)
+            guidance_start: Fraction of the trajectory after which guidance turns on
+            guidance_stop: Fraction of the trajectory after which guidance turns off
             autoencoder: Optional autoencoder for latent diffusion
         """
         self.model = model
         self.noise_schedule = noise_schedule
         self.model_output_transform = model_output_transform
         self.guidance_scale = guidance_scale
+        self.guidance_start = guidance_start
+        self.guidance_stop = guidance_stop
         self.autoencoder = autoencoder
         self.input_config = input_config
-        
+
         unconditionals = input_config.get_unconditionals()
-        
+
         if self.guidance_scale > 0:
             # Classifier free guidance
             print("Using classifier-free guidance")
 
+            def guidance_at(t):
+                """Interval-limited guidance (Kynkaanniemi et al. 2024).
+
+                Guidance hurts at high noise and buys nothing at low noise, so
+                outside [guidance_start, guidance_stop] the scale drops to 1,
+                which is exactly the plain conditional prediction. Progress runs
+                from 0 at pure noise to 1 at the clean sample.
+                """
+                progress = 1.0 - t / self.noise_schedule.max_timesteps
+                inside = (progress >= guidance_start) & (progress <= guidance_stop)
+                return jnp.where(inside, guidance_scale, 1.0)
+
             def sample_model(params, x_t, t, *conditioning_inputs):
                 # Concatenate unconditional and conditional inputs
                 x_t_cat = jnp.concatenate([x_t] * 2, axis=0)
@@ -71,7 +89,8 @@ def sample_model(params, x_t, t, *conditioning_inputs):
                 
                 # Split model output into unconditional and conditional parts
                 model_output_cond, model_output_uncond = jnp.split(model_output, 2, axis=0)
-                model_output = model_output_uncond + guidance_scale * (model_output_cond - model_output_uncond)
+                scale = guidance_at(t).reshape(get_coeff_shapes_tuple(model_output_cond))
+                model_output = model_output_uncond + scale * (model_output_cond - model_output_uncond)
 
                 x_0, eps = self.model_output_transform(x_t, model_output, t, self.noise_schedule)
                 return x_0, eps, model_output
diff --git a/flaxdiff/schedulers/__init__.py b/flaxdiff/schedulers/__init__.py
index 32594e4..278bdb9 100644
--- a/flaxdiff/schedulers/__init__.py
+++ b/flaxdiff/schedulers/__init__.py
@@ -3,4 +3,5 @@
 from .cosine import *
 from .linear import *
 from .sqrt import *
-from .karras import *
\ No newline at end of file
+from .karras import *
+from .flow import *
diff --git a/flaxdiff/schedulers/common.py b/flaxdiff/schedulers/common.py
index 3061542..e4c2459 100644
--- a/flaxdiff/schedulers/common.py
+++ b/flaxdiff/schedulers/common.py
@@ -18,11 +18,17 @@ def __init__(self, timesteps,
                     dtype=jnp.float32,
                     clip_min=-1.0,
                     clip_max=1.0,
+                    min_snr_gamma: float = None,
+                    prediction_transform=None,
                     *args, **kwargs):
         self.max_timesteps = timesteps
         self.dtype = dtype
         self.clip_min = clip_min
         self.clip_max = clip_max
+        if min_snr_gamma is not None and prediction_transform is None:
+            raise ValueError("min_snr_gamma needs the prediction transform it will be trained against")
+        self.min_snr_gamma = min_snr_gamma
+        self.prediction_transform = prediction_transform
         if type(timesteps) == int and timesteps > 1:
             timestep_generator = lambda rng, batch_size, max_timesteps = timesteps: jax.random.randint(rng, (batch_size,), 0, max_timesteps)
         else:
@@ -33,10 +39,28 @@ def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.n
         state, rng = state.get_random_key()
         timesteps = self.timestep_generator(rng, batch_size, self.max_timesteps)
         return timesteps, state
-    
+
+    def get_snr(self, steps) -> jnp.ndarray:
+        signal_rates, noise_rates = self.get_rates(steps, shape=(-1,))
+        return (signal_rates / noise_rates) ** 2
+
     def get_weights(self, steps, shape=(-1, 1, 1, 1)):
+        """Per-sample loss weight, in the target space of the paired parameterization.
+
+        min-SNR-gamma (Hang et al. 2023) is defined as min(SNR, gamma) on the
+        x_0 loss; the transform converts that into whichever space the trainer
+        actually computes the loss in. It replaces the schedule's own weighting
+        rather than stacking on top of it.
+        """
+        if self.min_snr_gamma is None:
+            return self.get_schedule_weights(steps, shape)
+        snr = self.get_snr(steps)
+        weights = jnp.minimum(snr, self.min_snr_gamma) / self.prediction_transform.target_error_scale(snr)
+        return weights.reshape(shape)
+
+    def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)):
         raise NotImplementedError
-    
+
     def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
         raise NotImplementedError
     
@@ -77,7 +101,7 @@ def __init__(self, timesteps, sigma_min=0.002, sigma_max=80.0, sigma_data=1, *ar
         self.sigma_max = sigma_max
         self.sigma_data = sigma_data
     
-    def get_weights(self, steps, shape=(-1, 1, 1, 1)):
+    def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)):
         sigma = self.get_sigmas(steps)
         return (1 + (1 / (1 + ((1 - sigma ** 2)/(sigma ** 2)))) / (self.sigma_max ** 2)).reshape(shape)
     
diff --git a/flaxdiff/schedulers/cosine.py b/flaxdiff/schedulers/cosine.py
index 00de7dd..e1f0bd9 100644
--- a/flaxdiff/schedulers/cosine.py
+++ b/flaxdiff/schedulers/cosine.py
@@ -34,7 +34,7 @@ def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarra
         noise_rates = jnp.sin((jnp.pi * steps) / (2 * self.max_timesteps))
         return reshape_rates((signal_rates, noise_rates), shape=shape)
     
-    def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray:
+    def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray:
         alpha, sigma = self.get_rates(steps, shape=shape)
         return 1 / (1 + (alpha ** 2 / sigma ** 2))
     
\ No newline at end of file
diff --git a/flaxdiff/schedulers/discrete.py b/flaxdiff/schedulers/discrete.py
index ee1a01c..a7af264 100644
--- a/flaxdiff/schedulers/discrete.py
+++ b/flaxdiff/schedulers/discrete.py
@@ -47,7 +47,7 @@ def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.n
     def get_p2_weights(self, k, gamma):
         return (k + self.alpha_cumprod / (1 - self.alpha_cumprod)) ** -gamma
     
-    def get_weights(self, steps, shape=(-1, 1, 1, 1)):
+    def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)):
         steps = jnp.int16(steps)
         return self.p2_loss_weights[steps].reshape(shape)
 
diff --git a/flaxdiff/schedulers/flow.py b/flaxdiff/schedulers/flow.py
new file mode 100644
index 0000000..7909fd2
--- /dev/null
+++ b/flaxdiff/schedulers/flow.py
@@ -0,0 +1,55 @@
+import math
+import jax
+import jax.numpy as jnp
+from ..utils import RandomMarkovState
+from .continuous import ContinuousNoiseScheduler
+from .common import reshape_rates
+
+
+def compute_resolution_shift(sequence_length, base_seq_len=256, max_seq_len=4096,
+                             base_shift=0.5, max_shift=1.15) -> float:
+    """Flux-style resolution dependent timestep shift.
+
+    Longer token sequences carry more redundancy, so the trajectory has to
+    spend more of its budget at high noise for the global structure to settle.
+    mu is interpolated linearly in sequence length and the shift is exp(mu).
+    """
+    slope = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    mu = base_shift + slope * (sequence_length - base_seq_len)
+    return math.exp(mu)
+
+
+class FlowMatchingScheduler(ContinuousNoiseScheduler):
+    """Rectified flow / conditional flow matching on the linear path.
+
+    x_t = (1 - t) * x_0 + t * epsilon for t in [0, 1], so alpha + sigma = 1 and
+    the model input needs no scaling. Timesteps are drawn logit-normal as in
+    SD3, which concentrates training on the middle of the trajectory where the
+    velocity is hardest to predict.
+    """
+    def __init__(self, shift: float = 1.0, logit_mean: float = 0.0, logit_std: float = 1.0,
+                 *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.shift = shift
+        self.logit_mean = logit_mean
+        self.logit_std = logit_std
+
+    def shift_timesteps(self, steps) -> jnp.ndarray:
+        return self.shift * steps / (1 + (self.shift - 1) * steps)
+
+    def generate_timesteps(self, batch_size, state: RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]:
+        state, rng = state.get_random_key()
+        normal = jax.random.normal(rng, (batch_size,), dtype=self.dtype)
+        return jax.nn.sigmoid(normal * self.logit_std + self.logit_mean), state
+
+    def get_rates(self, steps, shape=(-1, 1, 1, 1)) -> tuple[jnp.ndarray, jnp.ndarray]:
+        t = self.shift_timesteps(jnp.asarray(steps, dtype=self.dtype))
+        return reshape_rates((1 - t, t), shape=shape)
+
+    def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray:
+        return jnp.ones_like(jnp.asarray(steps, dtype=self.dtype)).reshape(shape)
+
+    def transform_inputs(self, x, steps) -> tuple[jnp.ndarray, jnp.ndarray]:
+        # The Fourier time embedding is tuned for discrete-style timesteps, so
+        # the [0, 1] flow time is rescaled into that range
+        return x, self.shift_timesteps(jnp.asarray(steps, dtype=self.dtype)) * 1000
diff --git a/flaxdiff/schedulers/karras.py b/flaxdiff/schedulers/karras.py
index b9e09d5..078165e 100644
--- a/flaxdiff/schedulers/karras.py
+++ b/flaxdiff/schedulers/karras.py
@@ -17,7 +17,7 @@ def get_sigmas(self, steps) -> jnp.ndarray:
         sigmas = (self.max_inv_rho + ramp * (self.min_inv_rho - self.max_inv_rho)) ** self.rho
         return sigmas
         
-    def get_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray:
+    def get_schedule_weights(self, steps, shape=(-1, 1, 1, 1)) -> jnp.ndarray:
         sigma = self.get_sigmas(steps)
         # EDM lambda(sigma) = (sigma^2 + sd^2) / (sigma * sd)^2, written in a
         # form that needs no epsilon guard (the old guard halved the weight at
@@ -64,13 +64,21 @@ def get_sigmas(self, steps) -> jnp.ndarray:
         return self.sigmas[steps]
 
 class EDMNoiseScheduler(KarrasVENoiseScheduler):
-    def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5, *args, **kwargs):
+    """Training sigmas drawn from exp(N(P_mean, P_std^2)).
+
+    Defaults are EDM2's (Karras et al. 2024); EDM1's -1.2/1.2 concentrated too
+    much mass on low noise levels for larger models. Pass them explicitly to
+    reproduce an EDM1 run.
+    """
+    def __init__(self, timesteps, sigma_min=0.002, sigma_max=80, rho=7., sigma_data=0.5,
+                 P_mean=-0.4, P_std=1.0, *args, **kwargs):
         super().__init__(timesteps=timesteps, sigma_min=sigma_min, sigma_max=sigma_max, sigma_data=sigma_data, *args, **kwargs)
+        self.P_mean = P_mean
+        self.P_std = P_std
 
-    def get_sigmas(self, steps, std=1.2, mean=-1.2) -> jnp.ndarray:
+    def get_sigmas(self, steps) -> jnp.ndarray:
         space = steps / self.max_timesteps
-        # space = jax.scipy.special.erfinv(self.erf_sigma_min + steps * (self.erf_sigma_max - self.erf_sigma_min))
-        return jnp.exp(space * std + mean)
+        return jnp.exp(space * self.P_std + self.P_mean)
     
     def generate_timesteps(self, batch_size, state:RandomMarkovState) -> tuple[jnp.ndarray, RandomMarkovState]:
         state, rng = state.get_random_key()
diff --git a/flaxdiff/trainer/__init__.py b/flaxdiff/trainer/__init__.py
index 871700a..cc7bf1c 100644
--- a/flaxdiff/trainer/__init__.py
+++ b/flaxdiff/trainer/__init__.py
@@ -1,2 +1,3 @@
 from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
+from .objectives import Objective, EMASpec, DiffusionObjective
 from .general_diffusion_trainer import GeneralDiffusionTrainer, TrainState, ConditionalInputConfig
diff --git a/flaxdiff/trainer/general_diffusion_trainer.py b/flaxdiff/trainer/general_diffusion_trainer.py
index ece09d0..9a48d74 100644
--- a/flaxdiff/trainer/general_diffusion_trainer.py
+++ b/flaxdiff/trainer/general_diffusion_trainer.py
@@ -1,81 +1,59 @@
 import json
+import numpy as np
 import flax
 from flax import linen as nn
 import jax
-from typing import Callable, List, Dict, Tuple, Union, Any, Sequence, Type, Optional
+from typing import Callable, List, Dict, Tuple, Union, Any, Sequence, Optional
 from dataclasses import field, dataclass
 import jax.numpy as jnp
 import optax
 import functools
-from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple
+from ..schedulers import NoiseScheduler
 from ..predictors import DiffusionPredictionTransform, EpsilonPredictionTransform
-from ..samplers.common import DiffusionSampler
-from ..samplers.ddim import DDIMSampler
 
-from flaxdiff.utils import RandomMarkovState, serialize_model, get_latest_checkpoint
+from flaxdiff.utils import RandomMarkovState, serialize_model, get_latest_checkpoint, shard_batch
 from flaxdiff.inputs import ConditioningEncoder, ConditionalInputConfig, DiffusionInputConfig
 
-from flaxdiff.utils import shard_batch
-
 from .simple_trainer import SimpleTrainer, SimpleTrainState, Metrics
+from .objectives import Objective, DiffusionObjective
 
 from flaxdiff.models.autoencoder.autoencoder import AutoEncoder
 from flax.training import dynamic_scale as dynamic_scale_lib
 import shutil
 
 
+def _replace_subtree(tree, path, value):
+    if not path:
+        return value
+    return {**tree, path[0]: _replace_subtree(tree[path[0]], path[1:], value)}
+
+
+def _subtree(tree, path):
+    for key in path:
+        tree = tree[key]
+    return tree
+
+
 class TrainState(SimpleTrainState):
     rngs: jax.random.PRNGKey
     ema_params: dict
 
-    def apply_ema(self, decay: float = 0.999):
-        new_ema_params = jax.tree_util.tree_map(
+    def apply_ema(self, decay: float = 0.999, path: Tuple[str, ...] = ()):
+        """EMA over a subtree of the parameters (the whole tree by default).
+
+        JEPA's target encoder is the EMA of the context encoder alone, so the
+        predictor's slice of the tree must be left out of the average.
+        """
+        new_subtree = jax.tree_util.tree_map(
             lambda ema, param: decay * ema + (1 - decay) * param,
-            self.ema_params,
-            self.params,
+            _subtree(self.ema_params, path),
+            _subtree(self.params, path),
         )
-        return self.replace(ema_params=new_ema_params)
+        return self.replace(ema_params=_replace_subtree(self.ema_params, path, new_subtree))
 
 
 from flaxdiff.metrics.common import EvaluationMetric
 
-def generate_modelname(
-    dataset_name: str,
-    noise_schedule_name: str,
-    architecture_name: str,
-    model: nn.Module,
-    input_config: DiffusionInputConfig,
-    autoencoder: AutoEncoder = None,
-    frames_per_sample: int = None,
-) -> str:
-    """
-    Generate a model name based on the configuration.
-    
-    Args:
-        config: Configuration dictionary.
-        
-    Returns:
-        A string representing the model name.
-    """
-    import hashlib
-    import json
-    
-    # Extract key components for the name
-    
-    model_name = f"diffusion-{dataset_name}-res{input_config.sample_data_shape[-2]}"
-    
-    # model_name = f"diffusion-{dataset_name}-res{input_config.sample_data_shape[-2]}-{noise_schedule_name}-{architecture_name}"
-    
-    # if autoencoder is not None:
-    #     model_name += f"-vae"
-        
-    # if frames_per_sample is not None:
-    #     model_name += f"-frames_{frames_per_sample}"
-    
-    # model_name += f"-{'.'.join([cond.encoder.key for cond in input_config.conditions])}"
-    
-    return model_name
-
 class GeneralDiffusionTrainer(SimpleTrainer):
     """
     General trainer for diffusion models supporting both images and videos.
@@ -89,45 +67,66 @@ class GeneralDiffusionTrainer(SimpleTrainer):
     def __init__(self,
                  model: nn.Module,
                  optimizer: optax.GradientTransformation,
-                 noise_schedule: NoiseScheduler,
                  input_config: DiffusionInputConfig,
                  rngs: jax.random.PRNGKey,
+                 noise_schedule: NoiseScheduler = None,
+                 objective: Objective = None,
                  unconditional_prob: float = 0.12,
                  name: str = "GeneralDiffusion",
                  model_output_transform: DiffusionPredictionTransform = EpsilonPredictionTransform(),
                  autoencoder: AutoEncoder = None,
                  native_resolution: int = None,
-                 frames_per_sample: int = None,
                  wandb_config: Dict[str, Any] = None,
                  eval_metrics: List[EvaluationMetric] = None,
                  best_tracker_metric: str = "train/best_loss",
                  ema_decay: float = 0.999,
+                 loss_fn: Callable = optax.l2_loss,
                  **kwargs
                  ):
         """
         Initialize the general diffusion trainer.
-        
+
         Args:
             model: Neural network model
             optimizer: Optimization algorithm
-            noise_schedule: Noise scheduler for diffusion process
             input_config: Configuration for input data, including keys, shapes and conditioning inputs
             rngs: Random number generator keys
+            noise_schedule: Noise scheduler for the default diffusion objective
+            objective: What to optimize. Defaults to diffusion over `model`.
             unconditional_prob: Probability of training with unconditional samples
             name: Name of this trainer
             model_output_transform: Transform for model predictions
             autoencoder: Optional autoencoder for latent diffusion
             native_resolution: Native resolution of the data
-            frames_per_sample: Number of frames per video sample (for video only)
             **kwargs: Additional arguments for parent class
         """
         # Initialize with parent DiffusionTrainer but without encoder parameter
         input_shapes = input_config.get_input_shapes(
             autoencoder=autoencoder,
         )
-        self.input_config = input_config
         self.eval_metrics = eval_metrics
-        
+
+        if native_resolution is None:
+            sample_shape = input_config.sample_data_shape
+            native_resolution = sample_shape[-2]
+            if autoencoder is not None:
+                native_resolution = native_resolution * autoencoder.downscale_factor
+
+        if objective is None:
+            objective = DiffusionObjective(
+                model=model,
+                noise_schedule=noise_schedule,
+                model_output_transform=model_output_transform,
+                input_config=input_config,
+                input_shapes=input_shapes,
+                autoencoder=autoencoder,
+                unconditional_prob=unconditional_prob,
+                loss_fn=loss_fn,
+                ema_decay=ema_decay,
+                native_resolution=native_resolution,
+            )
+        self.objective = objective
+
         if wandb_config is not None:
             # If input_config is not in wandb_config, add it
             if 'input_config' not in wandb_config['config']:
@@ -138,33 +137,14 @@ def __init__(self,
             if 'autoencoder' not in wandb_config['config'] and autoencoder is not None:
                 wandb_config['config']['autoencoder'] = autoencoder.name
                 wandb_config['config']['autoencoder_opts'] = json.dumps(autoencoder.serialize())
-                
-            # Generate a model name based on the configuration
-            modelname = generate_modelname(
-                dataset_name=wandb_config['config']['arguments']['dataset'],
-                noise_schedule_name=wandb_config['config']['arguments']['noise_schedule'],
-                architecture_name=wandb_config['config']['arguments']['architecture'],
-                model=model,
-                input_config=input_config,
-                autoencoder=autoencoder,
-                frames_per_sample=frames_per_sample,
-            )
-            print("Model name:", modelname)
-            self.modelname = modelname
-            wandb_config['config']['modelname'] = modelname
-        
-        self.noise_schedule = noise_schedule
-        self.model_output_transform = model_output_transform
-        self.unconditional_prob = unconditional_prob
-        self.autoencoder = autoencoder
-        self.ema_decay = ema_decay
 
-        if native_resolution is None:
-            sample_shape = input_config.sample_data_shape
-            native_resolution = sample_shape[-2]
-            if autoencoder is not None:
-                native_resolution = native_resolution * autoencoder.downscale_factor
-        self.native_resolution = native_resolution
+            # Names the wandb model artifact, so runs of different objectives on
+            # the same dataset and resolution do not collide
+            dataset_name = wandb_config['config']['arguments']['dataset']
+            self.modelname = (
+                f"{objective.tag}-{dataset_name}-res{input_config.sample_data_shape[-2]}")
+            print("Model name:", self.modelname)
+            wandb_config['config']['modelname'] = self.modelname
 
         super().__init__(
             model=model,
@@ -173,18 +153,11 @@ def __init__(self,
             rngs=rngs,
             name=name,
             wandb_config=wandb_config,
+            loss_fn=loss_fn,
             **kwargs
         )
         
         self.best_tracker_metric = best_tracker_metric
-        
-        # Store video-specific parameters
-        self.frames_per_sample = frames_per_sample
-        
-        # List of conditional inputs
-        self.conditional_inputs = input_config.conditions
-        # Determine if we're working with video or images
-        self.is_video = self._is_video_data()
         self.best_val_metrics = {}
 
         # val/<metric_name> -> True if higher is better (default lower-is-better)
@@ -204,7 +177,9 @@ def generate_states(
 
         def init_fn():
             next_rngs, subkey = jax.random.split(rngs)
-            params = model.init(subkey, **self.get_input_ones())
+            # The objective owns the parameter tree; the sharding constraints
+            # come from its abstract shapes, so it never has to describe them.
+            params = self.objective.init_params(subkey)
             return TrainState.create(
                 apply_fn=model.apply,
                 params=params,
@@ -217,118 +192,43 @@ def init_fn():
 
         return self._build_state(init_fn)
 
-    def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None):
+    def fit(self, data, training_steps_per_epoch, epochs, val_steps_per_epoch=8, **validation_step_args):
         local_batch_size = data['local_batch_size']
-        validation_step_args = {
-            "sampler_class": sampler_class,
-            "sampling_noise_schedule": sampling_noise_schedule,
-        }
         return super().fit(
-            data, 
-            train_steps_per_epoch=training_steps_per_epoch, 
-            epochs=epochs, 
-            train_step_args={"batch_size": local_batch_size}, 
+            data,
+            train_steps_per_epoch=training_steps_per_epoch,
+            epochs=epochs,
+            train_step_args={"batch_size": local_batch_size},
             val_steps_per_epoch=val_steps_per_epoch,
             validation_step_args=validation_step_args,
         )
 
-    def _is_video_data(self):
-        sample_data_shape = self.input_config.sample_data_shape
-        return len(sample_data_shape) == 5
-        
     def _define_train_step(self, batch_size):
         """
-        Define the training step function for both image and video diffusion.
-        Optimized for efficient sharding and JIT compilation.
+        Define the training step: the objective supplies the loss, the trainer
+        supplies gradients, sharding, EMA and the state update.
         """
-        # Access class variables once for JIT optimization
-        noise_schedule = self.noise_schedule
-        model = self.model
-        model_output_transform = self.model_output_transform
-        loss_fn = self.loss_fn
-        autoencoder = self.autoencoder
-        unconditional_prob = self.unconditional_prob
-        
-        input_config = self.input_config
-        sample_data_key = input_config.sample_data_key
-        
-        # JIT-optimized function for processing conditional inputs
-        # @functools.partial(jax.jit, static_argnums=(2,))
-        def process_conditioning(batch, uncond_mask):
-            return input_config.process_conditioning(
-                batch,
-                uncond_mask=uncond_mask,
-            )
+        objective = self.objective
+        ema = objective.ema
 
-        # Main training step function - optimized for JIT compilation and sharding
         def train_step(train_state: TrainState, rng_state: RandomMarkovState, batch):
             """Training step over the global batch; GSPMD partitions it."""
             # One key per step: threefry is partitionable, so every device draws
             # its own slice of the same stream without folding in a device index.
             rng_state, step_key = rng_state.get_random_key()
-            local_rng_state = RandomMarkovState(step_key)
 
-            # Extract and normalize data (works for both images and videos)
-            data = batch[sample_data_key]
-            local_batch_size = data.shape[0]
-            data = (jnp.asarray(data, dtype=jnp.float32) - 127.5) / 127.5
-            
-            # Autoencoder step (handles both image and video data)
-            if autoencoder is not None:
-                local_rng_state, enc_key = local_rng_state.get_random_key()
-                data = autoencoder.encode(data, enc_key)
-            
-            # Determine number of unconditional samples per mini batch randomly
-            local_rng_state, uncond_key = local_rng_state.get_random_key()
-            # Determine unconditional samples 
-            uncond_mask = jax.random.bernoulli(
-                uncond_key,
-                shape=(local_batch_size,),
-                p=unconditional_prob
-            )
-            
-            # Process conditioning
-            all_conditional_inputs = process_conditioning(batch, uncond_mask)
-            
-            # Generate diffusion timesteps
-            noise_level, local_rng_state = noise_schedule.generate_timesteps(local_batch_size, local_rng_state)
-            
-            # Generate noise
-            local_rng_state, noise_key = local_rng_state.get_random_key()
-            noise = jax.random.normal(noise_key, shape=data.shape, dtype=jnp.float32)
+            def objective_loss(params):
+                return objective.loss(params, train_state.ema_params, batch,
+                                      step_key, train_state.step)
 
-            local_rng_state, dropout_key = local_rng_state.get_random_key()
-            
-            # Forward diffusion process
-            rates = noise_schedule.get_rates(noise_level, get_coeff_shapes_tuple(data))
-            noisy_data, c_in, expected_output = model_output_transform.forward_diffusion(data, noise, rates)
-
-            # Loss function
-            def model_loss(params):
-                # Apply model
-                inputs = noise_schedule.transform_inputs(noisy_data * c_in, noise_level)
-                preds = model.apply(
-                    params, *inputs, *all_conditional_inputs,
-                    train=True, rngs={'dropout': dropout_key},
-                )
-                
-                # Transform predictions and calculate loss
-                preds = model_output_transform.pred_transform(noisy_data, preds, rates)
-                sample_losses = loss_fn(preds, expected_output)
-                
-                # Apply loss weighting
-                weights = noise_schedule.get_weights(noise_level, get_coeff_shapes_tuple(sample_losses))
-                weighted_loss = sample_losses * weights
-                
-                return jnp.mean(weighted_loss)
-            
             # Compute gradients and apply updates. The loss is a mean over the
             # batch-sharded axis, so its gradient carries the cross-device
             # all-reduce on its own - no hand-written pmean.
             if train_state.dynamic_scale is not None:
                 # Mixed precision training with dynamic scale
-                grad_fn = train_state.dynamic_scale.value_and_grad(model_loss)
-                dynamic_scale, grads_finite, loss, grads = grad_fn(train_state.params)
+                grad_fn = train_state.dynamic_scale.value_and_grad(
+                    objective_loss, has_aux=True)
+                dynamic_scale, grads_finite, (loss, aux), grads = grad_fn(train_state.params)
 
                 train_state = train_state.replace(dynamic_scale=dynamic_scale)
                 new_state = train_state.apply_gradients(grads=grads)
@@ -340,90 +240,27 @@ def model_loss(params):
                     params=jax.tree.map(select_fn, new_state.params, train_state.params)
                 )
             else:
-                grad_fn = jax.value_and_grad(model_loss)
-                loss, grads = grad_fn(train_state.params)
+                grad_fn = jax.value_and_grad(objective_loss, has_aux=True)
+                (loss, aux), grads = grad_fn(train_state.params)
                 new_state = train_state.apply_gradients(grads=grads)
 
-            # Apply EMA update
-            new_state = new_state.apply_ema(self.ema_decay)
+            # The EMA copy is sharded like the params it tracks, so averaging a
+            # subtree stays a local read-modify-write on every device.
+            new_state = new_state.apply_ema(ema.decay(train_state.step), ema.path)
 
-            return new_state, loss, rng_state, jnp.isfinite(loss)
+            return new_state, loss, aux, rng_state, jnp.isfinite(loss)
 
         replicated = self.replicated
         return jax.jit(
             train_step,
             in_shardings=(self.state_sharding, replicated, self.batch_sharding),
-            out_shardings=(self.state_sharding, replicated, replicated, replicated),
+            out_shardings=(self.state_sharding, replicated, replicated, replicated,
+                           replicated),
             donate_argnums=(0,),
         )
 
-    def _define_validation_step(self, sampler_class: Type[DiffusionSampler]=DDIMSampler, sampling_noise_schedule: NoiseScheduler=None):
-        """
-        Define the validation step for both image and video diffusion models.
-        """
-        # Setup for validation
-        model = self.model
-        autoencoder = self.autoencoder
-        input_config = self.input_config
-        conditional_inputs = self.conditional_inputs
-        is_video = self.is_video
-        
-        # Get necessary parameters
-        image_size = self._get_image_size()
-        
-        # Get sequence length only for video data
-        sequence_length = self._get_sequence_length() if is_video else None
-        
-        # Initialize the sampler
-        sampler = sampler_class(
-            model=model,
-            noise_schedule=self.noise_schedule if sampling_noise_schedule is None else sampling_noise_schedule,
-            model_output_transform=self.model_output_transform,
-            input_config=input_config,
-            autoencoder=autoencoder,
-            guidance_scale=3.0,
-        )
-        
-        def generate_samples(
-            val_state: TrainState,
-            batch,
-            diffusion_steps: int,
-        ):
-            # Process all conditional inputs
-            model_conditioning_inputs = [cond_input(batch) for cond_input in conditional_inputs]
-            
-            # Determine batch size
-            batch_size = len(model_conditioning_inputs[0]) if model_conditioning_inputs else 4
-            
-            # Generate samples - works for both images and videos
-            return sampler.generate_samples(
-                params=val_state.ema_params,
-                resolution=image_size,
-                num_samples=batch_size,
-                sequence_length=sequence_length,  # Will be None for images
-                diffusion_steps=diffusion_steps,
-                end_step=0,
-                priors=None,
-                model_conditioning_inputs=tuple(model_conditioning_inputs),
-            )
-        
-        return generate_samples
-        
-    def _get_image_size(self):
-        """Helper to determine image size from available information."""
-        if self.native_resolution is not None:
-            return self.native_resolution
-            
-        sample_data_shape = self.input_config.sample_data_shape
-        return sample_data_shape[-2] # Assuming [..., H, W, C] format
-    
-    def _get_sequence_length(self):
-        """Helper to determine sequence length for video generation."""
-        if not self.is_video:
-            return None
-            
-        sample_data_shape = self.input_config.sample_data_shape
-        return sample_data_shape[1]  # Assuming [B,T,H,W,C] format
+    def _define_validation_step(self, **kwargs):
+        return self.objective.make_validation_step(**kwargs)
 
     def validation_loop(
         self,
@@ -432,16 +269,15 @@ def validation_loop(
         val_ds,
         val_steps_per_epoch,
         current_step,
-        diffusion_steps=200,
     ):
         """
-        Run validation and log samples for both image and video diffusion.
+        Score the objective's validation artifacts and let it visualize them.
         """
         process_index = jax.process_index()
-        generate_samples = val_step_fn
-        
+
         val_ds = iter(val_ds()) if val_ds else None
-        print(f"Validation loop started for process index {process_index} with {global_device_count} devices.")
+        print(f"Validation loop started for process index {process_index} "
+              f"with {jax.device_count()} devices.")
         # Evaluation step
         try:
             metrics = {metric.name: [] for metric in self.eval_metrics} if self.eval_metrics else {}
@@ -450,18 +286,13 @@ def validation_loop(
                     batch = None
                 else:
                     batch = shard_batch(self.batch_sharding, next(val_ds))
-                # Generate samples
-                samples = generate_samples(
-                    val_state,
-                    batch,
-                    diffusion_steps,
-                )
-                
+                artifacts = val_step_fn(val_state, batch)
+
                 if self.eval_metrics is not None:
                     for metric in self.eval_metrics:
                         try:
                             # Evaluate metrics
-                            metric_val = metric.function(samples, batch)
+                            metric_val = metric.function(artifacts, batch)
                             metrics[metric.name].append(metric_val)
                         except Exception as e:
                             print("Error in evaluation metrics:", e)
@@ -471,16 +302,10 @@ def validation_loop(
                     
                 if i == 0:
                     print(f"Evaluation started for process index {process_index}")
-                    # Log samples to wandb
                     if self.wandb is not None and self.wandb:
-                        import numpy as np
-                        
-                        # Process samples differently based on dimensionality
-                        if len(samples.shape) == 5:  # [B,T,H,W,C] - Video data
-                            self._log_video_samples(samples, current_step)
-                        else:  # [B,H,W,C] - Image data
-                            self._log_image_samples(samples, current_step)
-                    
+                        self.objective.log_validation_artifacts(
+                            self.wandb, artifacts, current_step)
+
             # Flatten the metrics
             if metrics:
                 metrics = {k: np.mean(v) for k, v in metrics.items()}
@@ -520,46 +345,7 @@ def validation_loop(
             import traceback
             traceback.print_exc()
 
-            
-    def _log_video_samples(self, samples, current_step):
-        """Helper to log video samples to wandb."""
-        import numpy as np
-        from wandb import Video as wandbVideo
-        
-        for i in range(samples.shape[0]):
-            # Convert to numpy, denormalize and clip
-            sample = np.array(samples[i])
-            sample = (sample + 1) * 127.5
-            sample = np.clip(sample, 0, 255).astype(np.uint8)
-            
-            # Log as video
-            self.wandb.log({
-                f"video_sample_{i}": wandbVideo(
-                    sample, 
-                    fps=10, 
-                    caption=f"Video Sample {i} at step {current_step}"
-                )
-            }, step=current_step)
-            
-    def _log_image_samples(self, samples, current_step):
-        """Helper to log image samples to wandb."""
-        import numpy as np
-        from wandb import Image as wandbImage
-        
-        for i in range(samples.shape[0]):
-            # Convert to numpy, denormalize and clip
-            sample = np.array(samples[i])
-            sample = (sample + 1) * 127.5
-            sample = np.clip(sample, 0, 255).astype(np.uint8)
-            
-            # Log as image
-            self.wandb.log({
-                f"sample_{i}": wandbImage(
-                    sample, 
-                    caption=f"Sample {i} at step {current_step}"
-                )
-            }, step=current_step)
-            
+
     def push_to_registry(
         self,
         registry_name: str = 'wandb-registry-model',
diff --git a/flaxdiff/trainer/objectives.py b/flaxdiff/trainer/objectives.py
new file mode 100644
index 0000000..a6a508e
--- /dev/null
+++ b/flaxdiff/trainer/objectives.py
@@ -0,0 +1,200 @@
+"""What the trainer is optimizing.
+
+The trainer owns training mechanics - sharding, EMA bookkeeping, checkpoints,
+logging, the loops. An Objective owns what is being learned: the parameters it
+holds, the loss it computes from a batch, the telemetry that loss reports, and
+the artifacts validation scores. Swapping the objective swaps the research
+question without touching any of the mechanics.
+
+The loss always returns auxiliary metrics alongside the scalar, so an objective
+with several loss terms or with per-step diagnostics (JEPA's collapse
+telemetry) can surface them without the trainer knowing what they mean.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, Tuple, Type
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+
+from ..schedulers import NoiseScheduler, get_coeff_shapes_tuple
+from ..predictors import DiffusionPredictionTransform
+from ..samplers.common import DiffusionSampler
+from ..samplers.ddim import DDIMSampler
+from ..models.autoencoder.autoencoder import AutoEncoder
+from ..inputs import DiffusionInputConfig
+from ..utils import RandomMarkovState
+
+
+@dataclass
+class EMASpec:
+    """Which slice of the parameters the EMA copy tracks, and how fast.
+
+    decay is a step-indexed schedule rather than a constant because momentum
+    ramps are load-bearing for some objectives (I-JEPA anneals 0.996 -> 1.0).
+    path selects a subtree; the empty path means the whole parameter tree.
+    """
+    decay: Callable[[Any], Any]
+    path: Tuple[str, ...] = ()
+
+
+class Objective(ABC):
+    """The learning problem: parameters, loss, EMA policy, validation artifacts."""
+
+    tag: str = "objective"  # names the checkpoint artifact this run publishes
+    ema: EMASpec
+
+    @abstractmethod
+    def init_params(self, rng: jax.Array) -> Any:
+        """Build the full parameter tree, which may hold several sub-modules."""
+
+    @abstractmethod
+    def loss(self, params: Any, ema_params: Any, batch: Any, rng: jax.Array,
+             step: Any) -> Tuple[jax.Array, Dict[str, jax.Array]]:
+        """Scalar loss plus auxiliary metrics, differentiated w.r.t. params."""
+
+    @abstractmethod
+    def make_validation_step(self, **kwargs) -> Callable[[Any, Any], Any]:
+        """Build (val_state, batch) -> artifacts, which eval metrics score."""
+
+    def log_validation_artifacts(self, wandb, artifacts, step: int):
+        """Visualize the artifacts. Nothing to draw unless an objective says so."""
+
+
+class DiffusionObjective(Objective):
+    """Denoising diffusion: sample a noise level, corrupt, predict, weight."""
+
+    tag = "diffusion"
+
+    def __init__(
+        self,
+        model,
+        noise_schedule: NoiseScheduler,
+        model_output_transform: DiffusionPredictionTransform,
+        input_config: DiffusionInputConfig,
+        input_shapes: Dict[str, Tuple[int, ...]],
+        autoencoder: AutoEncoder = None,
+        unconditional_prob: float = 0.12,
+        loss_fn: Callable = optax.l2_loss,
+        ema_decay: float = 0.999,
+        native_resolution: int = None,
+        diffusion_steps: int = 200,
+    ):
+        self.model = model
+        self.noise_schedule = noise_schedule
+        self.model_output_transform = model_output_transform
+        self.input_config = input_config
+        self.input_shapes = input_shapes
+        self.autoencoder = autoencoder
+        self.unconditional_prob = unconditional_prob
+        self.loss_fn = loss_fn
+        self.native_resolution = native_resolution
+        self.diffusion_steps = diffusion_steps
+        self.ema = EMASpec(decay=optax.constant_schedule(ema_decay))
+
+    def init_params(self, rng):
+        input_ones = {k: jnp.ones((1, *v)) for k, v in self.input_shapes.items()}
+        return self.model.init(rng, **input_ones)
+
+    def loss(self, params, ema_params, batch, rng, step):
+        local_rng_state = RandomMarkovState(rng)
+
+        # Extract and normalize data (works for both images and videos)
+        data = batch[self.input_config.sample_data_key]
+        local_batch_size = data.shape[0]
+        data = (jnp.asarray(data, dtype=jnp.float32) - 127.5) / 127.5
+
+        if self.autoencoder is not None:
+            local_rng_state, enc_key = local_rng_state.get_random_key()
+            data = self.autoencoder.encode(data, enc_key)
+
+        local_rng_state, uncond_key = local_rng_state.get_random_key()
+        uncond_mask = jax.random.bernoulli(
+            uncond_key, shape=(local_batch_size,), p=self.unconditional_prob)
+        all_conditional_inputs = self.input_config.process_conditioning(
+            batch, uncond_mask=uncond_mask)
+
+        noise_level, local_rng_state = self.noise_schedule.generate_timesteps(
+            local_batch_size, local_rng_state)
+
+        local_rng_state, noise_key = local_rng_state.get_random_key()
+        noise = jax.random.normal(noise_key, shape=data.shape, dtype=jnp.float32)
+
+        local_rng_state, dropout_key = local_rng_state.get_random_key()
+
+        rates = self.noise_schedule.get_rates(noise_level, get_coeff_shapes_tuple(data))
+        noisy_data, c_in, expected_output = self.model_output_transform.forward_diffusion(
+            data, noise, rates)
+
+        inputs = self.noise_schedule.transform_inputs(noisy_data * c_in, noise_level)
+        preds = self.model.apply(
+            params, *inputs, *all_conditional_inputs,
+            train=True, rngs={'dropout': dropout_key},
+        )
+
+        preds = self.model_output_transform.pred_transform(noisy_data, preds, rates)
+        sample_losses = self.loss_fn(preds, expected_output)
+
+        weights = self.noise_schedule.get_weights(
+            noise_level, get_coeff_shapes_tuple(sample_losses))
+        return jnp.mean(sample_losses * weights), {}
+
+    def make_validation_step(
+        self,
+        sampler_class: Type[DiffusionSampler] = DDIMSampler,
+        sampling_noise_schedule: NoiseScheduler = None,
+    ):
+        sampler = sampler_class(
+            model=self.model,
+            noise_schedule=self.noise_schedule if sampling_noise_schedule is None else sampling_noise_schedule,
+            model_output_transform=self.model_output_transform,
+            input_config=self.input_config,
+            autoencoder=self.autoencoder,
+            guidance_scale=3.0,
+        )
+        conditional_inputs = self.input_config.conditions
+        image_size = self._image_size()
+        sequence_length = self._sequence_length()
+
+        def generate_samples(val_state, batch):
+            model_conditioning_inputs = [cond_input(batch) for cond_input in conditional_inputs]
+            batch_size = len(model_conditioning_inputs[0]) if model_conditioning_inputs else 4
+            return sampler.generate_samples(
+                params=val_state.ema_params,
+                resolution=image_size,
+                num_samples=batch_size,
+                sequence_length=sequence_length,  # None for images
+                diffusion_steps=self.diffusion_steps,
+                end_step=0,
+                priors=None,
+                model_conditioning_inputs=tuple(model_conditioning_inputs),
+            )
+
+        return generate_samples
+
+    def log_validation_artifacts(self, wandb, artifacts, step: int):
+        is_video = len(artifacts.shape) == 5
+        for i in range(artifacts.shape[0]):
+            sample = np.clip((np.array(artifacts[i]) + 1) * 127.5, 0, 255).astype(np.uint8)
+            if is_video:
+                from wandb import Video as wandbVideo
+                wandb.log({f"video_sample_{i}": wandbVideo(
+                    sample, fps=10, caption=f"Video Sample {i} at step {step}")}, step=step)
+            else:
+                from wandb import Image as wandbImage
+                wandb.log({f"sample_{i}": wandbImage(
+                    sample, caption=f"Sample {i} at step {step}")}, step=step)
+
+    def _is_video(self):
+        return len(self.input_config.sample_data_shape) == 4
+
+    def _image_size(self):
+        if self.native_resolution is not None:
+            return self.native_resolution
+        return self.input_config.sample_data_shape[-2]
+
+    def _sequence_length(self):
+        return self.input_config.sample_data_shape[0] if self._is_video() else None
diff --git a/flaxdiff/trainer/simple_trainer.py b/flaxdiff/trainer/simple_trainer.py
index a5ad219..9f19eb8 100644
--- a/flaxdiff/trainer/simple_trainer.py
+++ b/flaxdiff/trainer/simple_trainer.py
@@ -399,7 +399,8 @@ def train_loop(
             if i == 0 and self.profile_steps:
                 jax.profiler.start_trace(self.profile_path())
 
-            train_state, loss, rng_state, is_finite = train_step_fn(train_state, rng_state, batch)
+            train_state, loss, aux, rng_state, is_finite = train_step_fn(
+                train_state, rng_state, batch)
             # No stale alias may outlive the step: its buffers were donated.
             self.state, self.rngstate = train_state, rng_state
             self.dataset_state = getattr(train_ds, 'source_state', None)
@@ -437,6 +438,7 @@ def train_loop(
                         self.wandb.log({
                             "train/step": current_step,
                             "train/loss": loss,
+                            **{f"train/{k}": v for k, v in aux.items()},
                             **self._throughput_metrics(elapsed, steps_since_log),
                         }, step=current_step)
                     last_log_time, steps_since_log = now, 0
diff --git a/tests/test_flow_matching.py b/tests/test_flow_matching.py
new file mode 100644
index 0000000..b4efa6f
--- /dev/null
+++ b/tests/test_flow_matching.py
@@ -0,0 +1,249 @@
+"""Flow matching on the linear (rectified flow) path.
+
+Covers the schedule invariants, the exact velocity round-trip, the claim that
+the existing DDIM/Euler samplers already integrate the flow ODE, and a toy
+end-to-end run proving the objective actually learns a distribution.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import pytest
+from flax import linen as nn
+
+from flaxdiff.inputs import DiffusionInputConfig
+from flaxdiff.predictors import FlowMatchPredictionTransform, get_diffusion_preset
+from flaxdiff.samplers.ddim import DDIMSampler
+from flaxdiff.samplers.euler import EulerSampler
+from flaxdiff.schedulers import FlowMatchingScheduler
+from flaxdiff.schedulers.flow import compute_resolution_shift
+from flaxdiff.schedulers.common import get_coeff_shapes_tuple
+from flaxdiff.utils import RandomMarkovState
+
+STEPS = jnp.array([0.05, 0.3, 0.6, 0.95])
+
+
+def test_linear_path_rates():
+    schedule = FlowMatchingScheduler()
+    alpha, sigma = schedule.get_rates(STEPS, shape=(-1,))
+    assert jnp.allclose(alpha + sigma, 1.0, atol=1e-6)
+    assert jnp.allclose(sigma, STEPS, atol=1e-6)
+    # No input preconditioning on the linear path
+    assert FlowMatchPredictionTransform().get_input_scale((alpha, sigma)) == 1
+
+
+def test_endpoints_are_data_and_noise(rng):
+    schedule = FlowMatchingScheduler()
+    key0, key1 = jax.random.split(rng)
+    x0 = jax.random.normal(key0, (4, 8, 8, 3))
+    noise = jax.random.normal(key1, (4, 8, 8, 3))
+    assert jnp.allclose(schedule.add_noise(x0, noise, jnp.zeros((4,))), x0, atol=1e-6)
+    assert jnp.allclose(schedule.add_noise(x0, noise, jnp.ones((4,))), noise, atol=1e-6)
+
+
+def test_timesteps_are_logit_normal(rng):
+    schedule = FlowMatchingScheduler(logit_mean=-0.3, logit_std=1.4)
+    steps, _ = schedule.generate_timesteps(50000, RandomMarkovState(rng))
+    assert jnp.all((steps > 0) & (steps < 1))
+    logits = jnp.log(steps) - jnp.log1p(-steps)
+    assert abs(float(jnp.mean(logits)) - (-0.3)) < 0.05
+    assert abs(float(jnp.std(logits)) - 1.4) < 0.05
+
+
+def test_resolution_shift_is_identity_at_one():
+    schedule = FlowMatchingScheduler(shift=1.0)
+    assert jnp.allclose(schedule.shift_timesteps(STEPS), STEPS, atol=1e-7)
+
+
+@pytest.mark.parametrize("shift", [0.5, 1.0, 3.0])
+def test_resolution_shift_is_monotonic_and_fixes_endpoints(shift):
+    schedule = FlowMatchingScheduler(shift=shift)
+    t = jnp.linspace(0.0, 1.0, 101)
+    shifted = schedule.shift_timesteps(t)
+    assert jnp.all(jnp.diff(shifted) > 0)
+    assert float(shifted[0]) == pytest.approx(0.0, abs=1e-7)
+    assert float(shifted[-1]) == pytest.approx(1.0, abs=1e-7)
+    # A shift above 1 moves every interior timestep towards higher noise
+    if shift >= 1:
+        assert jnp.all(shifted >= t - 1e-7)
+    else:
+        assert jnp.all(shifted <= t + 1e-7)
+
+
+def test_resolution_shift_grows_with_sequence_length():
+    shifts = [compute_resolution_shift(n) for n in (256, 1024, 4096)]
+    assert shifts == sorted(shifts)
+    assert shifts[0] == pytest.approx(np.exp(0.5))
+    assert shifts[-1] == pytest.approx(np.exp(1.15))
+
+
+def test_timestep_conditioning_is_scaled_to_the_embedding_range():
+    schedule = FlowMatchingScheduler(shift=2.0)
+    x = jnp.zeros((4, 8, 8, 3))
+    _, temb = schedule.transform_inputs(x, STEPS)
+    assert jnp.allclose(temb, schedule.shift_timesteps(STEPS) * 1000)
+
+
+def test_velocity_roundtrip_is_exact(rng):
+    schedule = FlowMatchingScheduler()
+    transform = FlowMatchPredictionTransform()
+    key0, key1 = jax.random.split(rng)
+    x0 = jax.random.normal(key0, (4, 8, 8, 3))
+    noise = jax.random.normal(key1, (4, 8, 8, 3))
+    rates = schedule.get_rates(STEPS, get_coeff_shapes_tuple(x0))
+
+    xt, _, target = transform.forward_diffusion(x0, noise, rates)
+    assert jnp.allclose(target, noise - x0, atol=1e-6)
+
+    recovered_x0, recovered_noise = transform.backward_diffusion(xt, target, rates)
+    assert jnp.max(jnp.abs(recovered_x0 - x0)) < 1e-5
+    assert jnp.max(jnp.abs(recovered_noise - noise)) < 1e-5
+
+
+def test_preset_wires_flow_matching():
+    for name in ('flow', 'flow_matching'):
+        train, sample, transform = get_diffusion_preset(name, shift=2.0)
+        assert isinstance(train, FlowMatchingScheduler)
+        assert isinstance(sample, FlowMatchingScheduler)
+        assert isinstance(transform, FlowMatchPredictionTransform)
+        assert train.shift == 2.0 and sample.shift == 2.0
+
+
+############################################################################################################
+# The existing samplers already integrate the flow ODE
+############################################################################################################
+
+class ConstantVelocity(nn.Module):
+    """Stands in for a trained model; take_next_step never calls it."""
+    @nn.compact
+    def __call__(self, x, temb):
+        return x
+
+
+def _flow_sampler(sampler_class):
+    schedule = FlowMatchingScheduler()
+    return sampler_class(
+        model=ConstantVelocity(),
+        noise_schedule=schedule,
+        model_output_transform=FlowMatchPredictionTransform(),
+        input_config=DiffusionInputConfig(sample_data_key="image", sample_data_shape=(8, 8, 3), conditions=[]),
+        guidance_scale=0.0,
+    )
+
+
+@pytest.mark.parametrize("sampler_class", [EulerSampler, DDIMSampler])
+def test_sampler_step_is_the_flow_euler_step(sampler_class, rng):
+    """x_{t+dt} = x_t + u * dt exactly, for the unmodified samplers."""
+    sampler = _flow_sampler(sampler_class)
+    transform = FlowMatchPredictionTransform()
+    schedule = sampler.noise_schedule
+
+    key0, key1 = jax.random.split(rng)
+    x_t = jax.random.normal(key0, (4, 8, 8, 3))
+    velocity = jax.random.normal(key1, (4, 8, 8, 3))
+    current_step = jnp.full((4,), 0.8)
+    next_step = jnp.full((4,), 0.6)
+
+    rates = schedule.get_rates(current_step, get_coeff_shapes_tuple(x_t))
+    x0, eps = transform.backward_diffusion(x_t, velocity, rates)
+
+    stepped, _ = sampler.take_next_step(
+        current_samples=x_t,
+        reconstructed_samples=x0,
+        model_conditioning_inputs=(),
+        pred_noise=eps,
+        current_step=current_step,
+        state=RandomMarkovState(rng),
+        sample_model_fn=None,
+        next_step=next_step,
+    )
+    expected = x_t + velocity * (0.6 - 0.8)
+    assert jnp.max(jnp.abs(stepped - expected)) < 1e-5
+
+
+############################################################################################################
+# Toy end-to-end: a two-mode gaussian mixture in the plane
+############################################################################################################
+
+MODE_CENTERS = jnp.array([[-0.5, -0.5], [0.5, 0.5]])
+MODE_STD = 0.08
+
+
+def sample_mixture(key, n):
+    """Two well-separated modes in the leading two channels; the third channel
+    carries no mode information, so the sampler's fixed channel count does not
+    turn this into a harder problem."""
+    mode_key, noise_key = jax.random.split(key)
+    modes = jax.random.bernoulli(mode_key, 0.5, (n,)).astype(jnp.int32)
+    centers = jnp.concatenate([MODE_CENTERS[modes], jnp.zeros((n, 1))], axis=-1)
+    return (centers + MODE_STD * jax.random.normal(noise_key, (n, 3))).reshape(n, 1, 1, 3)
+
+
+class ToyVelocityMLP(nn.Module):
+    features: int = 128
+
+    @nn.compact
+    def __call__(self, x, temb):
+        t = jnp.reshape(temb, (-1, 1)) / 1000.0
+        freqs = jnp.arange(1, 5, dtype=jnp.float32) * jnp.pi
+        h = jnp.concatenate([x.reshape(x.shape[0], -1), t, jnp.sin(t * freqs), jnp.cos(t * freqs)], axis=-1)
+        h = nn.swish(nn.Dense(self.features)(h))
+        h = nn.swish(nn.Dense(self.features)(h))
+        return nn.Dense(3)(h).reshape(x.shape)
+
+
+def test_flow_matching_learns_a_two_mode_mixture():
+    schedule, sampling_schedule, transform = get_diffusion_preset('flow')
+    model = ToyVelocityMLP()
+    key = jax.random.PRNGKey(0)
+    params = model.init(key, jnp.zeros((1, 1, 1, 3)), jnp.zeros((1,)))
+    optimizer = optax.adam(3e-3)
+    opt_state = optimizer.init(params)
+
+    def loss_fn(params, x0, noise, steps):
+        rates = schedule.get_rates(steps, get_coeff_shapes_tuple(x0))
+        x_t, c_in, target = transform.forward_diffusion(x0, noise, rates)
+        preds = model.apply(params, *schedule.transform_inputs(x_t * c_in, steps))
+        weights = schedule.get_weights(steps, get_coeff_shapes_tuple(x0))
+        return jnp.mean(weights * (preds - target) ** 2)
+
+    @jax.jit
+    def train_step(params, opt_state, rng_state):
+        rng_state, data_key = rng_state.get_random_key()
+        rng_state, noise_key = rng_state.get_random_key()
+        x0 = sample_mixture(data_key, 512)
+        noise = jax.random.normal(noise_key, x0.shape)
+        steps, rng_state = schedule.generate_timesteps(512, rng_state)
+        loss, grads = jax.value_and_grad(loss_fn)(params, x0, noise, steps)
+        updates, opt_state = optimizer.update(grads, opt_state, params)
+        return optax.apply_updates(params, updates), opt_state, rng_state, loss
+
+    rng_state = RandomMarkovState(jax.random.PRNGKey(1))
+    for _ in range(1500):
+        params, opt_state, rng_state, loss = train_step(params, opt_state, rng_state)
+    assert float(loss) < 1.0, "flow matching loss did not come down"
+
+    sampler = EulerSampler(
+        model=model,
+        noise_schedule=sampling_schedule,
+        model_output_transform=transform,
+        input_config=DiffusionInputConfig(sample_data_key="image", sample_data_shape=(1, 1, 3), conditions=[]),
+        guidance_scale=0.0,
+    )
+    samples = sampler.generate_samples(
+        params, num_samples=2048, resolution=1, diffusion_steps=64,
+        rngstate=RandomMarkovState(jax.random.PRNGKey(2)),
+    ).reshape(-1, 3)
+
+    assignment = samples[:, 0] > 0
+    fraction = float(jnp.mean(assignment))
+    assert 0.35 < fraction < 0.65, f"modes not balanced: {fraction:.2f}"
+
+    for mode, center in enumerate(MODE_CENTERS):
+        members = samples[assignment == bool(mode)]
+        assert jnp.max(jnp.abs(jnp.mean(members[:, :2], axis=0) - center)) < 0.04
+        assert abs(float(jnp.std(members[:, :2])) - MODE_STD) < 0.04
+    # The third channel is a single zero-centred gaussian, not a mixture
+    assert abs(float(jnp.mean(samples[:, 2]))) < 0.04
+    assert abs(float(jnp.std(samples[:, 2])) - MODE_STD) < 0.04
diff --git a/tests/test_jepa.py b/tests/test_jepa.py
new file mode 100644
index 0000000..9b94d26
--- /dev/null
+++ b/tests/test_jepa.py
@@ -0,0 +1,482 @@
+"""I-JEPA and V-JEPA: masking, shapes, the objective, and collapse telemetry."""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import pytest
+
+from flaxdiff.inputs import DiffusionInputConfig
+from flaxdiff.jepa import (
+    JepaEncoder, JepaVideoEncoder, JepaObjective, multi_block_mask,
+    representation_health, normalize_targets, linear_probe, knn_probe,
+)
+from flaxdiff.jepa.models import JepaPredictor
+from flaxdiff.trainer import GeneralDiffusionTrainer
+from flaxdiff.utils import DevicePrefetchIterator
+
+RES = 32
+PATCH = 4
+GRID = (RES // PATCH, RES // PATCH)
+FRAMES = 3
+
+
+@pytest.fixture
+def mask():
+    return multi_block_mask(GRID, num_targets=4, scale=(0.15, 0.2))
+
+
+def make_encoder(**kwargs):
+    return JepaEncoder(patch_size=PATCH, emb_features=32, num_layers=2, num_heads=2,
+                       mlp_ratio=2, **kwargs)
+
+
+def make_predictor(**kwargs):
+    return JepaPredictor(grid=GRID, emb_features=32, predictor_features=16,
+                         num_layers=1, num_heads=2, mlp_ratio=2, **kwargs)
+
+
+def make_objective(mask, **kwargs):
+    return JepaObjective(make_encoder(), make_predictor(), mask,
+                         sample_data_key="image", sample_data_shape=(RES, RES, 3), **kwargs)
+
+
+def images(seed=0, batch=4):
+    return jnp.asarray(np.random.RandomState(seed).uniform(0, 255, (batch, RES, RES, 3)))
+
+
+def videos(seed=0, batch=2):
+    return jnp.asarray(
+        np.random.RandomState(seed).uniform(0, 255, (batch, FRAMES, RES, RES, 3)))
+
+
+# --- masking ---------------------------------------------------------------
+
+def test_context_and_targets_are_disjoint(mask):
+    context_idx, target_idx = mask.sample(jax.random.PRNGKey(3), 8)
+    for b in range(8):
+        context = set(np.array(context_idx[b]).tolist())
+        targets = set(np.array(target_idx[b]).reshape(-1).tolist())
+        assert not (context & targets), "context encoder can see a target patch"
+        assert max(targets) < mask.num_patches and min(context) >= 0
+
+
+def test_target_coverage_sits_inside_the_configured_scale():
+    scale = (0.15, 0.2)
+    mask = multi_block_mask(GRID, num_targets=4, scale=scale)
+    coverage = mask.block_area / mask.num_patches
+    assert scale[0] <= coverage <= scale[1]
+    for h, w in mask.block_shapes:
+        assert h * w == mask.block_area
+        assert 0.75 <= h / w <= 1.5
+
+
+def test_masks_are_reproducible_from_a_seed(mask):
+    a = mask.sample(jax.random.PRNGKey(11), 4)
+    b = mask.sample(jax.random.PRNGKey(11), 4)
+    c = mask.sample(jax.random.PRNGKey(12), 4)
+    assert all(jnp.array_equal(x, y) for x, y in zip(a, b))
+    assert not all(jnp.array_equal(x, y) for x, y in zip(a, c))
+
+
+def test_block_shapes_and_positions_actually_vary(mask):
+    blocks = np.array(mask.sample(jax.random.PRNGKey(5), 64)[1]).reshape(-1, mask.block_area)
+    corners = {int(block.min()) for block in blocks}
+    widths = {len(np.unique(block % GRID[1])) for block in blocks}
+    assert len(corners) > 1, "every block landed in the same place"
+    assert len(widths) > 1, "every block came out the same shape"
+
+
+def test_geometry_that_cannot_exist_is_rejected():
+    with pytest.raises(ValueError, match="aspect ratio"):
+        multi_block_mask((4, 4), num_targets=2, scale=(0.18, 0.19))
+    with pytest.raises(ValueError, match="no context"):
+        multi_block_mask(GRID, num_targets=6, scale=(0.15, 0.2))
+
+
+# --- forward shapes --------------------------------------------------------
+
+def test_image_encoder_shapes(mask, rng):
+    encoder = make_encoder()
+    context_idx, _ = mask.sample(rng, 4)
+    x = images()
+    params = encoder.init(rng, x, context_idx)
+
+    assert encoder.apply(params, x, context_idx).shape == (4, mask.num_context, 32)
+    assert encoder.apply(params, x).shape == (4, mask.num_patches, 32)
+
+
+def test_image_predictor_shapes(mask, rng):
+    encoder, predictor = make_encoder(), make_predictor()
+    context_idx, target_idx = mask.sample(rng, 4)
+    x = images()
+    encoder_params = encoder.init(rng, x, context_idx)
+    context = encoder.apply(encoder_params, x, context_idx)
+
+    block = target_idx[:, 0]
+    params = predictor.init(rng, context, context_idx, block)
+    assert predictor.apply(params, context, context_idx, block).shape == (4, mask.block_area, 32)
+
+
+def test_video_encoder_and_predictor_shapes(mask, rng):
+    encoder = JepaVideoEncoder(patch_size=PATCH, emb_features=32, num_layers=1,
+                               num_heads=2, mlp_ratio=2)
+    predictor = make_predictor(factorized=True)
+    context_idx, target_idx = mask.sample(rng, 2)
+    x = videos()
+
+    encoder_params = encoder.init(rng, x, context_idx)
+    context = encoder.apply(encoder_params, x, context_idx)
+    assert context.shape == (2, FRAMES, mask.num_context, 32)
+    assert encoder.apply(encoder_params, x).shape == (2, FRAMES, mask.num_patches, 32)
+
+    block = target_idx[:, 0]
+    params = predictor.init(rng, context, context_idx, block)
+    out = predictor.apply(params, context, context_idx, block)
+    assert out.shape == (2, FRAMES, mask.block_area, 32)
+
+
+def test_video_encoder_carries_a_temporal_signal(mask, rng):
+    """Tubelet masking keeps the spatial layout, so time is the only axis the
+    encoder can mix along: perturbing frame 0 must move frame 2's embedding."""
+    encoder = JepaVideoEncoder(patch_size=PATCH, emb_features=32, num_layers=1,
+                               num_heads=2, mlp_ratio=2)
+    context_idx, _ = mask.sample(rng, 2)
+    x = videos()
+    params = jax.tree.map(lambda p: p + 0.02, encoder.init(rng, x, context_idx))
+
+    before = encoder.apply(params, x, context_idx)
+    after = encoder.apply(params, x.at[:, 0].add(1.0), context_idx)
+    assert not jnp.allclose(before[:, 2], after[:, 2]), "no information flow across frames"
+
+
+def test_encoder_runs_on_the_ssm_mixer(mask, rng):
+    """The hybrid SSM encoder is the point of the shared block: same interface."""
+    encoder = make_encoder(ssm_attention_ratio="3:1", ssm_state_dim=8)
+    context_idx, _ = mask.sample(rng, 4)
+    x = images()
+    params = encoder.init(rng, x, context_idx)
+    out = encoder.apply(params, x, context_idx)
+    assert out.shape == (4, mask.num_context, 32)
+    assert jnp.all(jnp.isfinite(out))
+
+
+# --- the objective ---------------------------------------------------------
+
+def test_fresh_loss_is_non_trivial_and_training_reduces_it(mask, rng):
+    objective = make_objective(mask)
+    params = objective.init_params(rng)
+    batch = {"image": images()}
+
+    def loss_of(p):
+        return objective.loss(p, params, batch, jax.random.PRNGKey(7), 0)[0]
+
+    initial = float(loss_of(params))
+    assert initial > 0.1, "a fresh model already predicts the targets"
+
+    optimizer = optax.adam(3e-3)
+    opt_state = optimizer.init(params)
+    for _ in range(20):
+        grads = jax.grad(loss_of)(params)
+        updates, opt_state = optimizer.update(grads, opt_state, params)
+        params = optax.apply_updates(params, updates)
+
+    assert float(loss_of(params)) < initial * 0.9, "the objective does not train"
+
+
+def test_video_objective_trains(mask, rng):
+    objective = JepaObjective(
+        JepaVideoEncoder(patch_size=PATCH, emb_features=32, num_layers=1,
+                         num_heads=2, mlp_ratio=2),
+        make_predictor(factorized=True), mask,
+        sample_data_key="video", sample_data_shape=(FRAMES, RES, RES, 3))
+    params = objective.init_params(rng)
+    batch = {"video": videos()}
+
+    def loss_of(p):
+        return objective.loss(p, params, batch, jax.random.PRNGKey(7), 0)[0]
+
+    initial = float(loss_of(params))
+    optimizer = optax.adam(3e-3)
+    opt_state = optimizer.init(params)
+    for _ in range(15):
+        grads = jax.grad(loss_of)(params)
+        updates, opt_state = optimizer.update(grads, opt_state, params)
+        params = optax.apply_updates(params, updates)
+
+    assert float(loss_of(params)) < initial * 0.9
+
+
+def solid_colour_images(rs, n):
+    """One random colour per image, so a target block is fully determined by
+    anything else in the same image and by nothing in any other image."""
+    colours = rs.uniform(20, 235, (n, 1, 1, 3))
+    return jnp.asarray(np.clip(colours + rs.uniform(-15, 15, (n, RES, RES, 3)), 0, 255))
+
+
+def context_ablation(objective, params, ema, data, mask, rng):
+    """Prediction error from the true context vs. from another image's context."""
+    n, blocks = data.shape[0], mask.num_targets
+    context_idx, target_idx = mask.sample(rng, n)
+    full = normalize_targets(objective.encode(ema["params"]["context_encoder"], data))
+    targets = jnp.take_along_axis(
+        full[:, None], target_idx.reshape(n, blocks, -1, 1), axis=-2)
+    context = objective.encode(params["params"]["context_encoder"], data, context_idx)
+
+    def error(ctx):
+        predictions = objective.predictor.apply(
+            {"params": params["params"]["predictor"]},
+            jnp.repeat(ctx, blocks, axis=0),
+            jnp.repeat(context_idx, blocks, axis=0),
+            target_idx.reshape(n * blocks, -1),
+        ).reshape(targets.shape)
+        return float(jnp.mean((predictions - targets) ** 2))
+
+    return error(context), error(jnp.roll(context, 1, axis=0))
+
+
+def test_training_makes_the_prediction_depend_on_the_context(tmp_path, mask):
+    """The loss can be driven down by predicting the average target while
+    ignoring the context entirely - which would teach the encoder nothing.
+    Swapping in another image's context must cost real accuracy."""
+    rs = np.random.RandomState(0)
+    train_x, test_x = solid_colour_images(rs, 128), solid_colour_images(rs, 32)
+    normalized_test = (test_x - 127.5) / 127.5
+
+    trainer = make_jepa_trainer(tmp_path, mask, learning_rate=3e-3,
+                                momentum=(0.9, 0.99), momentum_steps=150)
+    objective = trainer.objective
+    before = context_ablation(objective, trainer.state.params, trainer.state.ema_params,
+                              normalized_test, mask, jax.random.PRNGKey(9))
+    assert before[1] / before[0] < 1.5, "a fresh predictor should not favour any context"
+
+    def batches():
+        while True:
+            yield {"image": train_x[rs.randint(0, len(train_x), 16)]}
+
+    state = trainer.fit({"train": batches, "train_len": 128, "local_batch_size": 16},
+                        training_steps_per_epoch=150, epochs=1, val_steps_per_epoch=0)
+
+    after = context_ablation(objective, state.params, state.ema_params,
+                             normalized_test, mask, jax.random.PRNGKey(9))
+    assert after[0] < before[0] / 2, "held-out prediction error did not improve"
+    assert after[1] / after[0] > 5.0, "the predictor still ignores its context"
+
+
+def test_no_gradient_reaches_the_target_branch(mask, rng):
+    """stop_gradient on the target encoder is what stops the trivial solution."""
+    objective = make_objective(mask)
+    params = objective.init_params(rng)
+    batch = {"image": images()}
+
+    grads = jax.grad(
+        lambda ema: objective.loss(params, ema, batch, jax.random.PRNGKey(7), 0)[0]
+    )(params)
+    assert all(float(jnp.max(jnp.abs(g))) == 0.0 for g in jax.tree.leaves(grads))
+
+
+# --- collapse telemetry ----------------------------------------------------
+
+def test_representation_std_detects_collapse():
+    healthy = jax.random.normal(jax.random.PRNGKey(0), (32, 16))
+    collapsed = jnp.tile(healthy[:1], (32, 1))
+
+    assert float(representation_health(healthy)["repr_std"]) > 0.5
+    assert float(representation_health(collapsed)["repr_std"]) < 1e-6
+
+
+def test_collapse_telemetry_flows_through_a_degenerate_encoder(mask, rng):
+    """A constant encoder is the failure mode; the aux dict must show it."""
+    class ConstantEncoder(JepaEncoder):
+        def __call__(self, x, token_idx=None, train: bool = False):
+            out = super().__call__(x, token_idx, train)
+            return jnp.zeros_like(out) + jnp.arange(out.shape[-1], dtype=out.dtype)
+
+    batch = {"image": images()}
+    collapsed = JepaObjective(
+        ConstantEncoder(patch_size=PATCH, emb_features=32, num_layers=2, num_heads=2,
+                        mlp_ratio=2),
+        make_predictor(), mask, "image", (RES, RES, 3))
+    collapsed_params = collapsed.init_params(rng)
+    _, collapsed_aux = collapsed.loss(
+        collapsed_params, collapsed_params, batch, jax.random.PRNGKey(7), 0)
+
+    healthy = make_objective(mask)
+    healthy_params = healthy.init_params(rng)
+    _, healthy_aux = healthy.loss(
+        healthy_params, healthy_params, batch, jax.random.PRNGKey(7), 0)
+
+    assert float(collapsed_aux["repr_std"]) < 1e-5
+    assert float(healthy_aux["repr_std"]) > 1e-3
+
+
+def test_offdiagonal_covariance_rises_with_redundant_dimensions():
+    base = jax.random.normal(jax.random.PRNGKey(0), (64, 1))
+    redundant = jnp.tile(base, (1, 8))                       # every dim identical
+    independent = jax.random.normal(jax.random.PRNGKey(1), (64, 8))
+    assert (float(representation_health(redundant)["repr_cov_offdiag"])
+            > 5 * float(representation_health(independent)["repr_cov_offdiag"]))
+
+
+# --- EMA target encoder ----------------------------------------------------
+
+def test_momentum_schedule_endpoints(mask):
+    objective = make_objective(mask, momentum=(0.996, 1.0), momentum_steps=1000)
+    assert float(objective.ema.decay(0)) == pytest.approx(0.996)
+    assert float(objective.ema.decay(1000)) == pytest.approx(1.0)
+    assert objective.ema.path == ("params", "context_encoder")
+
+
+def make_jepa_trainer(tmp_path, mask, learning_rate=1e-3, fsdp_size=1, **kwargs):
+    encoder = make_encoder()
+    return GeneralDiffusionTrainer(
+        model=encoder,
+        optimizer=optax.adam(learning_rate),
+        input_config=DiffusionInputConfig(
+            sample_data_key="image", sample_data_shape=(RES, RES, 3), conditions=[]),
+        rngs=jax.random.PRNGKey(0),
+        objective=JepaObjective(encoder, make_predictor(), mask, "image",
+                                (RES, RES, 3), **kwargs),
+        name="jepa-smoke", wandb_config=None,
+        distributed_training=fsdp_size > 1,
+        fsdp_size=fsdp_size,
+        # This encoder's parameters are far below the production shard
+        # threshold, so lower it or "FSDP on" would mean "all replicated"
+        fsdp_min_param_size=256,
+        checkpoint_base_path=str(tmp_path),
+    )
+
+
+def test_target_encoder_tracks_the_context_encoder(tmp_path, mask):
+    trainer = make_jepa_trainer(tmp_path, mask, momentum=(0.5, 0.5), momentum_steps=1)
+    # Copied to the host: the train step donates the state, so the device
+    # buffers behind this reference are gone once fit has run.
+    initial = jax.tree.map(np.asarray, trainer.state.ema_params)
+
+    def batches():
+        while True:
+            yield {"image": images()}
+
+    state = trainer.fit({"train": batches, "train_len": 16, "local_batch_size": 4},
+                        training_steps_per_epoch=3, epochs=1, val_steps_per_epoch=0)
+
+    context_moved = any(
+        not np.allclose(a, b) for a, b in zip(
+            jax.tree.leaves(state.ema_params["params"]["context_encoder"]),
+            jax.tree.leaves(initial["params"]["context_encoder"])))
+    predictor_frozen = all(
+        np.allclose(a, b) for a, b in zip(
+            jax.tree.leaves(state.ema_params["params"]["predictor"]),
+            jax.tree.leaves(initial["params"]["predictor"])))
+    assert context_moved, "the target encoder never followed the context encoder"
+    assert predictor_frozen, "EMA leaked outside the context encoder subtree"
+
+    # and it followed rather than jumped: still between where it started and now
+    ema = jax.tree.leaves(state.ema_params["params"]["context_encoder"])
+    live = jax.tree.leaves(state.params["params"]["context_encoder"])
+    assert any(not np.allclose(a, b) for a, b in zip(ema, live)), "EMA is not lagging"
+
+
+def test_jepa_trains_under_fsdp(tmp_path, mask):
+    """Where the two halves of the trainer meet: an objective that owns a
+    multi-encoder parameter tree, run through the sharded, donating train step.
+
+    Compiling is not the claim - the parameters have to be genuinely split
+    across the fsdp axis, the EMA target encoder has to follow their layout,
+    and the loss has to come back finite with its collapse telemetry intact.
+    """
+    trainer = make_jepa_trainer(tmp_path, mask, fsdp_size=2)
+    specs = [p.sharding.spec for p in jax.tree.leaves(trainer.state.params)]
+
+    sharded = [p for p in jax.tree.leaves(trainer.state.params)
+               if 'fsdp' in str(p.sharding.spec)]
+    assert sharded, "no JEPA parameter was sharded over the fsdp axis"
+    for param in sharded:
+        assert param.addressable_shards[0].data.size == param.size // 2
+
+    # The target encoder is a second copy of the same tree, so it must land on
+    # the mesh the same way rather than being gathered onto every device
+    assert specs == [p.sharding.spec for p in jax.tree.leaves(trainer.state.ema_params)]
+
+    def batches():
+        while True:
+            yield {"image": np.asarray(images(batch=jax.device_count()))}
+
+    train_step = trainer._define_train_step(batch_size=jax.device_count())
+    source = DevicePrefetchIterator(batches(), trainer.batch_sharding)
+    state, rng = trainer.state, trainer.rngstate
+    for _ in range(2):
+        state, loss, aux, rng, is_finite = train_step(state, rng, next(source))
+        assert bool(is_finite)
+        assert float(aux["repr_std"]) > 0, "collapse telemetry was lost in the step"
+
+    assert int(state.step) == 2
+    assert specs == [p.sharding.spec for p in jax.tree.leaves(state.params)]
+
+
+def test_validation_step_returns_pooled_embeddings(tmp_path, mask):
+    trainer = make_jepa_trainer(tmp_path, mask)
+    embed = trainer._define_validation_step()
+    out = embed(trainer.state, {"image": images()})
+    assert out.shape == (4, 32)
+
+
+# --- probes ----------------------------------------------------------------
+
+def separable_embeddings(num_classes=4, per_class=8, dim=6, noise=0.05):
+    rng = np.random.RandomState(0)
+    centers = rng.normal(size=(num_classes, dim)) * 3
+    labels = np.repeat(np.arange(num_classes), per_class)
+    x = centers[labels] + rng.normal(size=(len(labels), dim)) * noise
+    order = rng.permutation(len(labels))
+    return jnp.asarray(x[order], dtype=jnp.float32), jnp.asarray(labels[order])
+
+
+def test_probes_separate_clustered_embeddings():
+    x, y = separable_embeddings()
+    assert float(linear_probe(x, y, num_classes=4, steps=200)) > 0.9
+    assert float(knn_probe(x, y, num_classes=4, k=3)) > 0.9
+
+
+def test_probes_are_at_chance_on_noise():
+    x = jax.random.normal(jax.random.PRNGKey(0), (64, 6))
+    y = jnp.asarray(np.random.RandomState(1).randint(0, 4, 64))
+    assert float(knn_probe(x, y, num_classes=4, k=5)) < 0.6
+
+
+# --- registry and entrypoint ----------------------------------------------
+
+@pytest.mark.parametrize("architecture", ["jepa_encoder", "jepa_video_encoder"])
+def test_registry_builds_the_jepa_models(architecture):
+    from flaxdiff.models.registry import build_model
+    model = build_model(f"{architecture}+hilbert",
+                        {"patch_size": PATCH, "emb_features": 32, "num_layers": 1})
+    assert model.emb_features == 32 and model.scan_order == 'hilbert'
+
+
+def test_training_entrypoint_runs_end_to_end(tmp_path, monkeypatch):
+    """Registry -> mask -> objective -> trainer -> probes, as the CLI wires them."""
+    import training_jepa
+
+    monkeypatch.setenv("WANDB_MODE", "disabled")
+    classes = 4
+
+    def fake_dataset(*args, **kwargs):
+        def batches():
+            rs = np.random.RandomState(0)
+            while True:
+                yield {"image": jnp.asarray(rs.uniform(0, 255, (4, RES, RES, 3))),
+                       "label": jnp.asarray(rs.randint(0, classes, 4))}
+        return {"train": batches, "val": batches, "train_len": 16, "local_batch_size": 4}
+
+    monkeypatch.setattr(training_jepa, "get_dataset_grain", fake_dataset)
+    args = training_jepa.parser.parse_args([
+        '--image_size', str(RES), '--patch_size', str(PATCH), '--batch_size', '4',
+        '--emb_features', '32', '--num_layers', '1', '--num_heads', '2', '--mlp_ratio', '2',
+        '--predictor_features', '16', '--predictor_layers', '1', '--predictor_heads', '2',
+        '--epochs', '1', '--steps_per_epoch', '2', '--val_steps_per_epoch', '1',
+        '--probe_classes', str(classes), '--distributed_training', 'False',
+        '--ssm_attention_ratio', '3:1', '--ssm_state_dim', '8',
+        '--checkpoint_dir', str(tmp_path),
+    ])
+    training_jepa.main(args)
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
new file mode 100644
index 0000000..3801da9
--- /dev/null
+++ b/tests/test_metrics.py
@@ -0,0 +1,53 @@
+"""Evaluation metric tests.
+
+The Frechet distance itself is checked against closed forms that need no
+weights; the end-to-end InceptionV3 path downloads the FID checkpoint and is
+network-marked.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from flaxdiff.metrics.common import EvaluationMetric
+from flaxdiff.metrics.fid import frechet_distance, get_fid_metric
+
+
+def test_frechet_distance_of_a_distribution_with_itself_is_zero(rng):
+    features = np.asarray(jax.random.normal(rng, (256, 16)))
+    mu, sigma = features.mean(axis=0), np.cov(features, rowvar=False)
+    assert frechet_distance(mu, sigma, mu, sigma) == pytest.approx(0.0, abs=1e-6)
+
+
+def test_frechet_distance_of_shifted_gaussians_is_the_squared_mean_gap():
+    sigma = np.eye(8)
+    mu_a = np.zeros(8)
+    mu_b = np.full(8, 0.5)
+    assert frechet_distance(mu_a, sigma, mu_b, sigma) == pytest.approx(8 * 0.25, abs=1e-6)
+
+
+def test_frechet_distance_grows_with_covariance_mismatch():
+    mu = np.zeros(8)
+    identity = np.eye(8)
+    near = frechet_distance(mu, identity, mu, identity * 1.5)
+    far = frechet_distance(mu, identity, mu, identity * 4.0)
+    assert 0 < near < far
+
+
+@pytest.mark.network
+def test_fid_metric_scores_real_images_better_than_noise(rng):
+    metric = get_fid_metric()
+    assert isinstance(metric, EvaluationMetric)
+    assert metric.name == 'fid' and metric.higher_is_better is False
+
+    key_real, key_noise = jax.random.split(rng)
+    real = jax.random.randint(key_real, (8, 64, 64, 3), 0, 256, dtype=jnp.int32).astype(jnp.uint8)
+    batch = {'image': real}
+
+    # Generated samples live in [-1, 1]; the same images should score far
+    # closer to the batch than unrelated noise does
+    matching = metric.function((jnp.asarray(real, jnp.float32) - 127.5) / 127.5, batch)
+    unrelated = metric.function(jax.random.normal(key_noise, (8, 64, 64, 3)), batch)
+    assert np.isfinite(matching) and np.isfinite(unrelated)
+    assert matching < unrelated
diff --git a/tests/test_objectives.py b/tests/test_objectives.py
new file mode 100644
index 0000000..29d6c7e
--- /dev/null
+++ b/tests/test_objectives.py
@@ -0,0 +1,162 @@
+"""The Objective seam: the trainer owns mechanics, the objective owns the loss.
+
+The diffusion objective was lifted out of the trainer's train step, so the
+first test here is a golden fingerprint of the parameters, EMA and optimizer
+state after five real steps, captured from the implementation that inlined the
+objective. It must not move.
+
+The fingerprint is keyed to the step's RNG derivation, which draws one
+partitionable key per step rather than folding in a device index. Numbers
+captured against the older per-device folding do not carry over.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import pytest
+
+from flaxdiff.inputs import DiffusionInputConfig
+from flaxdiff.models.simple_dit import SimpleDiT
+from flaxdiff.predictors import get_diffusion_preset
+from flaxdiff.trainer import GeneralDiffusionTrainer
+from flaxdiff.trainer.general_diffusion_trainer import TrainState
+from flaxdiff.trainer.objectives import EMASpec, Objective
+
+RES = 8
+
+
+def tree_fingerprint(tree):
+    # per-leaf sums accumulated in python floats, so the golden values below
+    # do not depend on float32 reduction order
+    return sum(float(jnp.sum(leaf)) for leaf in jax.tree.leaves(tree)
+               if jnp.issubdtype(jnp.asarray(leaf).dtype, jnp.floating))
+
+
+def make_trainer(tmp_path, **kwargs):
+    train_schedule, _, transform = get_diffusion_preset("edm")
+    return GeneralDiffusionTrainer(
+        model=SimpleDiT(patch_size=4, emb_features=16, num_layers=1, num_heads=2, mlp_ratio=1),
+        optimizer=optax.adam(1e-3),
+        noise_schedule=train_schedule,
+        model_output_transform=transform,
+        input_config=DiffusionInputConfig(
+            sample_data_key="image", sample_data_shape=(RES, RES, 3), conditions=[]),
+        rngs=jax.random.PRNGKey(0),
+        name="objective-parity",
+        wandb_config=None,
+        distributed_training=False,
+        checkpoint_base_path=str(tmp_path),
+        **kwargs,
+    )
+
+
+def batch_iterator():
+    images = np.tile(np.linspace(0, 255, RES, dtype=np.float32)[None, :, None, None],
+                     (8, 1, RES, 3))
+    while True:
+        yield {"image": jnp.asarray(images)}
+
+
+def test_diffusion_objective_reproduces_the_inlined_train_step(tmp_path):
+    trainer = make_trainer(tmp_path)
+    data = {"train": batch_iterator, "train_len": 32, "local_batch_size": 8}
+    state = trainer.fit(data, training_steps_per_epoch=5, epochs=1, val_steps_per_epoch=0)
+
+    assert tree_fingerprint(state.params) == pytest.approx(8.209761425852776, abs=1e-9)
+    assert tree_fingerprint(state.ema_params) == pytest.approx(8.22020611886387, abs=1e-9)
+    assert tree_fingerprint(state.opt_state) == pytest.approx(-1.3229545587499768e-05, abs=1e-12)
+
+
+def test_train_step_returns_auxiliary_metrics(tmp_path):
+    """value_and_grad runs with has_aux so objectives can report diagnostics."""
+    trainer = make_trainer(tmp_path)
+    train_step = trainer._define_train_step(batch_size=8)
+    batch = next(batch_iterator())
+
+    new_state, loss, aux, rng_state, is_finite = train_step(
+        trainer.state, trainer.rngstate, batch)
+    assert bool(is_finite)
+    assert isinstance(aux, dict)
+
+
+class ConstantObjective(Objective):
+    """Two independent parameter subtrees, so EMA scoping is observable."""
+
+    def __init__(self, ema):
+        self.ema = ema
+
+    def init_params(self, rng):
+        return {"params": {"tracked": {"w": jnp.ones((2,))},
+                           "untracked": {"w": jnp.ones((2,))}}}
+
+    def loss(self, params, ema_params, batch, rng, step):
+        total = sum(jnp.sum(leaf ** 2) for leaf in jax.tree.leaves(params))
+        return total, {"probe": jnp.asarray(1.0)}
+
+    def make_validation_step(self, **kwargs):
+        return lambda val_state, batch: None
+
+
+def make_state(objective, params=None):
+    params = objective.init_params(jax.random.PRNGKey(0)) if params is None else params
+    return TrainState.create(
+        apply_fn=lambda *a, **k: None, params=params, ema_params=params,
+        tx=optax.sgd(0.0), rngs=jax.random.PRNGKey(0), metrics=None, dynamic_scale=None)
+
+
+def test_ema_over_a_subtree_leaves_its_siblings_alone():
+    objective = ConstantObjective(EMASpec(decay=optax.constant_schedule(0.5),
+                                          path=("params", "tracked")))
+    state = make_state(objective)
+    state = state.replace(params=jax.tree.map(lambda p: p * 3.0, state.params))
+
+    updated = state.apply_ema(0.5, objective.ema.path)
+    assert jnp.allclose(updated.ema_params["params"]["tracked"]["w"], 2.0)
+    assert jnp.allclose(updated.ema_params["params"]["untracked"]["w"], 1.0)
+
+
+def test_ema_over_the_whole_tree_is_the_default():
+    state = make_state(ConstantObjective(EMASpec(decay=optax.constant_schedule(0.5))))
+    state = state.replace(params=jax.tree.map(lambda p: p * 3.0, state.params))
+
+    updated = state.apply_ema(0.5)
+    assert jnp.allclose(updated.ema_params["params"]["tracked"]["w"], 2.0)
+    assert jnp.allclose(updated.ema_params["params"]["untracked"]["w"], 2.0)
+
+
+def test_ema_decay_follows_the_step_schedule():
+    """I-JEPA's momentum ramp: the trainer must read decay at the current step."""
+    ema = EMASpec(decay=optax.linear_schedule(0.996, 1.0, transition_steps=100))
+    assert float(ema.decay(0)) == pytest.approx(0.996)
+    assert float(ema.decay(100)) == pytest.approx(1.0)
+    assert float(ema.decay(50)) == pytest.approx(0.998)
+
+
+def test_trainer_drives_an_arbitrary_objective(tmp_path):
+    """The seam is real: a non-diffusion objective trains through the same loop."""
+    objective = ConstantObjective(EMASpec(decay=optax.constant_schedule(0.9),
+                                          path=("params", "tracked")))
+    trainer = make_trainer(tmp_path, objective=objective)
+    data = {"train": batch_iterator, "train_len": 32, "local_batch_size": 8}
+    state = trainer.fit(data, training_steps_per_epoch=3, epochs=1, val_steps_per_epoch=0)
+
+    assert int(state.step) == 3
+    # sum of squares under gradient descent must shrink
+    assert tree_fingerprint(state.params) < tree_fingerprint(objective.init_params(None))
+
+
+def test_mixed_precision_carries_the_objective_aux(tmp_path):
+    """The dynamic-scale branch differentiates the same has_aux loss, so an
+    objective's telemetry has to survive mixed precision as well."""
+    objective = ConstantObjective(EMASpec(decay=optax.constant_schedule(0.9)))
+    trainer = make_trainer(tmp_path, objective=objective, use_dynamic_scale=True)
+    train_step = trainer._define_train_step(batch_size=8)
+
+    state, loss, aux, _, is_finite = train_step(
+        trainer.state, trainer.rngstate, next(batch_iterator()))
+    # a live dynamic scale is exactly what selects that branch of the step
+    assert state.dynamic_scale is not None
+    assert bool(is_finite)
+    assert float(aux["probe"]) == 1.0
+    assert int(state.step) == 1
diff --git a/tests/test_parallelism.py b/tests/test_parallelism.py
index 931d639..989b141 100644
--- a/tests/test_parallelism.py
+++ b/tests/test_parallelism.py
@@ -61,7 +61,7 @@ def run_losses(trainer, steps):
     state, rng = trainer.state, trainer.rngstate
     losses = []
     for _ in range(steps):
-        state, loss, rng, is_finite = train_step(state, rng, next(source))
+        state, loss, _, rng, is_finite = train_step(state, rng, next(source))
         assert bool(is_finite)
         losses.append(float(loss))
     return losses
@@ -257,7 +257,7 @@ def snapshot(s):
 
     reference = snapshot(state)
     for micro in range(1, accum * 2 + 1):
-        state, _, rng, _ = train_step(state, rng, next(source))
+        state, _, _, rng, _ = train_step(state, rng, next(source))
         moved = any(not np.array_equal(a, b) for a, b in zip(reference, snapshot(state)))
         at_boundary = micro % accum == 0
         assert moved == at_boundary, f"micro-step {micro}: moved={moved}"
diff --git a/tests/test_samplers.py b/tests/test_samplers.py
index 9156c91..3fa8507 100644
--- a/tests/test_samplers.py
+++ b/tests/test_samplers.py
@@ -7,12 +7,15 @@
 shape bugs in one place, with no training involved.
 """
 
+from dataclasses import dataclass
+
 import jax
 import jax.numpy as jnp
 import pytest
 from flax import linen as nn
 
-from flaxdiff.inputs import DiffusionInputConfig
+from flaxdiff.inputs import ConditionalInputConfig, DiffusionInputConfig
+from flaxdiff.inputs.encoders import ConditioningEncoder
 from flaxdiff.predictors import EpsilonPredictionTransform, KarrasPredictionTransform
 from flaxdiff.samplers.ddim import DDIMSampler
 from flaxdiff.samplers.ddpm import DDPMSampler, SimpleDDPMSampler
@@ -186,3 +189,117 @@ def test_multistep_dpm_reentrant():
     assert jnp.allclose(first, second, atol=1e-5)
 
 
+############################################################################################################
+# Interval-limited classifier-free guidance (Kynkaanniemi et al. 2024)
+############################################################################################################
+
+@dataclass
+class LabelEncoder(ConditioningEncoder):
+    """Smallest conditioning seam that still gives CFG a real conditional and
+    unconditional branch to interpolate between."""
+    @property
+    def key(self):
+        return "label"
+
+    def tokenize(self, data):
+        return jnp.asarray(data, dtype=jnp.float32).reshape(-1, 1)
+
+    def encode_from_tokens(self, tokens):
+        return tokens
+
+    def serialize(self):
+        return {}
+
+    @staticmethod
+    def deserialize(serialized_config):
+        raise NotImplementedError
+
+
+conditional_input_config = DiffusionInputConfig(
+    sample_data_key="image",
+    sample_data_shape=(8, 8, 3),
+    conditions=[
+        ConditionalInputConfig(
+            encoder=LabelEncoder(model=None, tokenizer=None),
+            conditioning_data_key="label",
+            unconditional_input=0.0,
+        )
+    ],
+)
+
+
+class ConditionalVPOracle(nn.Module):
+    """VPOracle offset by the label, so the guided output reads back the scale
+    that was actually applied."""
+    schedule: CosineNoiseScheduler
+
+    @nn.compact
+    def __call__(self, x, temb, label):
+        alpha, sigma = self.schedule.get_rates(temb, get_coeff_shapes_tuple(x))
+        eps = x * sigma / (alpha**2 * DATA_STD**2 + sigma**2)
+        return eps + label.reshape(get_coeff_shapes_tuple(x))
+
+
+def make_guided_sampler(**kwargs):
+    schedule = CosineNoiseScheduler(1000)
+    model = ConditionalVPOracle(schedule=schedule)
+    sampler = DDIMSampler(
+        model=model,
+        noise_schedule=schedule,
+        model_output_transform=EpsilonPredictionTransform(),
+        input_config=conditional_input_config,
+        **kwargs,
+    )
+    params = model.init(
+        jax.random.PRNGKey(1), jnp.ones((1, 8, 8, 3)), jnp.ones((1,)), jnp.ones((1, 1))
+    )
+    return params, sampler
+
+
+@pytest.mark.parametrize("progress,inside", [(0.1, False), (0.5, True), (0.9, False)])
+def test_interval_cfg_applies_only_inside_the_interval(progress, inside):
+    params, full = make_guided_sampler(guidance_scale=3.0)
+    _, interval = make_guided_sampler(guidance_scale=3.0, guidance_start=0.4, guidance_stop=0.6)
+    _, unguided = make_guided_sampler(guidance_scale=0.0)
+
+    x_t = jax.random.normal(jax.random.PRNGKey(3), (4, 8, 8, 3))
+    labels = jnp.full((4, 1), 0.7)
+    t = jnp.full((4,), (1.0 - progress) * 1000)
+
+    def output(sampler):
+        return sampler.sample_model(params, x_t, t, labels)[2]
+
+    matches_full = bool(jnp.allclose(output(interval), output(full), atol=1e-5))
+    matches_unguided = bool(jnp.allclose(output(interval), output(unguided), atol=1e-5))
+    assert matches_full is inside
+    assert matches_unguided is not inside
+
+
+def test_interval_cfg_defaults_to_the_full_range():
+    params, default = make_guided_sampler(guidance_scale=3.0)
+    _, explicit = make_guided_sampler(guidance_scale=3.0, guidance_start=0.0, guidance_stop=1.0)
+
+    x_t = jax.random.normal(jax.random.PRNGKey(3), (4, 8, 8, 3))
+    labels = jnp.full((4, 1), 0.7)
+    t = jnp.full((4,), 500.0)
+    assert jnp.allclose(
+        default.sample_model(params, x_t, t, labels)[2],
+        explicit.sample_model(params, x_t, t, labels)[2],
+        atol=1e-6,
+    )
+
+
+def test_empty_guidance_interval_generates_the_unguided_samples():
+    params, empty = make_guided_sampler(guidance_scale=3.0, guidance_start=0.9, guidance_stop=0.1)
+    _, unguided = make_guided_sampler(guidance_scale=0.0)
+
+    def run(sampler):
+        return sampler.generate_samples(
+            params, num_samples=16, resolution=8, diffusion_steps=25,
+            rngstate=RandomMarkovState(jax.random.PRNGKey(2)),
+            model_conditioning_inputs=(jnp.full((16, 1), 0.7),),
+        )
+
+    assert jnp.allclose(run(empty), run(unguided), atol=1e-5)
+
+
diff --git a/tests/test_schedulers.py b/tests/test_schedulers.py
index 44e05fc..67c745e 100644
--- a/tests/test_schedulers.py
+++ b/tests/test_schedulers.py
@@ -9,6 +9,11 @@
 import jax.numpy as jnp
 import pytest
 
+from flaxdiff.predictors import (
+    EpsilonPredictionTransform,
+    VPredictionTransform,
+    get_diffusion_preset,
+)
 from flaxdiff.schedulers import (
     CosineNoiseScheduler,
     LinearNoiseSchedule,
@@ -75,12 +80,84 @@ def test_karras_weights_at_sigma_min():
     assert jnp.allclose(got, expected, rtol=1e-2)
 
 
-def test_edm_lognormal_sigma_distribution(rng):
-    """EDM training sigmas must follow exp(N(-1.2, 1.2^2)) when timesteps=1."""
+@pytest.mark.parametrize("P_mean,P_std", [(-0.4, 1.0), (-1.2, 1.2)])
+def test_edm_lognormal_sigma_distribution(rng, P_mean, P_std):
+    """EDM training sigmas follow exp(N(P_mean, P_std^2)), defaulting to EDM2."""
     from flaxdiff.utils import RandomMarkovState
 
-    schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5)
+    schedule = EDMNoiseScheduler(1, sigma_max=80, rho=7, sigma_data=0.5, P_mean=P_mean, P_std=P_std)
     steps, _ = schedule.generate_timesteps(20000, RandomMarkovState(rng))
     log_sigma = jnp.log(schedule.get_sigmas(steps))
-    assert abs(float(jnp.mean(log_sigma)) - (-1.2)) < 0.05
-    assert abs(float(jnp.std(log_sigma)) - 1.2) < 0.05
+    assert abs(float(jnp.mean(log_sigma)) - P_mean) < 0.05
+    assert abs(float(jnp.std(log_sigma)) - P_std) < 0.05
+
+
+def test_edm_defaults_to_edm2_distribution():
+    schedule = EDMNoiseScheduler(1)
+    assert (schedule.P_mean, schedule.P_std) == (-0.4, 1.0)
+
+
+############################################################################################################
+# min-SNR-gamma loss weighting (Hang et al. 2023)
+############################################################################################################
+
+MIN_SNR_STEPS = jnp.array([10, 200, 400, 600, 800, 990])
+
+
+def make_min_snr_schedule(transform, gamma):
+    return CosineNoiseScheduler(1000, min_snr_gamma=gamma, prediction_transform=transform)
+
+
+def test_min_snr_needs_a_parameterization():
+    with pytest.raises(ValueError):
+        CosineNoiseScheduler(1000, min_snr_gamma=5.0)
+
+
+def test_min_snr_weights_keep_the_requested_shape():
+    schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0)
+    assert schedule.get_weights(MIN_SNR_STEPS, shape=(-1, 1, 1, 1)).shape == (len(MIN_SNR_STEPS), 1, 1, 1)
+    assert schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)).shape == (len(MIN_SNR_STEPS),)
+
+
+def test_min_snr_epsilon_weights_match_the_paper():
+    schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0)
+    snr = schedule.get_snr(MIN_SNR_STEPS)
+    expected = jnp.minimum(snr, 5.0) / snr
+    assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), expected, rtol=1e-5)
+
+
+def test_min_snr_v_weights_match_the_paper():
+    schedule = make_min_snr_schedule(VPredictionTransform(), 5.0)
+    snr = schedule.get_snr(MIN_SNR_STEPS)
+    expected = jnp.minimum(snr, 5.0) / (snr + 1)
+    assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), expected, rtol=1e-5)
+
+
+def test_min_snr_weights_are_capped_and_non_increasing_in_snr():
+    """The whole point: high-SNR (low noise) steps stop dominating the gradient."""
+    schedule = make_min_snr_schedule(EpsilonPredictionTransform(), 5.0)
+    # ascending timesteps are descending SNR, so weights must be non-decreasing
+    weights = schedule.get_weights(jnp.arange(1, 1000, 10), shape=(-1,))
+    assert jnp.all(jnp.diff(weights) >= -1e-6)
+    assert jnp.all(weights <= 1.0 + 1e-6)
+
+
+def test_min_snr_gamma_infinity_is_the_unweighted_case():
+    schedule = make_min_snr_schedule(EpsilonPredictionTransform(), float('inf'))
+    assert jnp.allclose(schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)), 1.0, atol=1e-6)
+
+
+def test_min_snr_is_off_by_default():
+    schedule = CosineNoiseScheduler(1000)
+    assert jnp.allclose(
+        schedule.get_weights(MIN_SNR_STEPS, shape=(-1,)),
+        schedule.get_schedule_weights(MIN_SNR_STEPS, shape=(-1,)),
+    )
+
+
+@pytest.mark.parametrize("name", ['cosine', 'edm', 'karras', 'flow'])
+def test_preset_wires_min_snr_into_the_training_schedule_only(name):
+    train, sample, transform = get_diffusion_preset(name, min_snr_gamma=5.0)
+    assert train.min_snr_gamma == 5.0
+    assert train.prediction_transform is transform
+    assert sample.min_snr_gamma is None
diff --git a/tests/test_vae.py b/tests/test_vae.py
index 96904ec..3b648e0 100644
--- a/tests/test_vae.py
+++ b/tests/test_vae.py
@@ -1,7 +1,7 @@
-"""Vendored Stable Diffusion VAE tests.
+"""Vendored Stable Diffusion VAE tests, plus the latent normalization seam.
 
-Marked as network tests: they download the pretrained weights from the
-HuggingFace Hub on first run. Excluded in CI (-m "not network").
+Everything touching the pretrained weights is network-marked: they download
+from the HuggingFace Hub on first run. Excluded in CI (-m "not network").
 """
 
 import jax
@@ -9,7 +9,47 @@
 import numpy as np
 import pytest
 
-pytestmark = pytest.mark.network
+from flaxdiff.models.autoencoder import AutoEncoder
+
+
+class IdentityAutoEncoder(AutoEncoder):
+    """Latents are the input, so only the normalization seam is under test."""
+    def __encode__(self, x, **kwargs):
+        return x
+
+    def __decode__(self, z, **kwargs):
+        return z
+
+    def serialize(self):
+        return {}
+
+
+def test_latent_normalization_defaults_to_the_identity(rng):
+    autoencoder = IdentityAutoEncoder()
+    x = jax.random.normal(rng, (2, 8, 8, 4))
+    assert jnp.allclose(autoencoder.encode(x), x)
+    assert jnp.allclose(autoencoder.decode(x), x)
+
+
+@pytest.mark.parametrize("shape", [(2, 8, 8, 4), (2, 3, 8, 8, 4)])
+def test_latent_normalization_shifts_and_scales_roundtrip(rng, shape):
+    """SD3-style shift+scale: latents come out centred and rescaled, and
+    decoding inverts it exactly, for images and for video."""
+    autoencoder = IdentityAutoEncoder(latent_shift=0.25, latent_scale=4.0)
+    x = jax.random.normal(rng, shape)
+    latent = autoencoder.encode(x)
+    assert jnp.allclose(latent, (x - 0.25) * 4.0, atol=1e-6)
+    assert jnp.allclose(autoencoder.decode(latent), x, atol=1e-5)
+
+
+def test_latent_normalization_whitens_a_known_distribution(rng):
+    """The point of per-dataset stats: the diffusion model sees zero mean and
+    unit variance instead of whatever the encoder happens to produce."""
+    x = 3.0 + 5.0 * jax.random.normal(rng, (4096, 1, 1, 4))
+    autoencoder = IdentityAutoEncoder(latent_shift=float(jnp.mean(x)), latent_scale=1.0 / float(jnp.std(x)))
+    latent = autoencoder.encode(x)
+    assert abs(float(jnp.mean(latent))) < 1e-4
+    assert abs(float(jnp.std(latent)) - 1.0) < 1e-4
 
 
 @pytest.fixture(scope="module")
@@ -18,11 +58,21 @@ def vae():
     return StableDiffusionVAE(dtype=jnp.float32)
 
 
+@pytest.mark.network
 def test_vae_shapes(vae):
     assert vae.downscale_factor == 8
     assert vae.latent_channels == 4
 
 
+@pytest.mark.network
+def test_vae_uses_the_latent_normalization_seam(vae):
+    """The SD scaling factor rides on the shared seam, so there is one
+    normalization path a caller can override with dataset statistics."""
+    assert vae.latent_scale == pytest.approx(0.18215)
+    assert vae.latent_shift == 0.0
+
+
+@pytest.mark.network
 def test_vae_roundtrip_reconstructs(vae, rng):
     # A smooth image should survive the encode/decode roundtrip well
     ramp = jnp.linspace(-0.8, 0.8, 64)
diff --git a/training.py b/training.py
index f34a72b..f13febd 100644
--- a/training.py
+++ b/training.py
@@ -101,7 +101,11 @@ def boolean_string(s):
                     default='/home/mrwhite0racle/gcs_mount', help="Dataset location path")
 
 parser.add_argument('--noise_schedule', type=str, default='edm',
-                    choices=['cosine', 'karras', 'edm'], help='Noise schedule')
+                    choices=['cosine', 'karras', 'edm', 'flow', 'flow_matching'], help='Noise schedule')
+parser.add_argument('--min_snr_gamma', type=float, default=None,
+                    help='min-SNR-gamma loss weighting (Hang et al. 2023). 5.0 is the paper default; unset keeps the schedule own weighting.')
+parser.add_argument('--flow_shift', type=float, default=1.0,
+                    help='Resolution shift for the flow matching schedule. See flaxdiff.schedulers.flow.compute_resolution_shift.')
 
 # Any name from flaxdiff.models.registry, optionally with +2d/+hilbert/+zigzag
 # suffixes; validated by build_model against the registry itself
@@ -477,6 +481,10 @@ def main(args):
             from flaxdiff.metrics.images import get_clip_score_metric
             print("Using CLIPScore (val/clip_score, higher is better) for validation")
             eval_metrics.append(get_clip_score_metric())
+        if 'fid' in args.val_metrics:
+            from flaxdiff.metrics.fid import get_fid_metric
+            print("Using per-batch FID (val/fid) for validation")
+            eval_metrics.append(get_fid_metric())
     
     CONFIG = {
         "model": model_config,
@@ -501,7 +509,9 @@ def main(args):
     
     batches = batches if args.steps_per_epoch is None else args.steps_per_epoch
 
-    train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset(args.noise_schedule)
+    train_schedule, sampling_schedule, prediction_transform = get_diffusion_preset(
+        args.noise_schedule, shift=args.flow_shift, min_snr_gamma=args.min_snr_gamma,
+    )
     
     if args.experiment_name is not None:
         experiment_name = args.experiment_name
diff --git a/training_jepa.py b/training_jepa.py
new file mode 100644
index 0000000..9efd473
--- /dev/null
+++ b/training_jepa.py
@@ -0,0 +1,268 @@
+"""Train a JEPA encoder (I-JEPA over images, V-JEPA over video).
+
+A sibling of training.py rather than a flag on it: the two share the data
+pipeline, the registry and the trainer, but nothing else. A JEPA run has no
+noise schedule, no sampler, no text conditioning and no VAE, and folding an
+--objective switch into training.py would mean threading None through all of
+them. The Objective seam is what makes the same trainer serve both.
+"""
+
+import argparse
+import json
+import os
+import resource
+import time
+from datetime import datetime
+
+import jax
+import optax
+
+from flaxdiff.data.dataloaders import get_dataset_grain, get_dataset_online
+from flaxdiff.inputs import DiffusionInputConfig
+from flaxdiff.jepa import (
+    JepaObjective, multi_block_mask, get_linear_probe_metric, get_knn_probe_metric,
+)
+from flaxdiff.models.registry import build_model, canonicalize_architecture
+from flaxdiff.trainer import GeneralDiffusionTrainer
+
+os.environ['TOKENIZERS_PARALLELISM'] = "false"
+
+OPTIMIZER_MAP = {'adam': optax.adam, 'adamw': optax.adamw, 'lamb': optax.lamb}
+
+
+def boolean_string(s):
+    return s if isinstance(s, bool) else s == 'True'
+
+
+parser = argparse.ArgumentParser(description='Train a JEPA encoder')
+parser.add_argument('--GRAIN_WORKER_COUNT', type=int, default=32)
+parser.add_argument('--GRAIN_READ_THREAD_COUNT', type=int, default=140)
+parser.add_argument('--GRAIN_READ_BUFFER_SIZE', type=int, default=96)
+parser.add_argument('--GRAIN_WORKER_BUFFER_SIZE', type=int, default=100)
+
+parser.add_argument('--dataset', type=str, default='oxford_flowers102')
+parser.add_argument('--dataset_path', type=str, default='/home/mrwhite0racle/gcs_mount')
+parser.add_argument('--dataset_seed', type=int, default=0)
+parser.add_argument('--batch_size', type=int, default=64)
+parser.add_argument('--image_size', type=int, default=128)
+parser.add_argument('--frames_per_sample', type=int, default=None,
+                    help='Set for video (V-JEPA); leave unset for images (I-JEPA)')
+parser.add_argument('--epochs', type=int, default=100)
+parser.add_argument('--steps_per_epoch', type=int, default=None)
+parser.add_argument('--val_steps_per_epoch', type=int, default=4)
+
+parser.add_argument('--architecture', type=str, default='jepa_encoder',
+                    choices=['jepa_encoder', 'jepa_video_encoder'])
+parser.add_argument('--patch_size', type=int, default=16)
+parser.add_argument('--emb_features', type=int, default=384)
+parser.add_argument('--num_layers', type=int, default=12)
+parser.add_argument('--num_heads', type=int, default=6)
+parser.add_argument('--mlp_ratio', type=int, default=4)
+parser.add_argument('--ssm_attention_ratio', type=str, default='all-attn',
+                    help='Mixer pattern per layer: "all-attn", "all-ssm", "3:1", ...')
+parser.add_argument('--ssm_state_dim', type=int, default=64)
+parser.add_argument('--use_hilbert', type=boolean_string, default=False)
+parser.add_argument('--use_zigzag', type=boolean_string, default=False)
+parser.add_argument('--dropout_rate', type=float, default=0.0)
+parser.add_argument('--attention_impl', type=str, default=None,
+                    choices=[None, 'xla', 'cudnn', 'tpu'])
+parser.add_argument('--dtype', type=str, default=None)
+parser.add_argument('--precision', type=str, default='default',
+                    choices=['high', 'default', 'highest', 'None', None])
+
+parser.add_argument('--predictor_features', type=int, default=192)
+parser.add_argument('--predictor_layers', type=int, default=6)
+parser.add_argument('--predictor_heads', type=int, default=6)
+
+parser.add_argument('--num_target_blocks', type=int, default=4)
+parser.add_argument('--target_scale', type=float, nargs=2, default=[0.15, 0.2])
+parser.add_argument('--target_aspect', type=float, nargs=2, default=[0.75, 1.5])
+parser.add_argument('--momentum', type=float, nargs=2, default=[0.996, 1.0],
+                    help='Target-encoder EMA momentum, ramped over --momentum_steps')
+parser.add_argument('--momentum_steps', type=int, default=None,
+                    help='Defaults to the full training run')
+
+parser.add_argument('--optimizer', type=str, default='adamw', choices=list(OPTIMIZER_MAP))
+parser.add_argument('--optimizer_opts', type=str, default='{}')
+parser.add_argument('--learning_rate', type=float, default=1e-3)
+parser.add_argument('--learning_rate_schedule', type=str, default=None, choices=[None, 'cosine'])
+parser.add_argument('--learning_rate_peak', type=float, default=1.5e-3)
+parser.add_argument('--learning_rate_end', type=float, default=1e-6)
+parser.add_argument('--learning_rate_warmup_steps', type=int, default=10000)
+parser.add_argument('--learning_rate_decay_epochs', type=int, default=1)
+parser.add_argument('--clip_grads', type=float, default=0)
+
+parser.add_argument('--probe_classes', type=int, default=None,
+                    help='Number of classes for the frozen-encoder probes')
+parser.add_argument('--probe_label_key', type=str, default='label')
+parser.add_argument('--knn_k', type=int, default=20)
+
+parser.add_argument('--distributed_training', type=boolean_string, default=True)
+parser.add_argument('--experiment_name', type=str, default=None)
+parser.add_argument('--load_from_checkpoint', type=str, default=None)
+parser.add_argument('--resume_last_run', type=str, default=None)
+parser.add_argument('--checkpoint_dir', type=str, default='./checkpoints')
+parser.add_argument('--checkpoint_fs', type=str, default='local', choices=['local', 'gcs'])
+parser.add_argument('--max_checkpoints_to_keep', type=int, default=1)
+parser.add_argument('--wandb_project', type=str, default='mlops-msml605-project')
+parser.add_argument('--wandb_entity', type=str, default='umd-projects')
+
+
+def main(args):
+    resource.setrlimit(resource.RLIMIT_CORE, (resource.RLIM_INFINITY, resource.RLIM_INFINITY))
+    resource.setrlimit(resource.RLIMIT_OFILE, (65535, 65535))
+
+    print("Initializing JAX")
+    try:
+        jax.distributed.initialize()
+    except Exception:
+        pass
+    print(f"Number of devices: {jax.device_count()}")
+
+    checkpoint_dir = (f"gs://{args.checkpoint_dir}" if args.checkpoint_fs == 'gcs'
+                      else args.checkpoint_dir)
+
+    dataset_generator = (get_dataset_online if 'online' in args.dataset
+                         else get_dataset_grain)
+    data = dataset_generator(
+        args.dataset,
+        batch_size=args.batch_size, image_scale=args.image_size,
+        worker_count=args.GRAIN_WORKER_COUNT,
+        read_thread_count=args.GRAIN_READ_THREAD_COUNT,
+        read_buffer_size=args.GRAIN_READ_BUFFER_SIZE,
+        worker_buffer_size=args.GRAIN_WORKER_BUFFER_SIZE,
+        seed=args.dataset_seed,
+        dataset_source=args.dataset_path,
+    )
+    steps_per_epoch = args.steps_per_epoch or data['train_len'] // args.batch_size
+    total_steps = steps_per_epoch * args.epochs
+
+    architecture, suffix_flags = canonicalize_architecture(args.architecture)
+    is_video = args.frames_per_sample is not None
+    if is_video != (architecture == 'jepa_video_encoder'):
+        raise ValueError(
+            "--frames_per_sample and --architecture=jepa_video_encoder go together")
+
+    grid = (args.image_size // args.patch_size,) * 2
+    mask = multi_block_mask(
+        grid,
+        num_targets=args.num_target_blocks,
+        scale=tuple(args.target_scale),
+        aspect=tuple(args.target_aspect),
+    )
+    print(f"Mask geometry: {mask.block_area} tokens per target block "
+          f"({mask.block_shapes}), {mask.num_context} context tokens of {mask.num_patches}")
+
+    shared = {
+        "emb_features": args.emb_features,
+        "num_heads": args.num_heads,
+        "mlp_ratio": args.mlp_ratio,
+        "ssm_attention_ratio": args.ssm_attention_ratio,
+        "ssm_state_dim": args.ssm_state_dim,
+        "dropout_rate": args.dropout_rate,
+        "attention_impl": args.attention_impl,
+        "dtype": args.dtype,
+        "precision": args.precision,
+    }
+    encoder_config = {
+        **shared,
+        "patch_size": args.patch_size,
+        "num_layers": args.num_layers,
+        "use_hilbert": args.use_hilbert or suffix_flags.get('use_hilbert', False),
+        "use_zigzag": args.use_zigzag or suffix_flags.get('use_zigzag', False),
+    }
+    predictor_config = {
+        **shared,
+        "grid": grid,
+        "predictor_features": args.predictor_features,
+        "num_layers": args.predictor_layers,
+        "num_heads": args.predictor_heads,
+        "factorized": is_video,
+    }
+    encoder = build_model(architecture, encoder_config)
+    predictor_config["scan_order"] = encoder.scan_order
+    predictor = build_model('jepa_predictor', predictor_config)
+
+    sample_data_shape = ((args.frames_per_sample, args.image_size, args.image_size, 3)
+                         if is_video else (args.image_size, args.image_size, 3))
+    input_config = DiffusionInputConfig(
+        sample_data_key='video' if is_video else 'image',
+        sample_data_shape=sample_data_shape,
+        conditions=[],
+    )
+    objective = JepaObjective(
+        encoder=encoder,
+        predictor=predictor,
+        mask=mask,
+        sample_data_key=input_config.sample_data_key,
+        sample_data_shape=sample_data_shape,
+        momentum=tuple(args.momentum),
+        momentum_steps=args.momentum_steps or total_steps,
+    )
+
+    eval_metrics = []
+    if args.probe_classes:
+        eval_metrics = [
+            get_linear_probe_metric(args.probe_classes, label_key=args.probe_label_key),
+            get_knn_probe_metric(args.probe_classes, label_key=args.probe_label_key,
+                                 k=args.knn_k),
+        ]
+
+    learning_rate = args.learning_rate
+    if args.learning_rate_schedule == 'cosine':
+        learning_rate = optax.warmup_cosine_decay_schedule(
+            init_value=learning_rate, peak_value=args.learning_rate_peak,
+            warmup_steps=args.learning_rate_warmup_steps,
+            decay_steps=steps_per_epoch * args.learning_rate_decay_epochs,
+            end_value=args.learning_rate_end)
+    solver = OPTIMIZER_MAP[args.optimizer](learning_rate, **json.loads(args.optimizer_opts))
+    if args.clip_grads > 0:
+        solver = optax.chain(optax.clip_by_global_norm(args.clip_grads), solver)
+
+    experiment_name = args.experiment_name or (
+        f"jepa-{args.dataset}/res-{args.image_size}/patch-{args.patch_size}/"
+        f"mixer-{args.ssm_attention_ratio}/emb-{args.emb_features}/"
+        f"lr-{args.learning_rate}/date-{datetime.now().strftime('%Y-%m-%d_%H:%M:%S')}")
+    print("Experiment_Name:", experiment_name)
+
+    wandb_config = {
+        "project": args.wandb_project,
+        "entity": args.wandb_entity,
+        "name": experiment_name,
+        "config": {
+            "encoder": encoder_config,
+            "predictor": predictor_config,
+            "architecture": architecture,
+            "mask": {"grid": grid, "block_shapes": mask.block_shapes,
+                     "block_area": mask.block_area, "num_context": mask.num_context},
+            "dataset": {"name": args.dataset, "length": data['train_len']},
+            "arguments": vars(args),
+        },
+    }
+    if args.resume_last_run is not None:
+        wandb_config['id'] = args.resume_last_run
+
+    trainer = GeneralDiffusionTrainer(
+        model=encoder,
+        optimizer=solver,
+        input_config=input_config,
+        rngs=jax.random.PRNGKey(4),
+        objective=objective,
+        name=experiment_name,
+        wandb_config=wandb_config,
+        distributed_training=args.distributed_training,
+        checkpoint_base_path=checkpoint_dir,
+        load_from_checkpoint=args.load_from_checkpoint,
+        max_checkpoints_to_keep=args.max_checkpoints_to_keep,
+        eval_metrics=eval_metrics,
+        best_tracker_metric="val/knn_probe_accuracy" if eval_metrics else "train/best_loss",
+    )
+
+    start = time.time()
+    trainer.fit(data, training_steps_per_epoch=steps_per_epoch, epochs=args.epochs,
+                val_steps_per_epoch=args.val_steps_per_epoch)
+    print(f"Training finished in {time.time() - start:.0f}s")
+
+
+if __name__ == '__main__':
+    main(parser.parse_args())