dew.objectives.jepa

Name	Summary
FactorizedTokenStack	Spatial then temporal blocks over [B, T, N, F], as in VideoDiT.
JepaEncoder	ViT over an image, optionally restricted to a subset of its patches.
JepaObjective	Joint-embedding prediction over images (B,H,W,C) or video (B,T,H,W,C).
JepaPredictor	Narrow transformer from context embeddings to target embeddings.
JepaVideoEncoder	Factorized spatial-temporal encoder over (B, T, H, W, C).
KnnProbe	Report cosine k-NN accuracy over each validation batch’s representations.
LinearProbe	Report linear probe accuracy over each validation batch’s representations.
MultiBlockMask	Hold one patch grid’s static mask geometry, and sample masks over it.
TokenStack	A stack of unmodulated blocks over a token sequence.
knn_probe
knn_probe_accuracy	Score cosine k-NN accuracy, fitting on half the batch and scoring the rest.
linear_probe
linear_probe_accuracy	Fit a logistic regression on half the batch and score its accuracy on the rest.
multi_block_mask	Resolve the I-JEPA mask geometry for a patch grid.
normalize_targets	Feature-wise layer norm with no learned affine.
representation_health	Report collapse telemetry for pooled embeddings [B, D].

FactorizedTokenStack

Flax module source

class FactorizedTokenStack(
    mlp_ratio: int = 4,
    ssm_attention_ratio: str = 'all-attn',
    block_pattern: Sequence[str] | None = None,
    ssm_state_dim: int = 64,
    bidirectional_ssm: bool = True,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
)

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 is rotated by frame index while the spatial half runs unrotated.

FactorizedTokenStack.setup

def setup()

JepaEncoder

Flax module source

class JepaEncoder(
    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: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
    scan_order: Literal['raster', 'hilbert', 'zigzag'] = 'raster',
    stack_type: type[TokenStack | FactorizedTokenStack] = TokenStack,
)

ViT over an image, optionally restricted to a subset of its patches.

stack_type: type[TokenStack | FactorizedTokenStack]

The layers between the patches and the norm: TokenStack over one image’s tokens, FactorizedTokenStack over a clip’s frames.

JepaEncoder.setup

def setup()

JepaObjective

class source

class JepaObjective(
    encoder: nn.Module,
    predictor: nn.Module,
    mask: MultiBlockMask,
    sample: Field,
    momentum: tuple[float, float] = (0.996, 1.0),
    momentum_steps: int = 100000,
    label_key: str = LABEL_KEY,
)

Joint-embedding prediction over images (B,H,W,C) or video (B,T,H,W,C).

Evaluation returns the pooled target-encoder embeddings of a batch with its labels, which the probe metrics score.

JepaObjective.init

def init(key, variables: Variables | None = None)

JepaObjective.encode

def encode(encoder_params, samples, token_idx=None, train=False, rngs=None) -> jax.Array

JepaObjective.loss

def loss(params, batch, step: Step)

JepaObjective.evaluate

def evaluate(params, batch, step: Step)

The frozen target encoder’s pooled embeddings, with the batch labels.

JepaPredictor

Flax module source

class JepaPredictor(
    grid: tuple[int, int] = (14, 14),
    emb_features: int = 384,
    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: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
    scan_order: str = 'raster',
    factorized: bool = False,
)

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.

JepaPredictor.setup

def setup()

JepaVideoEncoder

class source

class JepaVideoEncoder(JepaEncoder)

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.

KnnProbe

dataclass source

class KnnProbe(num_classes: int, k: int = 20)

Report cosine k-NN accuracy over each validation batch’s representations.

KnnProbe.merge

def merge(
    accumulated: tuple[float, float],
    contribution: tuple[float, float],
) -> tuple[float, float]

KnnProbe.finalize

def finalize(accumulated: tuple[float, float]) -> float

LinearProbe

dataclass source

class LinearProbe(
    num_classes: int,
    steps: int = 100,
    learning_rate: float = 0.01,
    weight_decay: float = 0.0001,
)

Report linear probe accuracy over each validation batch’s representations.

LinearProbe.merge

def merge(
    accumulated: tuple[float, float],
    contribution: tuple[float, float],
) -> tuple[float, float]

LinearProbe.finalize

def finalize(accumulated: tuple[float, float]) -> float

MultiBlockMask

dataclass source

class MultiBlockMask(
    grid: tuple[int, int],
    num_targets: int,
    block_shapes: tuple[tuple[int, int], ...],
    num_context: int,
)

Hold one patch grid’s static mask geometry, and sample masks over it.

MultiBlockMask.sample

def sample(rng: jax.Array, batch_size: int)

Draw one batch of context and target token indices.

Returns the context indices, [B, num_context], and the target indices, [B, num_targets, block_area].

TokenStack

Flax module source

class TokenStack(
    mlp_ratio: int = 4,
    ssm_attention_ratio: str = 'all-attn',
    block_pattern: Sequence[str] | None = None,
    ssm_state_dim: int = 64,
    bidirectional_ssm: bool = True,
    dropout_rate: float = 0.0,
    dtype: Dtype | None = None,
    precision: PrecisionLike = None,
    force_fp32_for_softmax: bool = True,
    norm_epsilon: float = 1e-05,
    qk_norm: bool = False,
    attention_impl: str = 'auto',
)

A stack of unmodulated blocks over a token sequence.

TokenStack.setup

def setup()

knn_probe

function source

def knn_probe(num_classes: int, k: int = 20) -> KnnProbe

knn_probe_accuracy

function source

def knn_probe_accuracy(embeddings, labels, num_classes: int, k: int = 20)

Score cosine k-NN accuracy, fitting on half the batch and scoring the rest.

linear_probe

function source

def linear_probe(
    num_classes: int,
    steps: int = 100,
    learning_rate: float = 0.01,
    weight_decay: float = 0.0001,
) -> LinearProbe

linear_probe_accuracy

function source

def linear_probe_accuracy(
    embeddings,
    labels,
    num_classes: int,
    steps: int = 100,
    learning_rate: float = 0.01,
    weight_decay: float = 0.0001,
)

Fit a logistic regression on half the batch and score its accuracy on the rest.

multi_block_mask

function source

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.

normalize_targets

function source

def normalize_targets(x, epsilon: float = 1e-06)

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. Shrinking the representation then does not lower the loss.

representation_health

function source

def representation_health(z) -> dict[str, jax.Array]

Report 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) while repr_std holds.

Both are computed in fp32. A run’s compute dtype then does not set the noise floor of the drift, and bf16 and fp32 runs read off the same curves.