| """Small causal readiness head over frozen visual, state, and text features.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| import torch |
| from torch import nn |
|
|
|
|
| class ResidualCrossAttentionBlock(nn.Module): |
| def __init__(self, width: int, heads: int, ffn_width: int, dropout: float): |
| super().__init__() |
| self.query_norm = nn.LayerNorm(width) |
| self.visual_norm = nn.LayerNorm(width) |
| self.attention = nn.MultiheadAttention( |
| width, heads, dropout=dropout, batch_first=True |
| ) |
| self.ffn_norm = nn.LayerNorm(width) |
| self.ffn = nn.Sequential( |
| nn.Linear(width, ffn_width), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(ffn_width, width), |
| nn.Dropout(dropout), |
| ) |
|
|
| def forward( |
| self, query: torch.Tensor, visual: torch.Tensor, visual_mask: torch.Tensor |
| ) -> torch.Tensor: |
| normalized_visual = self.visual_norm(visual) |
| update, _ = self.attention( |
| self.query_norm(query), |
| normalized_visual, |
| normalized_visual, |
| key_padding_mask=~visual_mask, |
| need_weights=False, |
| ) |
| query = query + update |
| return query + self.ffn(self.ffn_norm(query)) |
|
|
|
|
| class CommitReadinessHead(nn.Module): |
| """Predict evidence for committing away from the active semantic phase.""" |
|
|
| def __init__( |
| self, |
| *, |
| image_dim: int = 1152, |
| text_dim: int = 4096, |
| state_dim: int = 8, |
| width: int = 256, |
| state_transformer_layers: int = 1, |
| cross_attention_layers: int = 2, |
| attention_heads: int = 4, |
| ffn_width: int = 1024, |
| head_hidden_width: int = 128, |
| dropout: float = 0.1, |
| max_relative_age: int = 2048, |
| state_recent_length: int = 16, |
| ): |
| super().__init__() |
| self.image_dim = int(image_dim) |
| self.text_dim = int(text_dim) |
| self.state_dim = int(state_dim) |
| self.width = int(width) |
| self.state_recent_length = int(state_recent_length) |
|
|
| self.visual_projection = nn.Linear(image_dim, width) |
| self.visual_role_embedding = nn.Embedding(4, width) |
| self.visual_view_embedding = nn.Embedding(2, width) |
| self.visual_spatial_embedding = nn.Embedding(16, width) |
| self.visual_age_embedding = nn.Embedding(max_relative_age + 1, width) |
|
|
| self.state_projection = nn.Linear(state_dim, width) |
| self.state_role_embedding = nn.Embedding(2, width) |
| self.state_time_embedding = nn.Embedding(state_recent_length, width) |
| state_layer = nn.TransformerEncoderLayer( |
| d_model=width, |
| nhead=attention_heads, |
| dim_feedforward=ffn_width, |
| dropout=dropout, |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.state_transformer = nn.TransformerEncoder( |
| state_layer, num_layers=state_transformer_layers, enable_nested_tensor=False |
| ) |
| self.state_fusion = nn.Sequential( |
| nn.Linear(2 * width, width), nn.GELU(), nn.Dropout(dropout) |
| ) |
| self.text_projection = nn.Linear(text_dim, width) |
| self.query = nn.Parameter(torch.empty(1, 1, width)) |
| nn.init.normal_(self.query, std=0.02) |
| self.cross_attention = nn.ModuleList( |
| [ |
| ResidualCrossAttentionBlock(width, attention_heads, ffn_width, dropout) |
| for _ in range(cross_attention_layers) |
| ] |
| ) |
| self.output_norm = nn.LayerNorm(width) |
| self.output = nn.Sequential( |
| nn.Linear(width, head_hidden_width), |
| nn.GELU(), |
| nn.Dropout(dropout), |
| nn.Linear(head_hidden_width, 1), |
| ) |
|
|
| @classmethod |
| def from_config(cls, cfg: Any) -> "CommitReadinessHead": |
| return cls( |
| image_dim=int(cfg.features.image_dim), |
| text_dim=int(cfg.features.text_dim), |
| state_dim=int(cfg.features.state_dim), |
| width=int(cfg.model.width), |
| state_transformer_layers=int(cfg.model.state_transformer_layers), |
| cross_attention_layers=int(cfg.model.cross_attention_layers), |
| attention_heads=int(cfg.model.attention_heads), |
| ffn_width=int(cfg.model.ffn_width), |
| head_hidden_width=int(cfg.model.head_hidden_width), |
| dropout=float(cfg.model.dropout), |
| max_relative_age=int(cfg.features.max_relative_age), |
| state_recent_length=int(cfg.features.state_recent_length), |
| ) |
|
|
| def forward(self, **features: torch.Tensor) -> dict[str, torch.Tensor]: |
| visual_tokens = features["visual_tokens"] |
| visual_mask = features["visual_mask"].bool() |
| batch, token_count, image_dim = visual_tokens.shape |
| if token_count != 608 or image_dim != self.image_dim: |
| raise ValueError( |
| f"Expected visual [B,608,{self.image_dim}], got {tuple(visual_tokens.shape)}" |
| ) |
| if ( |
| visual_mask.shape != (batch, token_count) |
| or not visual_mask.any(dim=1).all() |
| ): |
| raise ValueError("Each sample must have at least one valid visual token") |
| visual = ( |
| self.visual_projection(visual_tokens) |
| + self.visual_role_embedding(features["visual_role_ids"].long()) |
| + self.visual_view_embedding(features["visual_view_ids"].long()) |
| + self.visual_spatial_embedding(features["visual_spatial_ids"].long()) |
| + self.visual_age_embedding(features["visual_age_ids"].long()) |
| ) |
|
|
| active_state = features["active_start_state"] |
| recent_states = features["recent_states"] |
| recent_mask = features["recent_state_mask"].bool() |
| if active_state.shape != (batch, self.state_dim): |
| raise ValueError( |
| f"Bad active-start state shape: {tuple(active_state.shape)}" |
| ) |
| if recent_states.shape != (batch, self.state_recent_length, self.state_dim): |
| raise ValueError(f"Bad recent state shape: {tuple(recent_states.shape)}") |
| if ( |
| recent_mask.shape != (batch, self.state_recent_length) |
| or not recent_mask.any(dim=1).all() |
| ): |
| raise ValueError("Each sample must have at least one valid recent state") |
| active_summary = ( |
| self.state_projection(active_state) + self.state_role_embedding.weight[0] |
| ) |
| time_ids = torch.arange(self.state_recent_length, device=recent_states.device) |
| recent = ( |
| self.state_projection(recent_states) |
| + self.state_role_embedding.weight[1] |
| + self.state_time_embedding(time_ids)[None] |
| ) |
| recent = self.state_transformer(recent, src_key_padding_mask=~recent_mask) |
| recent_summary = (recent * recent_mask.unsqueeze(-1)).sum( |
| dim=1 |
| ) / recent_mask.sum(dim=1, keepdim=True).clamp_min(1) |
| state_summary = self.state_fusion( |
| torch.cat([active_summary, recent_summary], dim=-1) |
| ) |
|
|
| text = features["text_embedding"] |
| if text.shape != (batch, self.text_dim): |
| raise ValueError(f"Bad text embedding shape: {tuple(text.shape)}") |
| text_summary = self.text_projection(text) |
| query = ( |
| self.query.expand(batch, -1, -1) |
| + state_summary[:, None] |
| + text_summary[:, None] |
| ) |
| for block in self.cross_attention: |
| query = block(query, visual, visual_mask) |
| commit_logit = self.output(self.output_norm(query[:, 0])).squeeze(-1) |
| return { |
| "commit_logit": commit_logit, |
| "commit_score": torch.sigmoid(commit_logit), |
| } |
|
|
|
|
| def trainable_parameter_count(model: nn.Module) -> int: |
| return sum( |
| parameter.numel() for parameter in model.parameters() if parameter.requires_grad |
| ) |
|
|
