configuration_vjepa21.py

"""V-JEPA 2.1 model configuration"""

from typing import Optional

from transformers import PretrainedConfig


# Encoder hierarchical layers, from `app/vjepa_2_1/models/vision_transformer.py`
# in `facebookresearch/vjepa2`.
_ENCODER_LAYER_MAP = {
    4: [0, 1, 2, 3],
    8: [1, 3, 5, 7],
    12: [2, 5, 8, 11],
    20: [4, 9, 14, 19],
    24: [5, 11, 17, 23],
    40: [9, 19, 29, 39],
    48: [11, 23, 37, 47],
}

# Predictor hierarchical layers, from `app/vjepa_2_1/models/predictor.py`.
# NOTE: this is *not* the same table as the encoder's. At depth 24 the reference
# predictor uses [4, 11, 17, 23] while the encoder uses [5, 11, 17, 23], and the
# predictor table has no entry for depth 48. Sharing one table happens to work
# today because only `len(...)` is consumed, but it would silently produce wrong
# indices as soon as multi-level predictor outputs are exposed.
_PREDICTOR_LAYER_MAP = {
    4: [0, 1, 2, 3],
    8: [1, 3, 5, 7],
    12: [2, 5, 8, 11],
    20: [4, 9, 14, 19],
    24: [4, 11, 17, 23],
    40: [9, 19, 29, 39],
}


def _get_hierarchical_layers(depth: int) -> list[int]:
    """Encoder hierarchical layer indices for a given depth."""
    if depth not in _ENCODER_LAYER_MAP:
        raise ValueError(
            f"Unsupported encoder depth {depth}. Supported depths: "
            f"{list(_ENCODER_LAYER_MAP.keys())}"
        )
    return _ENCODER_LAYER_MAP[depth]


def _get_predictor_hierarchical_layers(depth: int) -> list[int]:
    """Predictor hierarchical layer indices for a given depth."""
    if depth not in _PREDICTOR_LAYER_MAP:
        raise ValueError(
            f"Unsupported predictor depth {depth}. Supported depths: "
            f"{list(_PREDICTOR_LAYER_MAP.keys())}"
        )
    return _PREDICTOR_LAYER_MAP[depth]


class VJEPA21Config(PretrainedConfig):
    r"""
    Configuration class for the V-JEPA 2.1 model.

    V-JEPA 2.1 extends V-JEPA 2 with:
    - Multi-modality support (image + video with modality embeddings)
    - Hierarchical output distillation across intermediate layers
    - Interpolatable RoPE for variable input resolutions
    - Dense predictive loss with context token prediction

    Args:
        patch_size (`int`, defaults to 16):
            Spatial patch size.
        crop_size (`int`, defaults to 384):
            Input resolution of the model.
        frames_per_clip (`int`, defaults to 64):
            Number of frames in a video clip used during pre-training. This is
            informational: the model accepts any number of frames at inference.
        tubelet_size (`int`, defaults to 2):
            Temporal patch size (number of frames per tubelet).
        hidden_size (`int`, defaults to 1024):
            Encoder embedding dimension.
        in_chans (`int`, defaults to 3):
            Number of input channels.
        num_attention_heads (`int`, defaults to 16):
            Number of attention heads in the encoder.
        num_hidden_layers (`int`, defaults to 24):
            Number of encoder transformer layers.
        drop_path_rate (`float`, defaults to 0.0):
            Stochastic depth rate.
        mlp_ratio (`float`, defaults to 4.0):
            Ratio of MLP hidden dim to embedding dim.
        layer_norm_eps (`float`, defaults to 1e-6):
            Layer normalization epsilon.
        qkv_bias (`bool`, defaults to True):
            Whether to use bias in QKV projection.
        hidden_act (`str`, defaults to "gelu"):
            Activation function in MLP. "silu" enables SwiGLU.
        wide_silu (`bool`, defaults to True):
            Whether to use wide SwiGLU (2/3 hidden features) when hidden_act is "silu".
        initializer_range (`float`, defaults to 0.02):
            Standard deviation for weight initialization.
        attention_probs_dropout_prob (`float`, defaults to 0.0):
            Dropout probability for attention weights.
        img_temporal_dim_size (`int` or `None`, defaults to 1):
            Temporal dimension for image inputs. When set, a separate patch embedding
            with tubelet_size=1 is used for images. Set to None to disable.
        interpolate_rope (`bool`, defaults to True):
            Whether to interpolate RoPE frequencies for variable input resolutions.
        modality_embedding (`bool`, defaults to True):
            Whether to add learned modality embeddings (image vs video).
        n_output_distillation (`int`, defaults to 4):
            Number of intermediate encoder layers for hierarchical output.
            Set to 1 to only use the final layer output.
        n_registers (`int`, defaults to 0):
            Number of register tokens (appended to sequence).
        has_cls_first (`bool`, defaults to False):
            Whether the sequence starts with a CLS token.
        num_pooler_layers (`int`, defaults to 3):
            Number of self-attention layers in the attentive pooler. Together with
            the cross-attention layer this reproduces `AttentivePooler(depth=4)`,
            which is `num_probe_blocks: 4` in the reference evaluation configs.
        num_pooler_heads (`int` or `None`, defaults to 16):
            Number of attention heads in the attentive pooler. 16 is the value used
            by every frozen-probe config under `configs/eval_2_1/` in the reference
            repository (`classifier.num_heads: 16`), for all four model sizes, so it
            is the default here rather than `num_attention_heads`. The pooler is
            always trained from scratch, so this only affects the probe you train.
        pred_hidden_size (`int`, defaults to 384):
            Predictor embedding dimension.
        pred_num_attention_heads (`int`, defaults to 12):
            Number of attention heads in the predictor.
        pred_num_hidden_layers (`int`, defaults to 12):
            Number of predictor transformer layers.
        pred_num_mask_tokens (`int`, defaults to 8):
            Number of learnable mask tokens in the predictor.
        pred_zero_init_mask_tokens (`bool`, defaults to True):
            Whether to zero-initialize mask tokens.
        pred_mlp_ratio (`float`, defaults to 4.0):
            MLP ratio in the predictor.
        pred_teacher_embed_dim (`int` or `None`, defaults to None):
            Teacher embedding dimension for predictor output projection.
            When set, predictor projects to teacher_embed_dim // n_hierarchical_layers per layer.
        pred_return_all_tokens (`bool`, defaults to False):
            Whether the predictor returns predictions for both masked and context tokens.
    """

    model_type = "vjepa21"

    def __init__(
        self,
        patch_size: int = 16,
        crop_size: int = 384,
        frames_per_clip: int = 64,
        tubelet_size: int = 2,
        hidden_size: int = 1024,
        in_chans: int = 3,
        num_attention_heads: int = 16,
        num_hidden_layers: int = 24,
        drop_path_rate: float = 0.0,
        mlp_ratio: float = 4.0,
        layer_norm_eps: float = 1e-6,
        qkv_bias: bool = True,
        hidden_act: str = "gelu",
        wide_silu: bool = True,
        initializer_range: float = 0.02,
        attention_probs_dropout_prob: float = 0.0,
        # V-JEPA 2.1 specific
        img_temporal_dim_size: Optional[int] = 1,
        interpolate_rope: bool = True,
        modality_embedding: bool = True,
        n_output_distillation: int = 4,
        n_registers: int = 0,
        has_cls_first: bool = False,
        # Pooler
        num_pooler_layers: int = 3,
        num_pooler_heads: Optional[int] = 16,
        # Predictor
        pred_hidden_size: int = 384,
        pred_num_attention_heads: int = 12,
        pred_num_hidden_layers: int = 12,
        pred_num_mask_tokens: int = 8,
        pred_zero_init_mask_tokens: bool = True,
        pred_mlp_ratio: float = 4.0,
        pred_teacher_embed_dim: Optional[int] = None,
        pred_return_all_tokens: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.patch_size = patch_size
        self.crop_size = crop_size
        self.frames_per_clip = frames_per_clip
        self.tubelet_size = tubelet_size
        self.hidden_size = hidden_size
        self.in_chans = in_chans
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.drop_path_rate = drop_path_rate
        self.mlp_ratio = mlp_ratio
        self.layer_norm_eps = layer_norm_eps
        self.qkv_bias = qkv_bias
        self.hidden_act = hidden_act
        self.wide_silu = wide_silu
        self.initializer_range = initializer_range
        self.attention_probs_dropout_prob = attention_probs_dropout_prob

        # V-JEPA 2.1 specific
        self.img_temporal_dim_size = img_temporal_dim_size
        self.interpolate_rope = interpolate_rope
        self.modality_embedding = modality_embedding
        self.n_output_distillation = n_output_distillation
        self.n_registers = n_registers
        self.has_cls_first = has_cls_first

        # Pooler
        self.num_pooler_layers = num_pooler_layers
        self.num_pooler_heads = num_pooler_heads if num_pooler_heads is not None else 16

        # Predictor
        self.pred_hidden_size = pred_hidden_size
        self.pred_num_attention_heads = pred_num_attention_heads
        self.pred_num_hidden_layers = pred_num_hidden_layers
        self.pred_num_mask_tokens = pred_num_mask_tokens
        self.pred_zero_init_mask_tokens = pred_zero_init_mask_tokens
        self.pred_mlp_ratio = pred_mlp_ratio
        self.pred_teacher_embed_dim = pred_teacher_embed_dim
        self.pred_return_all_tokens = pred_return_all_tokens

        self._validate()

    def _validate(self) -> None:
        n_levels = len(_get_hierarchical_layers(self.num_hidden_layers))
        n_pred_levels = len(_get_predictor_hierarchical_layers(self.pred_num_hidden_layers))
        if not 1 <= self.n_output_distillation <= n_levels:
            raise ValueError(
                f"n_output_distillation must be in [1, {n_levels}] for a model with "
                f"{self.num_hidden_layers} layers, got {self.n_output_distillation}."
            )
        if self.n_output_distillation > n_pred_levels:
            raise ValueError(
                f"n_output_distillation ({self.n_output_distillation}) exceeds the "
                f"{n_pred_levels} hierarchical levels available in a predictor with "
                f"{self.pred_num_hidden_layers} layers."
            )
        if self.hidden_size % self.num_attention_heads != 0:
            raise ValueError(
                f"hidden_size ({self.hidden_size}) must be divisible by "
                f"num_attention_heads ({self.num_attention_heads})."
            )
        if self.hidden_size % self.num_pooler_heads != 0:
            raise ValueError(
                f"hidden_size ({self.hidden_size}) must be divisible by "
                f"num_pooler_heads ({self.num_pooler_heads})."
            )
        if self.pred_hidden_size % self.pred_num_attention_heads != 0:
            raise ValueError(
                f"pred_hidden_size ({self.pred_hidden_size}) must be divisible by "
                f"pred_num_attention_heads ({self.pred_num_attention_heads})."
            )
        if self.pred_teacher_embed_dim is not None:
            if self.pred_teacher_embed_dim % self.n_output_distillation != 0:
                raise ValueError(
                    f"pred_teacher_embed_dim ({self.pred_teacher_embed_dim}) must be "
                    f"divisible by n_output_distillation ({self.n_output_distillation})."
                )
        if self.tubelet_size < 1:
            raise ValueError(f"tubelet_size must be >= 1, got {self.tubelet_size}.")
        if self.pred_num_mask_tokens < 1:
            raise ValueError(
                f"pred_num_mask_tokens must be >= 1, got {self.pred_num_mask_tokens}."
            )

    @property
    def encoder_hierarchical_layers(self) -> list[int]:
        """Layer indices at which the encoder carries a per-level LayerNorm."""
        return _get_hierarchical_layers(self.num_hidden_layers)

    @property
    def encoder_distillation_layers(self) -> list[int]:
        """Encoder layer indices contributing to the hierarchical output."""
        all_layers = _get_hierarchical_layers(self.num_hidden_layers)
        return all_layers[-self.n_output_distillation :]

    @property
    def predictor_hierarchical_layers(self) -> list[int]:
        """Predictor layer indices for hierarchical output.

        Uses the predictor's own depth table, which differs from the encoder's at
        depth 24 (`[4, 11, 17, 23]` vs `[5, 11, 17, 23]`).
        """
        all_layers = _get_predictor_hierarchical_layers(self.pred_num_hidden_layers)
        return all_layers[-self.n_output_distillation :]

    @property
    def pretrained_grid_size(self) -> int:
        """Grid size used during pre-training (for RoPE interpolation)."""
        if self.patch_size == 14:
            return int(252 / self.patch_size)
        return int(256 / self.patch_size)


__all__ = ["VJEPA21Config"]