"""Neural memory modules used by TrackingAnything.
Shapes follow the notation below:
B: batch size
M: memory slots (fixed to 6)
P: projected visual tokens per crop (256 for a 448 crop)
Q: resampled tokens per slot (fixed to 8)
D: LocateAnything language width (2048 for the 3B checkpoint)
"""
from __future__ import annotations
import math
from typing import Dict, Optional, Sequence, Tuple
import torch
from torch import nn
import torch.nn.functional as F
def _inverse_softplus(value: float) -> float:
if value <= 0:
raise ValueError("softplus target must be positive")
return math.log(math.expm1(value))
def _inverse_sigmoid(value: float) -> float:
if not 0.0 < value < 1.0:
raise ValueError("sigmoid target must be strictly between zero and one")
return math.log(value / (1.0 - value))
def normalized_log_memory_scores(
memory_scores: torch.Tensor,
slot_valid_mask: torch.Tensor,
*,
std_floor: float = 0.08,
clip_value: float = 2.0,
) -> torch.Tensor:
"""Return per-sample normalized log scores with invalid slots at zero.
The normalization is causal and deployment-safe: it uses only the valid
memories in the current bank, never the sequence length or future frames.
A standard-deviation floor prevents a nearly uniform bank from amplifying
numerical noise, while clipping gives a strict bound to every downstream
score prior.
"""
if memory_scores.ndim != 2:
raise ValueError("memory_scores must have shape [B, M]")
valid = slot_valid_mask.to(
device=memory_scores.device,
dtype=torch.bool,
)
if valid.shape != memory_scores.shape:
raise ValueError("slot_valid_mask must match memory_scores")
if std_floor <= 0.0:
raise ValueError("std_floor must be positive")
if clip_value <= 0.0:
raise ValueError("clip_value must be positive")
scores = torch.where(
valid & torch.isfinite(memory_scores),
memory_scores.float().clamp(1e-6, 1.0),
torch.ones_like(memory_scores, dtype=torch.float32),
)
log_scores = scores.log()
weights = valid.to(log_scores.dtype)
count = weights.sum(dim=-1, keepdim=True).clamp_min(1.0)
mean = (log_scores * weights).sum(dim=-1, keepdim=True) / count
centered = (log_scores - mean) * weights
variance = centered.square().sum(dim=-1, keepdim=True) / count
# Clamp the variance before sqrt. ``sqrt(0).clamp_min(floor)`` has the
# right forward value but an undefined/infinite derivative at a uniform
# differentiable score vector (V9 initializes every gate at 0.5), which
# turns the first optimizer update into NaNs. This ordering is forward-
# equivalent and gives a zero, finite gradient below the variance floor.
denominator = variance.clamp_min(float(std_floor) ** 2).sqrt()
normalized = (centered / denominator).clamp(
-float(clip_value),
float(clip_value),
)
return normalized * weights
def relative_time_bucket(
frame_gaps: torch.Tensor,
num_buckets: int = 8,
) -> torch.Tensor:
"""Map non-negative frame gaps to causal logarithmic time buckets.
Bucket 0 is the same frame (normally unused for a causal memory). Positive
gaps use powers of two: 1, 2--3, 4--7, ..., with the tail clipped to the
final bucket. This has no dependency on the future sequence length.
"""
if num_buckets < 2:
raise ValueError("num_buckets must be at least 2")
gaps = torch.as_tensor(frame_gaps)
if torch.any(gaps < 0):
raise ValueError("relative frame gaps must be non-negative")
positive = gaps.clamp_min(1).to(torch.float32)
buckets = torch.floor(torch.log2(positive)).to(torch.long) + 1
buckets = torch.where(gaps == 0, torch.zeros_like(buckets), buckets)
return buckets.clamp_max(num_buckets - 1)
class SwiGLU(nn.Module):
def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0) -> None:
super().__init__()
self.gate_up = nn.Linear(dim, 2 * hidden_dim)
self.down = nn.Linear(hidden_dim, dim)
self.dropout = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
gate, value = self.gate_up(x).chunk(2, dim=-1)
return self.down(self.dropout(F.silu(gate) * value))
class ResamplerLayer(nn.Module):
"""One Q-Former/Perceiver-style query resampling block."""
def __init__(
self,
dim: int,
num_heads: int,
ffn_dim: int,
dropout: float,
) -> None:
super().__init__()
self.query_norm = nn.LayerNorm(dim)
self.context_norm = nn.LayerNorm(dim)
self.ffn_norm = nn.LayerNorm(dim)
self.self_attn = nn.MultiheadAttention(
dim, num_heads, dropout=dropout, batch_first=True
)
self.cross_attn = nn.MultiheadAttention(
dim, num_heads, dropout=dropout, batch_first=True
)
self.ffn = SwiGLU(dim, ffn_dim, dropout)
self.dropout = nn.Dropout(dropout)
def forward(
self,
queries: torch.Tensor,
context: torch.Tensor,
context_padding_mask: Optional[torch.Tensor],
) -> torch.Tensor:
normalized = self.query_norm(queries)
update = self.self_attn(
normalized, normalized, normalized, need_weights=False
)[0]
queries = queries + self.dropout(update)
update = self.cross_attn(
self.query_norm(queries),
self.context_norm(context),
self.context_norm(context),
key_padding_mask=context_padding_mask,
need_weights=False,
)[0]
queries = queries + self.dropout(update)
return queries + self.dropout(self.ffn(self.ffn_norm(queries)))
class PerFrameResampler(nn.Module):
"""Compress every memory crop independently to exactly eight tokens."""
def __init__(
self,
dim: int = 2048,
num_queries: int = 8,
num_layers: int = 2,
num_heads: int = 16,
ffn_dim: int = 8192,
dropout: float = 0.1,
) -> None:
super().__init__()
self.num_queries = int(num_queries)
self.queries = nn.Parameter(torch.empty(num_queries, dim))
self.layers = nn.ModuleList(
[
ResamplerLayer(dim, num_heads, ffn_dim, dropout)
for _ in range(num_layers)
]
)
self.output_norm = nn.LayerNorm(dim)
nn.init.trunc_normal_(self.queries, std=0.02)
def forward(
self,
visual_tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
patch_valid_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Return ``[B, M, 8, D]`` tokens, with invalid slots exactly zero."""
if visual_tokens.ndim != 4:
raise ValueError("visual_tokens must have shape [B, M, P, D]")
batch, slots, patches, dim = visual_tokens.shape
slot_valid_mask = slot_valid_mask.to(
device=visual_tokens.device, dtype=torch.bool
)
if slot_valid_mask.shape != (batch, slots):
raise ValueError("slot_valid_mask must have shape [B, M]")
if patch_valid_mask is None:
patch_valid_mask = slot_valid_mask.unsqueeze(-1).expand(
batch, slots, patches
)
else:
patch_valid_mask = patch_valid_mask.to(
device=visual_tokens.device, dtype=torch.bool
)
if patch_valid_mask.shape != (batch, slots, patches):
raise ValueError("patch_valid_mask must have shape [B, M, P]")
flat_context = visual_tokens.reshape(batch * slots, patches, dim)
flat_valid = patch_valid_mask.reshape(batch * slots, patches)
padding_mask = ~flat_valid
# MultiheadAttention produces NaNs when every key is masked. Invalid
# slots receive one zero sentinel key and are zeroed again below.
all_invalid = ~flat_valid.any(dim=-1)
if all_invalid.any():
padding_mask = padding_mask.clone()
flat_context = flat_context.clone()
padding_mask[all_invalid, 0] = False
flat_context[all_invalid, 0] = 0
queries = self.queries.to(dtype=visual_tokens.dtype)
queries = queries.unsqueeze(0).expand(batch * slots, -1, -1)
for layer in self.layers:
queries = layer(queries, flat_context, padding_mask)
queries = self.output_norm(queries)
queries = queries.reshape(batch, slots, self.num_queries, dim)
return queries * slot_valid_mask[:, :, None, None].to(queries.dtype)
class TargetContextResamplerLayer(nn.Module):
"""Shared-parameter target/context query block with separated self-attention.
The first four and final four queries use the same attention and FFN
parameters. A block-diagonal self-attention mask prevents the two roles
from becoming indistinguishable before the per-frame visual compression is
complete, while a per-query soft region bias controls which image patches
each role may read.
"""
def __init__(
self,
dim: int,
num_heads: int,
ffn_dim: int,
dropout: float,
target_queries: int,
context_queries: int,
) -> None:
super().__init__()
self.num_heads = int(num_heads)
self.target_queries = int(target_queries)
self.context_queries = int(context_queries)
self.query_norm = nn.LayerNorm(dim)
self.context_norm = nn.LayerNorm(dim)
self.ffn_norm = nn.LayerNorm(dim)
self.self_attn = nn.MultiheadAttention(
dim, num_heads, dropout=dropout, batch_first=True
)
self.cross_attn = nn.MultiheadAttention(
dim, num_heads, dropout=dropout, batch_first=True
)
self.ffn = SwiGLU(dim, ffn_dim, dropout)
self.dropout = nn.Dropout(dropout)
total = self.target_queries + self.context_queries
separation = torch.ones(total, total, dtype=torch.bool)
separation[: self.target_queries, : self.target_queries] = False
separation[self.target_queries :, self.target_queries :] = False
self.register_buffer(
"query_group_separation_mask",
separation,
persistent=True,
)
def forward(
self,
queries: torch.Tensor,
context: torch.Tensor,
cross_attention_bias: torch.Tensor,
*,
return_attention: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
normalized = self.query_norm(queries)
update = self.self_attn(
normalized,
normalized,
normalized,
attn_mask=self.query_group_separation_mask,
need_weights=False,
)[0]
queries = queries + self.dropout(update)
if cross_attention_bias.shape != (
queries.shape[0], queries.shape[1], context.shape[1]
):
raise ValueError(
"cross_attention_bias must have shape [N,Q,P]"
)
# torch MultiheadAttention accepts one additive mask per batch/head.
attention_mask = cross_attention_bias.repeat_interleave(
self.num_heads, dim=0
).to(dtype=queries.dtype)
normalized_context = self.context_norm(context)
update, attention = self.cross_attn(
self.query_norm(queries),
normalized_context,
normalized_context,
attn_mask=attention_mask,
need_weights=return_attention,
average_attn_weights=False,
)
queries = queries + self.dropout(update)
queries = queries + self.dropout(self.ffn(self.ffn_norm(queries)))
return queries, attention if return_attention else None
class TargetContextPerFrameResampler(nn.Module):
"""V9 4-target/4-context resampler with shared neural parameters."""
def __init__(
self,
dim: int = 2048,
num_queries: int = 8,
num_layers: int = 2,
num_heads: int = 16,
ffn_dim: int = 8192,
dropout: float = 0.1,
region_bias_clip: float = 4.0,
region_bias_scale: float = 1.0,
) -> None:
super().__init__()
if num_queries != 8:
raise ValueError("V9 requires exactly eight queries per slot")
if region_bias_clip <= 0.0 or region_bias_scale <= 0.0:
raise ValueError("V9 region-bias constants must be positive")
self.num_queries = int(num_queries)
self.target_queries = self.num_queries // 2
self.context_queries = self.num_queries - self.target_queries
self.region_bias_clip = float(region_bias_clip)
self.region_bias_scale = float(region_bias_scale)
self.queries = nn.Parameter(torch.empty(num_queries, dim))
self.layers = nn.ModuleList(
[
TargetContextResamplerLayer(
dim,
num_heads,
ffn_dim,
dropout,
self.target_queries,
self.context_queries,
)
for _ in range(num_layers)
]
)
self.output_norm = nn.LayerNorm(dim)
nn.init.trunc_normal_(self.queries, std=0.02)
def forward(
self,
visual_tokens: torch.Tensor,
region_coverage: torch.Tensor,
slot_valid_mask: torch.Tensor,
patch_valid_mask: Optional[torch.Tensor] = None,
*,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if visual_tokens.ndim != 4:
raise ValueError("visual_tokens must have shape [B,M,P,D]")
batch, slots, patches, dim = visual_tokens.shape
valid_slots = slot_valid_mask.to(
device=visual_tokens.device, dtype=torch.bool
)
if valid_slots.shape != (batch, slots):
raise ValueError("slot_valid_mask must have shape [B,M]")
coverage = region_coverage.to(
device=visual_tokens.device, dtype=torch.float32
)
if coverage.shape != (batch, slots, patches):
raise ValueError("region_coverage must have shape [B,M,P]")
coverage = coverage.clamp(0.0, 1.0)
if patch_valid_mask is None:
patch_valid = valid_slots[:, :, None].expand(-1, -1, patches)
else:
patch_valid = patch_valid_mask.to(
device=visual_tokens.device, dtype=torch.bool
)
if patch_valid.shape != (batch, slots, patches):
raise ValueError("patch_valid_mask must have shape [B,M,P]")
flat_context = visual_tokens.reshape(batch * slots, patches, dim)
flat_valid = patch_valid.reshape(batch * slots, patches).clone()
all_invalid = ~flat_valid.any(dim=-1)
if all_invalid.any():
flat_context = flat_context.clone()
flat_context[all_invalid, 0] = 0
flat_valid[all_invalid, 0] = True
eps = 1e-6
target_bias = coverage.add(eps).log().clamp(
min=-self.region_bias_clip, max=0.0
)
context_bias = (1.0 - coverage).add(eps).log().clamp(
min=-self.region_bias_clip, max=0.0
)
target_bias = target_bias[:, :, None, :].expand(
-1, -1, self.target_queries, -1
)
context_bias = context_bias[:, :, None, :].expand(
-1, -1, self.context_queries, -1
)
attention_bias = self.region_bias_scale * torch.cat(
(target_bias, context_bias), dim=2
).reshape(batch * slots, self.num_queries, patches)
# A literal -inf mask gives finite forward values but can produce NaN
# gradients in torch MultiheadAttention for the all-empty physical
# slots. A finite floor is numerically equivalent for softmax while
# keeping the backward graph well defined in bf16/fp32.
mask_floor = -1.0e4
attention_bias = attention_bias.masked_fill(
~flat_valid[:, None, :], mask_floor
)
if all_invalid.any():
attention_bias = attention_bias.clone()
attention_bias[all_invalid] = mask_floor
attention_bias[all_invalid, :, 0] = 0.0
queries = self.queries.to(dtype=visual_tokens.dtype)[None].expand(
batch * slots, -1, -1
)
last_attention = None
for layer_index, layer in enumerate(self.layers):
queries, attention = layer(
queries,
flat_context,
attention_bias,
return_attention=(
return_diagnostics and layer_index == len(self.layers) - 1
),
)
if attention is not None:
last_attention = attention
queries = self.output_norm(queries).reshape(
batch, slots, self.num_queries, dim
)
queries = queries * valid_slots[:, :, None, None].to(queries.dtype)
diagnostics: Dict[str, torch.Tensor] = {}
if return_diagnostics:
diagnostics["target_context_region_bias"] = attention_bias.reshape(
batch, slots, self.num_queries, patches
)
if last_attention is not None:
diagnostics["target_context_cross_attention"] = (
last_attention.reshape(
batch,
slots,
last_attention.shape[1],
self.num_queries,
patches,
)
)
return queries, diagnostics
class DirectDynamicMemoryGate(nn.Module):
"""Current-aware V9 memory reliability predictor.
Four target and four context tokens independently search the full current
image patch grid. Their log-sum-exp retrieval margin, LoRAT confidence,
causal relative time and slot role are the only gate inputs.
"""
def __init__(
self,
dim: int,
projection_dim: int = 256,
hidden_dim: int = 256,
num_time_buckets: int = 8,
structural_embedding_dim: int = 32,
temperature: float = 0.07,
) -> None:
super().__init__()
if projection_dim <= 0 or hidden_dim <= 0:
raise ValueError("V9 gate dimensions must be positive")
if temperature <= 0.0:
raise ValueError("V9 gate temperature must be positive")
self.projection_dim = int(projection_dim)
self.temperature = float(temperature)
self.memory_norm = nn.LayerNorm(dim)
self.current_norm = nn.LayerNorm(dim)
self.memory_projection = nn.Linear(dim, projection_dim, bias=False)
self.current_projection = nn.Linear(dim, projection_dim, bias=False)
self.time_embedding = nn.Embedding(
num_time_buckets, structural_embedding_dim
)
self.role_embedding = nn.Embedding(3, structural_embedding_dim)
gate_input_dim = 2 + 2 * structural_embedding_dim
self.gate_mlp = nn.Sequential(
nn.LayerNorm(gate_input_dim),
nn.Linear(gate_input_dim, hidden_dim),
nn.GELU(),
nn.Linear(hidden_dim, 1),
)
# Direct gate output starts at 0.5 for every valid slot. Training uses
# a confidence-to-gate transition outside this module.
nn.init.zeros_(self.gate_mlp[-1].weight)
nn.init.zeros_(self.gate_mlp[-1].bias)
def forward(
self,
visual_memory_tokens: torch.Tensor,
current_visual_tokens: torch.Tensor,
current_token_valid_mask: torch.Tensor,
memory_confidence: torch.Tensor,
relative_frame_gaps: torch.Tensor,
slot_role_ids: torch.Tensor,
slot_valid_mask: torch.Tensor,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if visual_memory_tokens.ndim != 4:
raise ValueError("visual_memory_tokens must be [B,M,8,D]")
batch, slots, queries, _ = visual_memory_tokens.shape
if queries != 8:
raise ValueError("V9 dynamic gate requires 4+4 memory tokens")
if current_visual_tokens.ndim != 3:
raise ValueError("current_visual_tokens must be [B,P,D]")
valid = slot_valid_mask.to(
device=visual_memory_tokens.device, dtype=torch.bool
)
current_valid = current_token_valid_mask.to(
device=visual_memory_tokens.device, dtype=torch.bool
)
if valid.shape != (batch, slots):
raise ValueError("slot_valid_mask must be [B,M]")
if current_valid.shape != current_visual_tokens.shape[:2]:
raise ValueError("current_token_valid_mask must be [B,P]")
memory_projected = self.memory_projection(
self.memory_norm(visual_memory_tokens)
)
current_projected = self.current_projection(
self.current_norm(current_visual_tokens)
)
similarities = torch.einsum(
"bmqd,bpd->bmqp", memory_projected, current_projected
).float() / math.sqrt(self.projection_dim)
similarities = similarities.masked_fill(
~current_valid[:, None, None, :], -torch.inf
)
if (~current_valid.any(dim=-1)).any():
raise ValueError("every current image requires at least one patch")
retrieval = self.temperature * torch.logsumexp(
similarities / self.temperature, dim=-1
)
target_similarity = retrieval[:, :, :4].mean(dim=-1)
context_similarity = retrieval[:, :, 4:].mean(dim=-1)
target_context_margin = target_similarity - context_similarity
confidence = memory_confidence.to(
device=visual_memory_tokens.device, dtype=torch.float32
)
confidence = torch.where(
valid & torch.isfinite(confidence),
confidence.clamp(0.05, 0.95),
torch.full_like(confidence, 0.5),
)
confidence_logit = torch.logit(confidence)
time_ids = relative_time_bucket(
relative_frame_gaps.to(visual_memory_tokens.device),
self.time_embedding.num_embeddings,
)
time_features = self.time_embedding(time_ids)
roles = slot_role_ids.to(visual_memory_tokens.device)
if roles.ndim == 1:
roles = roles[None].expand(batch, -1)
role_features = self.role_embedding(roles)
gate_input = torch.cat(
(
target_context_margin.unsqueeze(-1),
confidence_logit.unsqueeze(-1),
time_features.float(),
role_features.float(),
),
dim=-1,
)
logits = self.gate_mlp(gate_input).squeeze(-1)
gate_scores = torch.sigmoid(logits)
gate_scores = gate_scores * valid.to(gate_scores.dtype)
diagnostics = {
"memory_gate_logits": logits,
"memory_gate_scores": gate_scores,
"memory_gate_target_similarity": target_similarity,
"memory_gate_context_similarity": context_similarity,
"memory_gate_target_context_margin": target_context_margin,
}
return gate_scores, diagnostics
class ScoreBiasedSelfAttention(nn.Module):
"""Self-attention whose keys receive a per-memory log-score prior."""
def __init__(
self,
dim: int,
num_heads: int,
dropout: float,
initial_score_lambda: float = 1.0,
max_score_lambda: float = 4.0,
) -> None:
super().__init__()
if dim % num_heads:
raise ValueError("dim must be divisible by num_heads")
self.dim = dim
self.num_heads = num_heads
self.head_dim = dim // num_heads
self.max_score_lambda = float(max_score_lambda)
initial_raw = _inverse_softplus(initial_score_lambda)
self.raw_score_lambda = nn.Parameter(
torch.full((num_heads,), initial_raw)
)
self.qkv = nn.Linear(dim, 3 * dim)
self.output = nn.Linear(dim, dim)
self.attention_dropout = float(dropout)
self.output_dropout = nn.Dropout(dropout)
@property
def score_lambda(self) -> torch.Tensor:
return F.softplus(self.raw_score_lambda).clamp_max(
self.max_score_lambda
)
def forward(
self,
x: torch.Tensor,
token_valid_mask: torch.Tensor,
token_scores: torch.Tensor,
return_attention: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
batch, length, dim = x.shape
valid = token_valid_mask.to(device=x.device, dtype=torch.bool)
scores = token_scores.to(device=x.device, dtype=torch.float32)
if valid.shape != (batch, length) or scores.shape != (batch, length):
raise ValueError("token masks and scores must have shape [B, L]")
qkv = self.qkv(x).view(
batch, length, 3, self.num_heads, self.head_dim
)
q, k, v = qkv.unbind(dim=2)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
logits = torch.matmul(q.float(), k.float().transpose(-2, -1))
logits = logits / math.sqrt(self.head_dim)
log_scores = torch.log(scores.clamp_min(1e-6))
score_bias = (
self.score_lambda.float()[None, :, None, None]
* log_scores[:, None, None, :]
)
logits = logits + score_bias
logits = logits.masked_fill(~valid[:, None, None, :], -torch.inf)
no_valid_keys = ~valid.any(dim=-1)
if no_valid_keys.any():
logits = logits.clone()
logits[no_valid_keys, :, :, 0] = 0.0
attention = torch.softmax(logits, dim=-1)
attention = torch.nan_to_num(attention, nan=0.0)
attention = attention * valid[:, None, None, :].to(
attention.dtype
)
attention = F.dropout(
attention,
p=self.attention_dropout,
training=self.training,
)
update = torch.matmul(attention.to(v.dtype), v)
update = update.transpose(1, 2).reshape(batch, length, dim)
update = self.output_dropout(self.output(update))
update = update * valid.unsqueeze(-1).to(update.dtype)
return update, attention if return_attention else None
class TemporalMemoryMixer(nn.Module):
"""Mix all 48 tokens with score competition and a score-gated residual."""
def __init__(
self,
dim: int = 2048,
num_slots: int = 6,
tokens_per_slot: int = 8,
num_heads: int = 16,
ffn_dim: int = 8192,
dropout: float = 0.1,
initial_score_lambda: float = 1.0,
initial_gate_beta: float = 1.0,
) -> None:
super().__init__()
self.num_slots = int(num_slots)
self.tokens_per_slot = int(tokens_per_slot)
self.input_norm = nn.LayerNorm(dim)
self.attention = ScoreBiasedSelfAttention(
dim,
num_heads,
dropout,
initial_score_lambda=initial_score_lambda,
)
self.ffn_norm = nn.LayerNorm(dim)
self.ffn = SwiGLU(dim, ffn_dim, dropout)
self.dropout = nn.Dropout(dropout)
self.raw_gate_beta = nn.Parameter(
torch.tensor(_inverse_softplus(initial_gate_beta))
)
self.output_norm = nn.LayerNorm(dim)
@property
def gate_beta(self) -> torch.Tensor:
return F.softplus(self.raw_gate_beta)
def forward(
self,
tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if tokens.ndim != 4:
raise ValueError("tokens must have shape [B, M, Q, D]")
batch, slots, per_slot, dim = tokens.shape
if (slots, per_slot) != (self.num_slots, self.tokens_per_slot):
raise ValueError(
f"expected [B,{self.num_slots},{self.tokens_per_slot},D], "
f"got {tuple(tokens.shape)}"
)
valid = slot_valid_mask.to(device=tokens.device, dtype=torch.bool)
scores = memory_scores.to(device=tokens.device, dtype=torch.float32)
if valid.shape != (batch, slots) or scores.shape != (batch, slots):
raise ValueError("slot_valid_mask and memory_scores must be [B, M]")
# Scores are reliability priors. Invalid and non-finite values cannot
# leak into attention or the gated residual.
scores = torch.where(
valid & torch.isfinite(scores),
scores.clamp(1e-6, 1.0),
torch.zeros_like(scores),
)
token_valid = valid.repeat_interleave(per_slot, dim=1)
token_scores = scores.repeat_interleave(per_slot, dim=1)
flat = tokens.reshape(batch, slots * per_slot, dim)
update, attention = self.attention(
self.input_norm(flat),
token_valid,
token_scores,
return_attention=return_diagnostics,
)
slot_gates = scores.pow(self.gate_beta)
token_gates = slot_gates.repeat_interleave(per_slot, dim=1)
# Requested score-gated residual: low-confidence memories preserve less
# of their original per-frame feature before cross-frame mixing.
mixed = token_gates.unsqueeze(-1).to(flat.dtype) * flat + update
mixed = mixed + self.dropout(self.ffn(self.ffn_norm(mixed)))
mixed = self.output_norm(mixed)
mixed = mixed * token_valid.unsqueeze(-1).to(mixed.dtype)
diagnostics: Dict[str, torch.Tensor] = {}
if return_diagnostics:
diagnostics = {
"attention": attention,
"slot_gates": slot_gates,
"score_lambda": self.attention.score_lambda,
"gate_beta": self.gate_beta,
"token_valid_mask": token_valid,
}
return mixed, diagnostics
class CurrentLanguageQueryEncoder(nn.Module):
"""Build target queries from the current image with weak language context."""
def __init__(
self,
dim: int,
num_queries: int,
num_heads: int,
dropout: float,
language_gate_init: float = 0.1,
) -> None:
super().__init__()
self.num_queries = int(num_queries)
self.queries = nn.Parameter(torch.empty(num_queries, dim))
self.query_norm = nn.LayerNorm(dim)
self.current_norm = nn.LayerNorm(dim)
self.current_cross_attention = nn.MultiheadAttention(
dim,
num_heads,
dropout=dropout,
batch_first=True,
)
self.language_norm = nn.LayerNorm(dim)
self.language_projection = nn.Linear(dim, dim)
self.raw_language_gate = nn.Parameter(
torch.tensor(_inverse_sigmoid(language_gate_init))
)
self.dropout = nn.Dropout(dropout)
self.output_norm = nn.LayerNorm(dim)
nn.init.trunc_normal_(self.queries, std=0.02)
@property
def language_gate(self) -> torch.Tensor:
return torch.sigmoid(self.raw_language_gate)
def forward(
self,
current_visual_tokens: torch.Tensor,
current_token_valid_mask: torch.Tensor,
language_tokens: Optional[torch.Tensor] = None,
language_token_mask: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
if current_visual_tokens.ndim != 3:
raise ValueError(
"current_visual_tokens must have shape [B, P_current, D]"
)
batch, current_length, _ = current_visual_tokens.shape
current_valid = current_token_valid_mask.to(
device=current_visual_tokens.device,
dtype=torch.bool,
)
if current_valid.shape != (batch, current_length):
raise ValueError(
"current_token_valid_mask must have shape [B, P_current]"
)
# Current images are required by TrackingAnything, but retain a stable
# sentinel path so malformed/all-masked diagnostic batches do not
# produce NaNs inside MultiheadAttention.
current_context = current_visual_tokens
current_padding_mask = ~current_valid
all_current_invalid = ~current_valid.any(dim=-1)
if all_current_invalid.any():
current_context = current_context.clone()
current_padding_mask = current_padding_mask.clone()
current_context[all_current_invalid, 0] = 0
current_padding_mask[all_current_invalid, 0] = False
queries = self.queries.to(dtype=current_visual_tokens.dtype)
queries = queries.unsqueeze(0).expand(batch, -1, -1)
update = self.current_cross_attention(
self.query_norm(queries),
self.current_norm(current_context),
self.current_norm(current_context),
key_padding_mask=current_padding_mask,
need_weights=False,
)[0]
queries = queries + self.dropout(update)
language_present = torch.zeros(
batch,
device=current_visual_tokens.device,
dtype=torch.bool,
)
if language_tokens is not None:
if language_tokens.ndim != 3 or language_tokens.shape[0] != batch:
raise ValueError("language_tokens must have shape [B, L, D]")
if language_tokens.shape[-1] != current_visual_tokens.shape[-1]:
raise ValueError(
"language and current-image token widths must match"
)
if language_token_mask is None:
language_valid = torch.ones(
language_tokens.shape[:2],
device=language_tokens.device,
dtype=torch.bool,
)
else:
language_valid = language_token_mask.to(
device=language_tokens.device,
dtype=torch.bool,
)
if language_valid.shape != language_tokens.shape[:2]:
raise ValueError(
"language_token_mask must have shape [B, L]"
)
language_present = language_valid.any(dim=-1)
weights = language_valid.unsqueeze(-1).to(language_tokens.dtype)
pooled_language = (language_tokens * weights).sum(dim=1)
pooled_language = pooled_language / weights.sum(dim=1).clamp_min(1)
language_update = self.language_projection(
self.language_norm(pooled_language)
)
effective_gate = (
self.language_gate.to(language_update.dtype)
* language_present.to(language_update.dtype)
)
queries = queries + (
effective_gate[:, None, None]
* language_update[:, None, :]
)
queries = self.output_norm(queries)
queries = queries * (~all_current_invalid)[:, None, None].to(
queries.dtype
)
return queries, (
self.language_gate
* language_present.to(self.language_gate.dtype)
)
class CurrentMemoryRelevance(nn.Module):
"""Score all 48 memory tokens using current queries and V6 score priors."""
def __init__(
self,
dim: int,
num_heads: int,
initial_static_score_lambda: float = 0.1,
max_static_score_lambda: float = 4.0,
) -> None:
super().__init__()
if dim % num_heads:
raise ValueError("dim must be divisible by num_heads")
self.num_heads = int(num_heads)
self.head_dim = dim // num_heads
self.max_static_score_lambda = float(max_static_score_lambda)
self.current_norm = nn.LayerNorm(dim)
self.memory_norm = nn.LayerNorm(dim)
self.query_projection = nn.Linear(dim, dim)
self.key_projection = nn.Linear(dim, dim)
self.raw_static_score_lambda = nn.Parameter(
torch.full(
(num_heads,),
_inverse_softplus(initial_static_score_lambda),
)
)
@property
def static_score_lambda(self) -> torch.Tensor:
return F.softplus(self.raw_static_score_lambda).clamp_max(
self.max_static_score_lambda
)
def forward(
self,
current_queries: torch.Tensor,
memory_tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
return_attention: bool = False,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
if current_queries.ndim != 3 or memory_tokens.ndim != 4:
raise ValueError(
"current queries must be [B,Qc,D] and memories [B,M,Qm,D]"
)
batch, slots, per_slot, dim = memory_tokens.shape
if current_queries.shape[0] != batch or current_queries.shape[-1] != dim:
raise ValueError("current-query and memory shapes are incompatible")
valid = slot_valid_mask.to(
device=memory_tokens.device,
dtype=torch.bool,
)
scores = memory_scores.to(
device=memory_tokens.device,
dtype=torch.float32,
)
if valid.shape != (batch, slots) or scores.shape != (batch, slots):
raise ValueError("slot_valid_mask and memory_scores must be [B, M]")
scores = torch.where(
valid & torch.isfinite(scores),
scores.clamp(1e-6, 1.0),
torch.zeros_like(scores),
)
memory_flat = memory_tokens.reshape(batch, slots * per_slot, dim)
query = self.query_projection(
self.current_norm(current_queries)
).view(
batch,
current_queries.shape[1],
self.num_heads,
self.head_dim,
)
key = self.key_projection(self.memory_norm(memory_flat)).view(
batch,
slots * per_slot,
self.num_heads,
self.head_dim,
)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
logits = torch.matmul(
query.float(),
key.float().transpose(-2, -1),
) / math.sqrt(self.head_dim)
token_valid = valid.repeat_interleave(per_slot, dim=1)
token_scores = scores.repeat_interleave(per_slot, dim=1)
static_bias = (
self.static_score_lambda.float()[None, :, None, None]
* torch.log(token_scores.clamp_min(1e-6))[:, None, None, :]
)
logits = logits + static_bias
logits = logits.masked_fill(
~token_valid[:, None, None, :],
-torch.inf,
)
no_valid_memories = ~token_valid.any(dim=-1)
if no_valid_memories.any():
logits = logits.clone()
logits[no_valid_memories, :, :, 0] = 0.0
attention_probabilities = torch.softmax(logits, dim=-1)
attention_probabilities = torch.nan_to_num(
attention_probabilities, nan=0.0
)
attention_probabilities = attention_probabilities * token_valid[:, None, None, :].to(
attention_probabilities.dtype
)
attention = attention_probabilities
# Attention mass is first averaged across current queries and heads,
# then summed over a frame's eight tokens. Every non-empty sample is
# therefore a probability distribution across its valid memory slots.
token_relevance = attention.mean(dim=(1, 2))
slot_relevance = token_relevance.reshape(
batch, slots, per_slot
).sum(dim=-1)
slot_relevance = slot_relevance * valid.to(slot_relevance.dtype)
normalizer = slot_relevance.sum(dim=-1, keepdim=True)
slot_relevance = torch.where(
normalizer > 0,
slot_relevance / normalizer.clamp_min(1e-12),
torch.zeros_like(slot_relevance),
)
return slot_relevance, attention if return_attention else None
class CurrentToMemoryMixer(nn.Module):
"""Current-conditioned one-layer temporal mixing over all 48 memories."""
def __init__(
self,
dim: int = 2048,
num_slots: int = 6,
tokens_per_slot: int = 8,
num_current_queries: int = 8,
num_heads: int = 16,
ffn_dim: int = 8192,
dropout: float = 0.1,
language_gate_init: float = 0.1,
initial_static_score_lambda: float = 0.1,
initial_dynamic_relevance_lambda: float = 1.0,
) -> None:
super().__init__()
self.num_slots = int(num_slots)
self.tokens_per_slot = int(tokens_per_slot)
self.current_query_encoder = CurrentLanguageQueryEncoder(
dim=dim,
num_queries=num_current_queries,
num_heads=num_heads,
dropout=dropout,
language_gate_init=language_gate_init,
)
self.current_memory_relevance = CurrentMemoryRelevance(
dim=dim,
num_heads=num_heads,
initial_static_score_lambda=initial_static_score_lambda,
)
self.input_norm = nn.LayerNorm(dim)
# The current-memory relevance already combines visual match and the
# static V6 score. This attention therefore consumes only that
# combined dynamic prior and does not apply the static score twice.
self.attention = ScoreBiasedSelfAttention(
dim,
num_heads,
dropout,
initial_score_lambda=initial_dynamic_relevance_lambda,
)
self.ffn_norm = nn.LayerNorm(dim)
self.ffn = SwiGLU(dim, ffn_dim, dropout)
self.dropout = nn.Dropout(dropout)
self.output_norm = nn.LayerNorm(dim)
def forward(
self,
tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
current_visual_tokens: torch.Tensor,
current_token_valid_mask: torch.Tensor,
language_tokens: Optional[torch.Tensor] = None,
language_token_mask: Optional[torch.Tensor] = None,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if tokens.ndim != 4:
raise ValueError("tokens must have shape [B, M, Q, D]")
batch, slots, per_slot, dim = tokens.shape
if (slots, per_slot) != (self.num_slots, self.tokens_per_slot):
raise ValueError(
f"expected [B,{self.num_slots},{self.tokens_per_slot},D], "
f"got {tuple(tokens.shape)}"
)
valid = slot_valid_mask.to(device=tokens.device, dtype=torch.bool)
if valid.shape != (batch, slots):
raise ValueError("slot_valid_mask must have shape [B, M]")
current_queries, language_gate = self.current_query_encoder(
current_visual_tokens=current_visual_tokens,
current_token_valid_mask=current_token_valid_mask,
language_tokens=language_tokens,
language_token_mask=language_token_mask,
)
slot_relevance, current_memory_attention = (
self.current_memory_relevance(
current_queries=current_queries,
memory_tokens=tokens,
slot_valid_mask=valid,
memory_scores=memory_scores,
return_attention=return_diagnostics,
)
)
token_valid = valid.repeat_interleave(per_slot, dim=1)
dynamic_token_scores = slot_relevance.repeat_interleave(
per_slot, dim=1
)
flat = tokens.reshape(batch, slots * per_slot, dim)
update, temporal_attention = self.attention(
self.input_norm(flat),
token_valid,
dynamic_token_scores,
return_attention=return_diagnostics,
)
# Preserve every valid per-frame token. Relevance controls which keys
# are read during the update instead of destructively attenuating the
# original token as the legacy score-gated residual did.
mixed = flat + update
mixed = mixed + self.dropout(self.ffn(self.ffn_norm(mixed)))
mixed = self.output_norm(mixed)
mixed = mixed * token_valid.unsqueeze(-1).to(mixed.dtype)
diagnostics: Dict[str, torch.Tensor] = {}
if return_diagnostics:
diagnostics = {
"attention": temporal_attention,
"current_memory_attention": current_memory_attention,
"slot_dynamic_relevance": slot_relevance,
"current_queries": current_queries,
"static_score_lambda": (
self.current_memory_relevance.static_score_lambda
),
"dynamic_relevance_lambda": self.attention.score_lambda,
"language_gate": language_gate,
"token_valid_mask": token_valid,
}
return mixed, diagnostics
class TokenLevelCurrentToMemoryMixer(nn.Module):
"""V4 token-level current-to-memory fusion with one shared slot reader.
The same Cross-Attention and FFN parameters are applied to all six memory
slots by folding the slot axis into the batch axis. Current queries read
each eight-token memory independently, preserving token-level evidence
instead of reducing it to one scalar relevance value per frame. A second
one-layer mixer then lets all 48 tokens compete under the static V6 score
prior.
"""
def __init__(
self,
dim: int = 2048,
num_slots: int = 6,
tokens_per_slot: int = 8,
num_current_queries: int = 8,
num_heads: int = 16,
temporal_ffn_dim: int = 8192,
cross_ffn_dim: int = 4096,
dropout: float = 0.1,
language_gate_init: float = 0.1,
cross_residual_gate_init: float = 0.1,
initial_score_lambda: float = 0.1,
) -> None:
super().__init__()
if num_current_queries != tokens_per_slot:
raise ValueError(
"V4 requires one current query per compressed memory token"
)
self.num_slots = int(num_slots)
self.tokens_per_slot = int(tokens_per_slot)
self.num_current_queries = int(num_current_queries)
self.num_heads = int(num_heads)
self.current_query_encoder = CurrentLanguageQueryEncoder(
dim=dim,
num_queries=num_current_queries,
num_heads=num_heads,
dropout=dropout,
language_gate_init=language_gate_init,
)
# This is deliberately one module, not a six-element ModuleList. The
# [B,M] axes are flattened below so every slot uses identical weights.
self.current_query_norm = nn.LayerNorm(dim)
self.slot_memory_norm = nn.LayerNorm(dim)
self.slot_cross_attention = nn.MultiheadAttention(
dim,
num_heads,
dropout=dropout,
batch_first=True,
)
self.slot_cross_ffn_norm = nn.LayerNorm(dim)
self.slot_cross_ffn = SwiGLU(dim, cross_ffn_dim, dropout)
self.cross_dropout = nn.Dropout(dropout)
self.raw_cross_residual_gate = nn.Parameter(
torch.tensor(_inverse_sigmoid(cross_residual_gate_init))
)
# Inference-only diagnostic. Keeping this buffer empty preserves the
# trained V4 path exactly; it is non-persistent so old and new
# checkpoints have identical state-dict contracts.
self.register_buffer(
"_runtime_cross_readout_permutation",
torch.empty(0, dtype=torch.long),
persistent=False,
)
self.temporal_input_norm = nn.LayerNorm(dim)
self.temporal_attention = ScoreBiasedSelfAttention(
dim,
num_heads,
dropout,
initial_score_lambda=initial_score_lambda,
)
self.temporal_ffn_norm = nn.LayerNorm(dim)
self.temporal_ffn = SwiGLU(dim, temporal_ffn_dim, dropout)
self.temporal_dropout = nn.Dropout(dropout)
self.output_norm = nn.LayerNorm(dim)
@property
def cross_residual_gate(self) -> torch.Tensor:
return torch.sigmoid(self.raw_cross_residual_gate)
def set_runtime_cross_readout_permutation(
self,
permutation: Optional[Sequence[int]],
) -> None:
"""Set a fixed inference-only permutation of the eight readouts.
``None`` disables the intervention. A non-empty permutation must be
a bijection over the current-query/readout axis. It is deliberately
not serialized in checkpoints because this is an evaluation control,
not a learned model parameter.
"""
if permutation is None:
value = torch.empty(
0,
dtype=torch.long,
device=self.raw_cross_residual_gate.device,
)
else:
values = tuple(int(index) for index in permutation)
expected = tuple(range(self.num_current_queries))
if len(values) != self.num_current_queries:
raise ValueError(
"V4 readout permutation must have "
f"{self.num_current_queries} entries"
)
if tuple(sorted(values)) != expected:
raise ValueError(
"V4 readout permutation must be a bijection over "
f"{expected}"
)
value = torch.tensor(
values,
dtype=torch.long,
device=self.raw_cross_residual_gate.device,
)
self._runtime_cross_readout_permutation = value
def _apply_runtime_cross_readout_permutation(
self,
cross_readout: torch.Tensor,
) -> torch.Tensor:
permutation = self._runtime_cross_readout_permutation
if permutation.numel() == 0:
return cross_readout
return cross_readout.index_select(
2,
permutation.to(device=cross_readout.device),
)
def forward(
self,
tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
current_visual_tokens: torch.Tensor,
current_token_valid_mask: torch.Tensor,
language_tokens: Optional[torch.Tensor] = None,
language_token_mask: Optional[torch.Tensor] = None,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if tokens.ndim != 4:
raise ValueError("tokens must have shape [B, M, Q, D]")
batch, slots, per_slot, dim = tokens.shape
if (slots, per_slot) != (self.num_slots, self.tokens_per_slot):
raise ValueError(
f"expected [B,{self.num_slots},{self.tokens_per_slot},D], "
f"got {tuple(tokens.shape)}"
)
valid = slot_valid_mask.to(device=tokens.device, dtype=torch.bool)
scores = memory_scores.to(device=tokens.device, dtype=torch.float32)
if valid.shape != (batch, slots) or scores.shape != (batch, slots):
raise ValueError("slot_valid_mask and memory_scores must be [B, M]")
scores = torch.where(
valid & torch.isfinite(scores),
scores.clamp(1e-6, 1.0),
torch.zeros_like(scores),
)
current_queries, language_gate = self.current_query_encoder(
current_visual_tokens=current_visual_tokens,
current_token_valid_mask=current_token_valid_mask,
language_tokens=language_tokens,
language_token_mask=language_token_mask,
)
# Six independent reads with one shared parameter set. Invalid slots
# receive a zero sentinel key to avoid all-masked MHA rows, and are
# exactly zeroed again after fusion.
flat_current = current_queries[:, None, :, :].expand(
batch, slots, self.num_current_queries, dim
).reshape(batch * slots, self.num_current_queries, dim)
flat_memory = tokens.reshape(batch * slots, per_slot, dim)
flat_valid = valid.reshape(batch * slots)
key_padding_mask = (~flat_valid)[:, None].expand(
batch * slots, per_slot
)
if (~flat_valid).any():
flat_memory = flat_memory.clone()
key_padding_mask = key_padding_mask.clone()
flat_memory[~flat_valid, 0] = 0
key_padding_mask[~flat_valid, 0] = False
cross_update, slot_cross_attention = self.slot_cross_attention(
self.current_query_norm(flat_current),
self.slot_memory_norm(flat_memory),
self.slot_memory_norm(flat_memory),
key_padding_mask=key_padding_mask,
need_weights=return_diagnostics,
average_attn_weights=False,
)
cross_readout = cross_update + self.cross_dropout(
self.slot_cross_ffn(
self.slot_cross_ffn_norm(cross_update)
)
)
cross_readout = cross_readout.reshape(
batch, slots, per_slot, dim
)
cross_readout = self._apply_runtime_cross_readout_permutation(
cross_readout
)
conditioned = tokens + (
self.cross_residual_gate.to(tokens.dtype) * cross_readout
)
conditioned = conditioned * valid[:, :, None, None].to(
conditioned.dtype
)
token_valid = valid.repeat_interleave(per_slot, dim=1)
token_scores = scores.repeat_interleave(per_slot, dim=1)
flat = conditioned.reshape(batch, slots * per_slot, dim)
temporal_update, temporal_attention = self.temporal_attention(
self.temporal_input_norm(flat),
token_valid,
token_scores,
return_attention=return_diagnostics,
)
# V4 keeps a normal residual. The score affects which memory keys are
# read, but it does not destructively scale the original token stream.
mixed = flat + temporal_update
mixed = mixed + self.temporal_dropout(
self.temporal_ffn(self.temporal_ffn_norm(mixed))
)
mixed = self.output_norm(mixed)
mixed = mixed * token_valid.unsqueeze(-1).to(mixed.dtype)
diagnostics: Dict[str, torch.Tensor] = {}
if return_diagnostics:
diagnostics = {
"attention": temporal_attention,
"slot_cross_attention": slot_cross_attention.reshape(
batch,
slots,
self.num_heads,
self.num_current_queries,
per_slot,
),
"current_queries": current_queries,
"cross_residual_gate": self.cross_residual_gate,
"score_lambda": self.temporal_attention.score_lambda,
"language_gate": language_gate,
"token_valid_mask": token_valid,
}
if self._runtime_cross_readout_permutation.numel():
diagnostics["cross_readout_permutation"] = (
self._runtime_cross_readout_permutation
)
return mixed, diagnostics
class TokenAlignedCurrentToMemoryMixer(nn.Module):
"""V5 current-conditioned memory mixing without index-wise readout fusion.
Eight current queries summarize the current image and weak language
context. Every compressed memory token then queries those current
summaries, so the cross-attention output is aligned with the memory token
that receives it. The same cross-attention parameters are shared by all
six slots. A small learned residual gate keeps the new conditioning path
non-destructive, after which all 48 tokens compete in one score-biased
temporal self-attention layer. The V6 score is used exactly once.
"""
def __init__(
self,
dim: int = 2048,
num_slots: int = 6,
tokens_per_slot: int = 8,
num_current_queries: int = 8,
num_heads: int = 16,
ffn_dim: int = 8192,
dropout: float = 0.1,
language_gate_init: float = 0.1,
current_residual_gate_init: float = 0.05,
initial_score_lambda: float = 0.1,
max_score_lambda: float = 1.0,
) -> None:
super().__init__()
if not 0.0 < current_residual_gate_init < 1.0:
raise ValueError(
"current_residual_gate_init must be strictly between zero "
"and one"
)
self.num_slots = int(num_slots)
self.tokens_per_slot = int(tokens_per_slot)
self.num_current_queries = int(num_current_queries)
self.num_heads = int(num_heads)
self.current_query_encoder = CurrentLanguageQueryEncoder(
dim=dim,
num_queries=num_current_queries,
num_heads=num_heads,
dropout=dropout,
language_gate_init=language_gate_init,
)
# One shared reader for all six slots. Memory tokens are the queries,
# so each output is semantically aligned with the token receiving the
# residual; V4's unsupported query-index/token-index correspondence is
# deliberately absent.
self.memory_query_norm = nn.LayerNorm(dim)
self.current_context_norm = nn.LayerNorm(dim)
self.current_cross_attention = nn.MultiheadAttention(
dim,
num_heads,
dropout=dropout,
batch_first=True,
)
self.cross_dropout = nn.Dropout(dropout)
self.raw_current_residual_gate = nn.Parameter(
torch.tensor(_inverse_sigmoid(current_residual_gate_init))
)
self.temporal_input_norm = nn.LayerNorm(dim)
self.temporal_attention = ScoreBiasedSelfAttention(
dim,
num_heads,
dropout,
initial_score_lambda=initial_score_lambda,
max_score_lambda=max_score_lambda,
)
self.temporal_ffn_norm = nn.LayerNorm(dim)
self.temporal_ffn = SwiGLU(dim, ffn_dim, dropout)
self.temporal_dropout = nn.Dropout(dropout)
self.output_norm = nn.LayerNorm(dim)
@property
def current_residual_gate(self) -> torch.Tensor:
return torch.sigmoid(self.raw_current_residual_gate)
def forward(
self,
tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
current_visual_tokens: torch.Tensor,
current_token_valid_mask: torch.Tensor,
language_tokens: Optional[torch.Tensor] = None,
language_token_mask: Optional[torch.Tensor] = None,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if tokens.ndim != 4:
raise ValueError("tokens must have shape [B, M, Q, D]")
batch, slots, per_slot, dim = tokens.shape
if (slots, per_slot) != (self.num_slots, self.tokens_per_slot):
raise ValueError(
f"expected [B,{self.num_slots},{self.tokens_per_slot},D], "
f"got {tuple(tokens.shape)}"
)
valid = slot_valid_mask.to(device=tokens.device, dtype=torch.bool)
scores = memory_scores.to(device=tokens.device, dtype=torch.float32)
if valid.shape != (batch, slots) or scores.shape != (batch, slots):
raise ValueError("slot_valid_mask and memory_scores must be [B, M]")
scores = torch.where(
valid & torch.isfinite(scores),
scores.clamp(1e-6, 1.0),
torch.zeros_like(scores),
)
current_queries, language_gate = self.current_query_encoder(
current_visual_tokens=current_visual_tokens,
current_token_valid_mask=current_token_valid_mask,
language_tokens=language_tokens,
language_token_mask=language_token_mask,
)
current_available = current_token_valid_mask.to(
device=tokens.device,
dtype=torch.bool,
).any(dim=-1)
flat_memory = tokens.reshape(batch * slots, per_slot, dim)
flat_current = current_queries[:, None, :, :].expand(
batch,
slots,
self.num_current_queries,
dim,
).reshape(batch * slots, self.num_current_queries, dim)
cross_update, current_cross_attention = self.current_cross_attention(
self.memory_query_norm(flat_memory),
self.current_context_norm(flat_current),
self.current_context_norm(flat_current),
need_weights=return_diagnostics,
average_attn_weights=False,
)
cross_update = cross_update.reshape(batch, slots, per_slot, dim)
cross_update = cross_update * (
valid & current_available[:, None]
)[:, :, None, None].to(cross_update.dtype)
conditioned = tokens + (
self.current_residual_gate.to(tokens.dtype)
* self.cross_dropout(cross_update)
)
conditioned = conditioned * valid[:, :, None, None].to(
conditioned.dtype
)
token_valid = valid.repeat_interleave(per_slot, dim=1)
token_scores = scores.repeat_interleave(per_slot, dim=1)
flat = conditioned.reshape(batch, slots * per_slot, dim)
temporal_update, temporal_attention = self.temporal_attention(
self.temporal_input_norm(flat),
token_valid,
token_scores,
return_attention=return_diagnostics,
)
mixed = flat + temporal_update
mixed = mixed + self.temporal_dropout(
self.temporal_ffn(self.temporal_ffn_norm(mixed))
)
mixed = self.output_norm(mixed)
mixed = mixed * token_valid.unsqueeze(-1).to(mixed.dtype)
diagnostics: Dict[str, torch.Tensor] = {}
if return_diagnostics:
diagnostics = {
"attention": temporal_attention,
"current_cross_attention": (
current_cross_attention.reshape(
batch,
slots,
self.num_heads,
per_slot,
self.num_current_queries,
)
),
"current_queries": current_queries,
"current_residual_gate": self.current_residual_gate,
"score_lambda": self.temporal_attention.score_lambda,
"language_gate": language_gate,
"token_valid_mask": token_valid,
}
return mixed, diagnostics
class LateMemoryBridge(nn.Module):
"""Inject a score-biased 48-token readout into ``<TRACK_QUERY>``.
The bridge is intentionally placed immediately before the language model.
Consequently autoregressive answer positions can causally attend to the
bridge-conditioned TRACK_QUERY, while the continuous box and presence
heads read the corresponding final TRACK_QUERY hidden state.
"""
def __init__(
self,
dim: int,
num_slots: int = 6,
tokens_per_slot: int = 8,
num_heads: int = 16,
dropout: float = 0.1,
residual_gate_init: float = 0.1,
residual_gate_max: float = 1.0,
score_std_floor: float = 0.08,
score_clip: float = 2.0,
max_score_prior_ratio: float = 3.0,
) -> None:
super().__init__()
if dim % num_heads:
raise ValueError("dim must be divisible by num_heads")
if max_score_prior_ratio <= 1.0:
raise ValueError("max_score_prior_ratio must exceed one")
if score_clip <= 0.0:
raise ValueError("score_clip must be positive")
if not 0.0 < residual_gate_max <= 1.0:
raise ValueError("residual_gate_max must be in (0, 1]")
if not 0.0 < residual_gate_init < residual_gate_max:
raise ValueError(
"residual_gate_init must be in (0, residual_gate_max)"
)
self.dim = int(dim)
self.num_slots = int(num_slots)
self.tokens_per_slot = int(tokens_per_slot)
self.num_heads = int(num_heads)
self.head_dim = self.dim // self.num_heads
self.score_std_floor = float(score_std_floor)
self.score_clip = float(score_clip)
self.max_score_prior_ratio = float(max_score_prior_ratio)
self.residual_gate_max = float(residual_gate_max)
# z lies in [-clip,+clip], so this fixed beta guarantees that the
# largest possible bias-only odds ratio is exactly max_prior_ratio.
self.register_buffer(
"score_bias_beta",
torch.tensor(
math.log(self.max_score_prior_ratio)
/ (2.0 * self.score_clip),
dtype=torch.float32,
),
persistent=True,
)
self.track_norm = nn.LayerNorm(dim)
self.memory_norm = nn.LayerNorm(dim)
self.query_projection = nn.Linear(dim, dim)
self.key_projection = nn.Linear(dim, dim)
self.value_projection = nn.Linear(dim, dim)
self.output_projection = nn.Linear(dim, dim)
self.attention_dropout = float(dropout)
self.output_dropout = nn.Dropout(dropout)
self.raw_residual_gate = nn.Parameter(
torch.tensor(
_inverse_sigmoid(residual_gate_init / residual_gate_max)
)
)
@property
def residual_gate(self) -> torch.Tensor:
return self.residual_gate_max * torch.sigmoid(
self.raw_residual_gate
)
def forward(
self,
track_embedding: torch.Tensor,
memory_tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
*,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if track_embedding.ndim != 2:
raise ValueError("track_embedding must have shape [B, D]")
if memory_tokens.ndim != 3:
raise ValueError("memory_tokens must have shape [B, 48, D]")
batch, length, dim = memory_tokens.shape
if track_embedding.shape != (batch, dim) or dim != self.dim:
raise ValueError("TRACK_QUERY and memory dimensions must match")
expected_length = self.num_slots * self.tokens_per_slot
if length != expected_length:
raise ValueError(f"expected exactly {expected_length} memory tokens")
valid_slots = slot_valid_mask.to(
device=memory_tokens.device,
dtype=torch.bool,
)
if valid_slots.shape != (batch, self.num_slots):
raise ValueError("slot_valid_mask must have shape [B, M]")
normalized_z = normalized_log_memory_scores(
memory_scores.to(memory_tokens.device),
valid_slots,
std_floor=self.score_std_floor,
clip_value=self.score_clip,
)
token_valid = valid_slots.repeat_interleave(
self.tokens_per_slot,
dim=1,
)
token_z = normalized_z.repeat_interleave(
self.tokens_per_slot,
dim=1,
)
query = self.query_projection(
self.track_norm(track_embedding)
).view(batch, self.num_heads, 1, self.head_dim)
normalized_memory = self.memory_norm(memory_tokens)
key = self.key_projection(normalized_memory).view(
batch, length, self.num_heads, self.head_dim
).transpose(1, 2)
value = self.value_projection(normalized_memory).view(
batch, length, self.num_heads, self.head_dim
).transpose(1, 2)
logits = torch.matmul(query.float(), key.float().transpose(-2, -1))
logits = logits / math.sqrt(self.head_dim)
score_bias = self.score_bias_beta.float() * token_z[:, None, None, :]
logits = logits + score_bias
logits = logits.masked_fill(
~token_valid[:, None, None, :],
-torch.inf,
)
no_valid_memory = ~token_valid.any(dim=-1)
if no_valid_memory.any():
logits = logits.clone()
logits[no_valid_memory, :, :, 0] = 0.0
attention_probabilities = torch.softmax(logits, dim=-1)
attention_probabilities = torch.nan_to_num(
attention_probabilities, nan=0.0
)
attention_probabilities = attention_probabilities * token_valid[:, None, None, :].to(
attention_probabilities.dtype
)
attention = attention_probabilities
attention = F.dropout(
attention,
p=self.attention_dropout,
training=self.training,
)
bridge = torch.matmul(attention.to(value.dtype), value)
bridge = bridge.transpose(1, 2).reshape(batch, dim)
bridge = self.output_dropout(self.output_projection(bridge))
bridge = bridge * (~no_valid_memory).unsqueeze(-1).to(bridge.dtype)
updated_track = track_embedding + self.residual_gate.to(
bridge.dtype
) * bridge
diagnostics: Dict[str, torch.Tensor] = {}
if return_diagnostics:
diagnostics = {
"late_bridge_attention": attention,
"late_bridge_attention_probabilities": attention_probabilities,
"late_bridge_readout": bridge,
"late_bridge_normalized_z": normalized_z,
"late_bridge_score_bias_beta": self.score_bias_beta,
"late_bridge_residual_gate": self.residual_gate,
"late_bridge_token_valid_mask": token_valid,
}
return updated_track, diagnostics
class DirectCurrentIndexedMemoryReadoutMixer(nn.Module):
"""Direct current-indexed retrieval from six dense memory slots.
Current queries independently read every slot through one shared
Cross-Attention module. Its values come from memory, so the resulting
readouts are passed directly to temporal mixing; they are never added to
the per-frame resampler tokens by an unsupported positional association.
The V6 bank score is applied exactly once, as a key bias in the following
score-biased self-attention layer.
"""
def __init__(
self,
dim: int = 2048,
num_slots: int = 6,
tokens_per_slot: int = 8,
num_current_queries: int = 8,
num_heads: int = 16,
temporal_ffn_dim: int = 8192,
readout_ffn_dim: int = 4096,
dropout: float = 0.1,
language_gate_init: float = 0.1,
initial_score_lambda: float = 0.1,
max_score_lambda: float = 1.0,
enable_score_residual: bool = False,
score_std_floor: float = 0.08,
score_clip: float = 2.0,
score_residual_gate: float = 0.1,
additive_score_residual_scale: float = 0.05,
temporal_layers: int = 1,
enable_patch_location_fusion: bool = False,
patch_gate_hidden_dim: int = 256,
patch_gate_init: float = 0.8,
patch_fusion_mode: str = "convex",
) -> None:
super().__init__()
if tokens_per_slot <= 0 or num_current_queries <= 0:
raise ValueError("memory and current query counts must be positive")
if readout_ffn_dim <= 0:
raise ValueError("readout_ffn_dim must be positive")
self.num_slots = int(num_slots)
self.tokens_per_slot = int(tokens_per_slot)
self.num_current_queries = int(num_current_queries)
self.num_heads = int(num_heads)
self.enable_score_residual = bool(enable_score_residual)
self.score_std_floor = float(score_std_floor)
self.score_clip = float(score_clip)
self.score_residual_gate = float(score_residual_gate)
self.additive_score_residual_scale = float(
additive_score_residual_scale
)
self.temporal_layers = int(temporal_layers)
self.enable_patch_location_fusion = bool(
enable_patch_location_fusion
)
self.patch_fusion_mode = str(patch_fusion_mode)
if self.temporal_layers not in (1, 2):
raise ValueError("temporal_layers must be 1 or 2")
if patch_gate_hidden_dim <= 0:
raise ValueError("V10 patch gate hidden dimension must be positive")
if not 0.0 < patch_gate_init < 1.0:
raise ValueError("V10 patch gate initialization must be in (0,1)")
if self.patch_fusion_mode not in ("convex", "residual"):
raise ValueError(
"V10 patch fusion mode must be 'convex' or 'residual'"
)
if self.score_std_floor <= 0.0 or self.score_clip <= 0.0:
raise ValueError("V7 score normalization constants must be positive")
if not 0.0 <= self.score_residual_gate <= 0.5:
raise ValueError("V7 score residual gate must be in [0, 0.5]")
if not 0.0 <= self.additive_score_residual_scale <= 0.3:
raise ValueError("V7 additive score scale must be in [0, 0.3]")
self.current_query_encoder = CurrentLanguageQueryEncoder(
dim=dim,
num_queries=num_current_queries,
num_heads=num_heads,
dropout=dropout,
language_gate_init=language_gate_init,
)
# The [B,M] axes are folded together, so all six slots use this one
# reader. Query positions are current-indexed and values are memory;
# the output is therefore a current-conditioned memory retrieval.
self.current_query_norm = nn.LayerNorm(dim)
self.memory_context_norm = nn.LayerNorm(dim)
self.memory_readout_attention = nn.MultiheadAttention(
dim,
num_heads,
dropout=dropout,
batch_first=True,
)
self.readout_ffn_norm = nn.LayerNorm(dim)
self.readout_ffn = SwiGLU(dim, readout_ffn_dim, dropout)
self.readout_dropout = nn.Dropout(dropout)
# V10 adds exactly one module after the original V7 identity readout.
# R_i is used as the query and the uncompressed current-image patch
# grid supplies keys/values. Consequently the returned C_i positions
# inherit R_i's query indexing and can be mixed pointwise without the
# unsupported Q_cur,j <-> M_i,j association avoided by V6/V7.
if self.enable_patch_location_fusion:
self.patch_query_norm = nn.LayerNorm(dim)
self.patch_context_norm = nn.LayerNorm(dim)
self.patch_location_attention = nn.MultiheadAttention(
dim,
num_heads,
dropout=dropout,
batch_first=True,
)
self.patch_gate_norm_identity = nn.LayerNorm(dim)
self.patch_gate_norm_location = nn.LayerNorm(dim)
self.patch_fusion_gate = nn.Sequential(
nn.Linear(2 * dim, int(patch_gate_hidden_dim)),
nn.GELU(),
nn.Linear(int(patch_gate_hidden_dim), 1),
)
# Zero weights make the configured scalar gate exact for all data.
# In convex mode A weights identity; in residual mode g weights
# only the additive patch-location update.
nn.init.zeros_(self.patch_fusion_gate[-1].weight)
nn.init.constant_(
self.patch_fusion_gate[-1].bias,
_inverse_sigmoid(float(patch_gate_init)),
)
self.temporal_input_norm = nn.LayerNorm(dim)
self.temporal_attention = ScoreBiasedSelfAttention(
dim,
num_heads,
dropout,
initial_score_lambda=initial_score_lambda,
max_score_lambda=max_score_lambda,
)
self.temporal_ffn_norm = nn.LayerNorm(dim)
self.temporal_ffn = SwiGLU(dim, temporal_ffn_dim, dropout)
self.temporal_dropout = nn.Dropout(dropout)
self.output_norm = nn.LayerNorm(dim)
# Preserve the trained V7 first-layer names exactly. The optional
# second block receives the same score prior, but does not repeat V7's
# score residual or introduce another Late Memory Bridge.
self.additional_temporal_input_norms = nn.ModuleList()
self.additional_temporal_attentions = nn.ModuleList()
self.additional_temporal_ffn_norms = nn.ModuleList()
self.additional_temporal_ffns = nn.ModuleList()
self.additional_temporal_dropouts = nn.ModuleList()
self.additional_temporal_output_norms = nn.ModuleList()
self.additional_score_directions = nn.ParameterList()
for _ in range(self.temporal_layers - 1):
self.additional_temporal_input_norms.append(nn.LayerNorm(dim))
self.additional_temporal_attentions.append(
ScoreBiasedSelfAttention(
dim,
num_heads,
dropout,
initial_score_lambda=initial_score_lambda,
max_score_lambda=max_score_lambda,
)
)
self.additional_temporal_ffn_norms.append(nn.LayerNorm(dim))
self.additional_temporal_ffns.append(
SwiGLU(dim, temporal_ffn_dim, dropout)
)
self.additional_temporal_dropouts.append(nn.Dropout(dropout))
self.additional_temporal_output_norms.append(nn.LayerNorm(dim))
if self.enable_score_residual:
direction = nn.Parameter(torch.empty(dim))
nn.init.trunc_normal_(direction, std=0.02)
self.additional_score_directions.append(direction)
if self.enable_score_residual:
self.score_direction = nn.Parameter(torch.empty(dim))
nn.init.trunc_normal_(self.score_direction, std=0.02)
def forward(
self,
tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
current_visual_tokens: torch.Tensor,
current_token_valid_mask: torch.Tensor,
language_tokens: Optional[torch.Tensor] = None,
language_token_mask: Optional[torch.Tensor] = None,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
if tokens.ndim != 4:
raise ValueError("tokens must have shape [B, M, Q, D]")
batch, slots, per_slot, dim = tokens.shape
if (slots, per_slot) != (self.num_slots, self.tokens_per_slot):
raise ValueError(
f"expected [B,{self.num_slots},{self.tokens_per_slot},D], "
f"got {tuple(tokens.shape)}"
)
valid = slot_valid_mask.to(device=tokens.device, dtype=torch.bool)
scores = memory_scores.to(device=tokens.device, dtype=torch.float32)
if valid.shape != (batch, slots) or scores.shape != (batch, slots):
raise ValueError("slot_valid_mask and memory_scores must be [B, M]")
scores = torch.where(
valid & torch.isfinite(scores),
scores.clamp(1e-6, 1.0),
torch.zeros_like(scores),
)
current_queries, language_gate = self.current_query_encoder(
current_visual_tokens=current_visual_tokens,
current_token_valid_mask=current_token_valid_mask,
language_tokens=language_tokens,
language_token_mask=language_token_mask,
)
current_available = current_token_valid_mask.to(
device=tokens.device,
dtype=torch.bool,
).any(dim=-1)
flat_current = current_queries[:, None, :, :].expand(
batch,
slots,
self.num_current_queries,
dim,
).reshape(batch * slots, self.num_current_queries, dim)
flat_memory = tokens.reshape(batch * slots, per_slot, dim)
flat_valid = (valid & current_available[:, None]).reshape(
batch * slots
)
memory_padding_mask = (~flat_valid)[:, None].expand(
batch * slots, per_slot
)
if (~flat_valid).any():
flat_memory = flat_memory.clone()
memory_padding_mask = memory_padding_mask.clone()
flat_memory[~flat_valid, 0] = 0
memory_padding_mask[~flat_valid, 0] = False
readout, readout_attention = self.memory_readout_attention(
self.current_query_norm(flat_current),
self.memory_context_norm(flat_memory),
self.memory_context_norm(flat_memory),
key_padding_mask=memory_padding_mask,
need_weights=return_diagnostics,
average_attn_weights=False,
)
# This is an FFN residual inside the readout stream. There is
# deliberately no M_j + R_j or Q_cur,j + R_j fusion.
readout = readout + self.readout_dropout(
self.readout_ffn(self.readout_ffn_norm(readout))
)
identity_readout = readout
patch_readout = None
patch_attention = None
patch_fusion_gate = None
if self.enable_patch_location_fusion:
current_context = current_visual_tokens[:, None, :, :].expand(
batch,
slots,
current_visual_tokens.shape[1],
dim,
).reshape(batch * slots, current_visual_tokens.shape[1], dim)
patch_valid = current_token_valid_mask.to(
device=tokens.device,
dtype=torch.bool,
)[:, None, :].expand(
batch,
slots,
current_visual_tokens.shape[1],
)
patch_valid = patch_valid & valid[:, :, None]
patch_valid = patch_valid.reshape(
batch * slots, current_visual_tokens.shape[1]
)
patch_padding_mask = ~patch_valid
patch_all_invalid = ~patch_valid.any(dim=-1)
if patch_all_invalid.any():
current_context = current_context.clone()
patch_padding_mask = patch_padding_mask.clone()
current_context[patch_all_invalid, 0] = 0
patch_padding_mask[patch_all_invalid, 0] = False
patch_readout, patch_attention = self.patch_location_attention(
self.patch_query_norm(readout),
self.patch_context_norm(current_context),
self.patch_context_norm(current_context),
key_padding_mask=patch_padding_mask,
need_weights=return_diagnostics,
average_attn_weights=False,
)
patch_readout = patch_readout * (
~patch_all_invalid
)[:, None, None].to(patch_readout.dtype)
gate_input = torch.cat(
(
self.patch_gate_norm_identity(readout),
self.patch_gate_norm_location(patch_readout),
),
dim=-1,
)
patch_fusion_gate = torch.sigmoid(
self.patch_fusion_gate(gate_input)
)
gate = patch_fusion_gate.to(readout.dtype)
if self.patch_fusion_mode == "convex":
readout = gate * readout + (1.0 - gate) * patch_readout
else:
# V10.1 preserves the complete V7 identity stream. Patch
# localization is an initially weak, learnable residual, so
# an unreliable current-patch match cannot attenuate R_i.
readout = readout + gate * patch_readout
readout = readout * (~patch_all_invalid)[:, None, None].to(
readout.dtype
)
readout = readout.reshape(
batch, slots, self.num_current_queries, dim
)
readout = readout * (
valid & current_available[:, None]
)[:, :, None, None].to(readout.dtype)
token_valid = valid.repeat_interleave(
self.num_current_queries, dim=1
)
token_valid = token_valid & current_available[:, None]
token_scores = scores.repeat_interleave(
self.num_current_queries, dim=1
)
flat = readout.reshape(
batch, slots * self.num_current_queries, dim
)
temporal_update, temporal_attention = self.temporal_attention(
self.temporal_input_norm(flat),
token_valid,
token_scores,
return_attention=return_diagnostics,
)
normalized_z = normalized_log_memory_scores(
scores,
valid,
std_floor=self.score_std_floor,
clip_value=self.score_clip,
)
token_z = normalized_z.repeat_interleave(
self.num_current_queries,
dim=1,
)
if self.enable_score_residual:
residual_gate = 1.0 + self.score_residual_gate * torch.tanh(
token_z
)
direction = self.score_direction / self.score_direction.float().pow(
2
).mean().clamp_min(1e-12).sqrt().to(self.score_direction.dtype)
mixed = (
residual_gate.unsqueeze(-1).to(flat.dtype) * flat
+ temporal_update
+ self.additive_score_residual_scale
* token_z.unsqueeze(-1).to(flat.dtype)
* direction[None, None, :].to(flat.dtype)
)
else:
residual_gate = torch.ones_like(token_z)
mixed = flat + temporal_update
mixed = mixed + self.temporal_dropout(
self.temporal_ffn(self.temporal_ffn_norm(mixed))
)
mixed = self.output_norm(mixed)
mixed = mixed * token_valid.unsqueeze(-1).to(mixed.dtype)
temporal_attentions = [temporal_attention]
temporal_residual_gates = [residual_gate]
for layer_index, (
input_norm,
attention_layer,
ffn_norm,
ffn,
layer_dropout,
layer_output_norm,
) in enumerate(
zip(
self.additional_temporal_input_norms,
self.additional_temporal_attentions,
self.additional_temporal_ffn_norms,
self.additional_temporal_ffns,
self.additional_temporal_dropouts,
self.additional_temporal_output_norms,
)
):
layer_update, layer_attention = attention_layer(
input_norm(mixed),
token_valid,
token_scores,
return_attention=return_diagnostics,
)
if self.enable_score_residual:
layer_residual_gate = (
1.0
+ self.score_residual_gate * torch.tanh(token_z)
)
layer_direction = self.additional_score_directions[
layer_index
]
layer_direction = layer_direction / layer_direction.float().pow(
2
).mean().clamp_min(1e-12).sqrt().to(layer_direction.dtype)
mixed = (
layer_residual_gate.unsqueeze(-1).to(mixed.dtype) * mixed
+ layer_update
+ self.additive_score_residual_scale
* token_z.unsqueeze(-1).to(mixed.dtype)
* layer_direction[None, None, :].to(mixed.dtype)
)
else:
layer_residual_gate = torch.ones_like(token_z)
mixed = mixed + layer_update
mixed = mixed + layer_dropout(ffn(ffn_norm(mixed)))
mixed = layer_output_norm(mixed)
mixed = mixed * token_valid.unsqueeze(-1).to(mixed.dtype)
temporal_attentions.append(layer_attention)
temporal_residual_gates.append(layer_residual_gate)
diagnostics: Dict[str, torch.Tensor] = {}
if return_diagnostics:
diagnostics = {
"attention": temporal_attention,
"memory_readout_attention": readout_attention.reshape(
batch,
slots,
self.num_heads,
self.num_current_queries,
per_slot,
),
"current_queries": current_queries,
"direct_memory_readouts": readout,
# Opt-in tensors used by controlled mixer ablations. They
# expose the two streams before residual fusion without
# changing the default V6 forward or checkpoint state.
"temporal_update": temporal_update,
"temporal_pre_ffn": flat + temporal_update,
"score_lambda": self.temporal_attention.score_lambda,
"temporal_attention_layers": torch.stack(
temporal_attentions, dim=1
),
"temporal_score_lambdas": torch.stack(
[
self.temporal_attention.score_lambda,
*[
layer.score_lambda
for layer in self.additional_temporal_attentions
],
],
dim=0,
),
"temporal_score_residual_gates": torch.stack(
temporal_residual_gates, dim=1
),
"normalized_score_z": normalized_z,
"score_residual_gate": residual_gate,
"additive_score_residual_scale": torch.as_tensor(
self.additive_score_residual_scale,
device=mixed.device,
),
"language_gate": language_gate,
"token_valid_mask": token_valid,
}
if self.enable_patch_location_fusion:
diagnostics.update(
{
"v10_identity_readouts": identity_readout.reshape(
batch,
slots,
self.num_current_queries,
dim,
),
"v10_fused_readouts": readout.reshape(
batch,
slots,
self.num_current_queries,
dim,
),
"patch_location_readouts": patch_readout.reshape(
batch,
slots,
self.num_current_queries,
dim,
),
"patch_location_attention": patch_attention.reshape(
batch,
slots,
self.num_heads,
self.num_current_queries,
current_visual_tokens.shape[1],
),
"patch_fusion_gate": patch_fusion_gate.reshape(
batch,
slots,
self.num_current_queries,
1,
),
}
)
return mixed, diagnostics
class MemoryTokenEncoder(nn.Module):
"""Encode dense per-frame memories and emit a fixed LLM token budget."""
def __init__(
self,
dim: int = 2048,
num_slots: int = 6,
num_queries: int = 8,
output_num_queries: Optional[int] = None,
num_time_buckets: int = 8,
num_heads: int = 16,
resampler_layers: int = 2,
ffn_dim: int = 8192,
dropout: float = 0.1,
mixer_version: str = "temporal_v2",
current_num_queries: int = 8,
language_gate_init: float = 0.1,
current_static_score_lambda_init: float = 0.1,
current_dynamic_relevance_lambda_init: float = 1.0,
v4_cross_ffn_dim: int = 4096,
v4_cross_residual_gate_init: float = 0.1,
v4_score_lambda_init: float = 0.1,
v5_current_residual_gate_init: float = 0.05,
v5_score_lambda_init: float = 0.1,
v5_score_lambda_max: float = 1.0,
v6_readout_ffn_dim: int = 4096,
v6_score_lambda_init: float = 0.1,
v6_score_lambda_max: float = 1.0,
v7_score_std_floor: float = 0.08,
v7_score_clip: float = 2.0,
v7_front_residual_gate: float = 0.1,
v7_additive_score_residual_scale: float = 0.05,
v7_temporal_layers: int = 1,
v7_late_bridge_residual_gate_init: float = 0.1,
v7_late_bridge_residual_gate_max: float = 1.0,
v7_late_bridge_max_score_prior_ratio: float = 3.0,
v10_patch_gate_hidden_dim: int = 256,
v10_patch_gate_init: float = 0.8,
v10_patch_fusion_mode: str = "convex",
v9_region_bias_clip: float = 4.0,
v9_region_bias_scale: float = 1.0,
v9_gate_projection_dim: int = 256,
v9_gate_hidden_dim: int = 256,
v9_gate_structural_embedding_dim: int = 32,
v9_gate_temperature: float = 0.07,
) -> None:
super().__init__()
self.num_slots = int(num_slots)
self.num_queries = int(num_queries)
self.output_num_queries = int(
num_queries if output_num_queries is None else output_num_queries
)
if self.num_queries <= 0 or self.output_num_queries <= 0:
raise ValueError("dense and output query counts must be positive")
self.num_time_buckets = int(num_time_buckets)
if mixer_version not in (
"temporal_v2",
"current_to_memory_v3",
"current_to_memory_v4",
"current_to_memory_v5",
"current_to_memory_v6",
"current_to_memory_v7",
"current_to_memory_v9",
"current_to_memory_v10",
"current_to_memory_v14",
):
raise ValueError(f"unsupported memory mixer version: {mixer_version}")
self.mixer_version = str(mixer_version)
self.target_region_embedding = nn.Parameter(torch.empty(dim))
self.context_region_embedding = nn.Parameter(torch.empty(dim))
# role ids: anchor, long-term, short-term
self.role_embedding = nn.Embedding(3, dim)
self.time_embedding = nn.Embedding(num_time_buckets, dim)
if self.mixer_version == "current_to_memory_v9":
self.resampler = TargetContextPerFrameResampler(
dim=dim,
num_queries=num_queries,
num_layers=resampler_layers,
num_heads=num_heads,
ffn_dim=ffn_dim,
dropout=dropout,
region_bias_clip=v9_region_bias_clip,
region_bias_scale=v9_region_bias_scale,
)
self.target_context_type_embedding = nn.Embedding(2, dim)
self.dynamic_memory_gate = DirectDynamicMemoryGate(
dim=dim,
projection_dim=v9_gate_projection_dim,
hidden_dim=v9_gate_hidden_dim,
num_time_buckets=num_time_buckets,
structural_embedding_dim=(
v9_gate_structural_embedding_dim
),
temperature=v9_gate_temperature,
)
else:
self.resampler = PerFrameResampler(
dim=dim,
num_queries=num_queries,
num_layers=resampler_layers,
num_heads=num_heads,
ffn_dim=ffn_dim,
dropout=dropout,
)
if self.mixer_version == "temporal_v2":
self.temporal_mixer = TemporalMemoryMixer(
dim=dim,
num_slots=num_slots,
tokens_per_slot=num_queries,
num_heads=num_heads,
ffn_dim=ffn_dim,
dropout=dropout,
)
elif self.mixer_version == "current_to_memory_v3":
self.current_to_memory_mixer = CurrentToMemoryMixer(
dim=dim,
num_slots=num_slots,
tokens_per_slot=num_queries,
num_current_queries=current_num_queries,
num_heads=num_heads,
ffn_dim=ffn_dim,
dropout=dropout,
language_gate_init=language_gate_init,
initial_static_score_lambda=(
current_static_score_lambda_init
),
initial_dynamic_relevance_lambda=(
current_dynamic_relevance_lambda_init
),
)
elif self.mixer_version == "current_to_memory_v4":
self.current_to_memory_mixer = TokenLevelCurrentToMemoryMixer(
dim=dim,
num_slots=num_slots,
tokens_per_slot=num_queries,
num_current_queries=current_num_queries,
num_heads=num_heads,
temporal_ffn_dim=ffn_dim,
cross_ffn_dim=v4_cross_ffn_dim,
dropout=dropout,
language_gate_init=language_gate_init,
cross_residual_gate_init=v4_cross_residual_gate_init,
initial_score_lambda=v4_score_lambda_init,
)
elif self.mixer_version == "current_to_memory_v5":
self.current_to_memory_mixer = TokenAlignedCurrentToMemoryMixer(
dim=dim,
num_slots=num_slots,
tokens_per_slot=num_queries,
num_current_queries=current_num_queries,
num_heads=num_heads,
ffn_dim=ffn_dim,
dropout=dropout,
language_gate_init=language_gate_init,
current_residual_gate_init=(
v5_current_residual_gate_init
),
initial_score_lambda=v5_score_lambda_init,
max_score_lambda=v5_score_lambda_max,
)
elif self.mixer_version == "current_to_memory_v6":
self.current_to_memory_mixer = (
DirectCurrentIndexedMemoryReadoutMixer(
dim=dim,
num_slots=num_slots,
tokens_per_slot=num_queries,
num_current_queries=current_num_queries,
num_heads=num_heads,
temporal_ffn_dim=ffn_dim,
readout_ffn_dim=v6_readout_ffn_dim,
dropout=dropout,
language_gate_init=language_gate_init,
initial_score_lambda=v6_score_lambda_init,
max_score_lambda=v6_score_lambda_max,
)
)
else:
self.current_to_memory_mixer = (
DirectCurrentIndexedMemoryReadoutMixer(
dim=dim,
num_slots=num_slots,
tokens_per_slot=num_queries,
num_current_queries=current_num_queries,
num_heads=num_heads,
temporal_ffn_dim=ffn_dim,
readout_ffn_dim=v6_readout_ffn_dim,
dropout=dropout,
language_gate_init=language_gate_init,
initial_score_lambda=v6_score_lambda_init,
max_score_lambda=v6_score_lambda_max,
enable_score_residual=True,
score_std_floor=v7_score_std_floor,
score_clip=v7_score_clip,
score_residual_gate=v7_front_residual_gate,
additive_score_residual_scale=(
v7_additive_score_residual_scale
),
temporal_layers=v7_temporal_layers,
enable_patch_location_fusion=(
self.mixer_version == "current_to_memory_v10"
),
patch_gate_hidden_dim=v10_patch_gate_hidden_dim,
patch_gate_init=v10_patch_gate_init,
patch_fusion_mode=v10_patch_fusion_mode,
)
)
self.late_memory_bridge = LateMemoryBridge(
dim=dim,
num_slots=num_slots,
tokens_per_slot=self.output_num_queries,
num_heads=num_heads,
dropout=dropout,
residual_gate_init=v7_late_bridge_residual_gate_init,
residual_gate_max=v7_late_bridge_residual_gate_max,
score_std_floor=v7_score_std_floor,
score_clip=v7_score_clip,
max_score_prior_ratio=(
v7_late_bridge_max_score_prior_ratio
),
)
nn.init.trunc_normal_(self.target_region_embedding, std=0.02)
nn.init.trunc_normal_(self.context_region_embedding, std=0.02)
nn.init.trunc_normal_(self.role_embedding.weight, std=0.02)
nn.init.trunc_normal_(self.time_embedding.weight, std=0.02)
if self.mixer_version == "current_to_memory_v9":
nn.init.trunc_normal_(
self.target_context_type_embedding.weight, std=0.02
)
if self.mixer_version in (
"current_to_memory_v7",
"current_to_memory_v9",
"current_to_memory_v10",
"current_to_memory_v14",
):
# Give the dedicated learnable score direction a semantic start:
# positive z initially moves tokens toward target rather than
# context. RMS normalization in the forward keeps its scale fixed.
initial_direction = (
self.target_region_embedding
- self.context_region_embedding
)
initial_direction = initial_direction / initial_direction.float().pow(
2
).mean().clamp_min(1e-12).sqrt().to(initial_direction.dtype)
with torch.no_grad():
self.current_to_memory_mixer.score_direction.copy_(
initial_direction
)
for direction in (
self.current_to_memory_mixer
.additional_score_directions
):
direction.copy_(initial_direction)
self.register_buffer(
"slot_role_ids",
torch.tensor([0, 1, 1, 2, 2, 2][:self.num_slots], dtype=torch.long),
persistent=True,
)
@property
def output_token_count(self) -> int:
return self.num_slots * self.output_num_queries
def initialize_additional_v7_temporal_blocks(self) -> None:
"""Initialize appended blocks from block one after loading V7.
Transformers initializes parameters absent from an older checkpoint
after the module constructor has run. Copy initialization keeps the
repeated module numerically stable and makes it a strict depth
ablation. The parameters are independent and diverge during training.
"""
if self.mixer_version not in (
"current_to_memory_v7",
"current_to_memory_v10",
"current_to_memory_v14",
):
return
mixer = self.current_to_memory_mixer
for index in range(len(mixer.additional_temporal_attentions)):
mixer.additional_temporal_input_norms[index].load_state_dict(
mixer.temporal_input_norm.state_dict()
)
mixer.additional_temporal_attentions[index].load_state_dict(
mixer.temporal_attention.state_dict()
)
mixer.additional_temporal_ffn_norms[index].load_state_dict(
mixer.temporal_ffn_norm.state_dict()
)
mixer.additional_temporal_ffns[index].load_state_dict(
mixer.temporal_ffn.state_dict()
)
mixer.additional_temporal_output_norms[index].load_state_dict(
mixer.output_norm.state_dict()
)
with torch.no_grad():
mixer.additional_score_directions[index].copy_(
mixer.score_direction
)
def forward(
self,
projected_visual_tokens: torch.Tensor,
region_coverage: torch.Tensor,
relative_frame_gaps: torch.Tensor,
memory_scores: torch.Tensor,
slot_valid_mask: torch.Tensor,
patch_valid_mask: Optional[torch.Tensor] = None,
current_visual_tokens: Optional[torch.Tensor] = None,
current_token_valid_mask: Optional[torch.Tensor] = None,
language_tokens: Optional[torch.Tensor] = None,
language_token_mask: Optional[torch.Tensor] = None,
memory_gate_blend: Optional[float | torch.Tensor] = None,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
batch, slots, patches, _ = projected_visual_tokens.shape
if slots != self.num_slots:
raise ValueError(f"expected {self.num_slots} memory slots")
coverage = region_coverage.to(
projected_visual_tokens.device, projected_visual_tokens.dtype
)
if coverage.shape != (batch, slots, patches):
raise ValueError("region_coverage must have shape [B, M, P]")
coverage = coverage.clamp(0.0, 1.0)
region = (
coverage.unsqueeze(-1) * self.target_region_embedding
+ (1.0 - coverage.unsqueeze(-1))
* self.context_region_embedding
)
role_ids = self.slot_role_ids.to(projected_visual_tokens.device)
role = self.role_embedding(role_ids)[None, :, None, :]
time_ids = relative_time_bucket(
relative_frame_gaps.to(projected_visual_tokens.device),
self.num_time_buckets,
)
time = self.time_embedding(time_ids)[:, :, None, :]
if self.mixer_version == "current_to_memory_v9":
# V9 visual matching sees region-aware appearance only. Time,
# slot role and target/context type are appended afterwards so
# they cannot fabricate a visual retrieval margin.
visual_enriched = projected_visual_tokens + region
visual_enriched = visual_enriched * (
slot_valid_mask[:, :, None, None].to(
visual_enriched.dtype
)
)
resampled_visual, resampler_diagnostics = self.resampler(
visual_enriched,
region_coverage=coverage,
slot_valid_mask=slot_valid_mask,
patch_valid_mask=patch_valid_mask,
return_diagnostics=return_diagnostics,
)
type_ids = torch.tensor(
[0, 0, 0, 0, 1, 1, 1, 1],
device=projected_visual_tokens.device,
dtype=torch.long,
)
token_type = self.target_context_type_embedding(type_ids)[
None, None, :, :
]
resampled = resampled_visual + role + time + token_type
resampled = resampled * slot_valid_mask[:, :, None, None].to(
resampled.dtype
)
else:
enriched = projected_visual_tokens + region + role + time
enriched = enriched * slot_valid_mask[:, :, None, None].to(
enriched.dtype
)
resampled = self.resampler(
enriched,
slot_valid_mask=slot_valid_mask,
patch_valid_mask=patch_valid_mask,
)
resampled_visual = None
resampler_diagnostics = {}
effective_memory_scores = memory_scores
gate_diagnostics: Dict[str, torch.Tensor] = {}
if self.mixer_version == "current_to_memory_v9":
if current_visual_tokens is None:
raise ValueError(
"current_to_memory_v9 requires current visual patches"
)
if current_token_valid_mask is None:
current_token_valid_mask = torch.ones(
current_visual_tokens.shape[:2],
device=current_visual_tokens.device,
dtype=torch.bool,
)
gate_scores, gate_diagnostics = self.dynamic_memory_gate(
visual_memory_tokens=resampled_visual,
current_visual_tokens=current_visual_tokens,
current_token_valid_mask=current_token_valid_mask,
memory_confidence=memory_scores,
relative_frame_gaps=relative_frame_gaps,
slot_role_ids=role_ids,
slot_valid_mask=slot_valid_mask,
)
if memory_gate_blend is None:
blend = torch.tensor(
1.0,
device=gate_scores.device,
dtype=torch.float32,
)
else:
blend = torch.as_tensor(
memory_gate_blend,
device=gate_scores.device,
dtype=torch.float32,
).clamp(0.0, 1.0)
confidence = memory_scores.to(
device=gate_scores.device, dtype=torch.float32
)
valid = slot_valid_mask.to(
device=gate_scores.device, dtype=torch.bool
)
confidence = torch.where(
valid & torch.isfinite(confidence),
confidence.clamp(0.05, 0.95),
torch.zeros_like(confidence),
)
effective_memory_scores = (
(1.0 - blend) * confidence + blend * gate_scores
)
effective_memory_scores = torch.where(
valid,
effective_memory_scores.clamp(0.05, 0.95),
torch.zeros_like(effective_memory_scores),
)
gate_diagnostics["memory_gate_blend"] = blend
gate_diagnostics[
"effective_memory_scores"
] = effective_memory_scores
if self.mixer_version == "temporal_v2":
mixed, diagnostics = self.temporal_mixer(
resampled,
slot_valid_mask=slot_valid_mask,
memory_scores=effective_memory_scores,
return_diagnostics=return_diagnostics,
)
else:
if current_visual_tokens is None:
raise ValueError(
f"{self.mixer_version} requires current_visual_tokens"
)
if current_token_valid_mask is None:
current_token_valid_mask = torch.ones(
current_visual_tokens.shape[:2],
device=current_visual_tokens.device,
dtype=torch.bool,
)
mixed, diagnostics = self.current_to_memory_mixer(
resampled,
slot_valid_mask=slot_valid_mask,
memory_scores=effective_memory_scores,
current_visual_tokens=current_visual_tokens,
current_token_valid_mask=current_token_valid_mask,
language_tokens=language_tokens,
language_token_mask=language_token_mask,
return_diagnostics=return_diagnostics,
)
# Gate scores and effective scores are tiny and required by the loss
# path even when large diagnostic tensors are disabled.
if self.mixer_version == "current_to_memory_v9":
diagnostics.update(gate_diagnostics)
if return_diagnostics:
diagnostics["resampled_tokens"] = resampled
diagnostics["time_bucket_ids"] = time_ids
diagnostics.update(resampler_diagnostics)
return mixed, diagnostics
def apply_late_memory_bridge(
self,
track_embedding: torch.Tensor,
memory_tokens: torch.Tensor,
slot_valid_mask: torch.Tensor,
memory_scores: torch.Tensor,
*,
return_diagnostics: bool = False,
) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
"""Apply V7's pre-LLM bridge; earlier versions remain exact no-ops."""
if self.mixer_version not in (
"current_to_memory_v7",
"current_to_memory_v9",
"current_to_memory_v10",
"current_to_memory_v14",
):
return track_embedding, {}
return self.late_memory_bridge(
track_embedding,
memory_tokens,
slot_valid_mask,
memory_scores,
return_diagnostics=return_diagnostics,
)
"""Neural memory modules used by TrackingAnything.
Shapes follow the notation below:
B: batch size
M: memory slots (fixed to 6)
P: projected visual tokens per crop (256 for a 448 crop)
Q: resampled tokens per slot (fixed to 8)
D: LocateAnything language width (2048 for the 3B checkpoint)
"""
from __future__ import annotations
import math
from typing import Dict, Optional, Sequence, Tuple
import torch
from torch import nn
import torch.nn.functional as F
def _inverse_softplus(value: float) -> float:
if value <= 0:
raise ValueError("softplus target must be positive")
return math.log(math.expm1(value))
def _inverse_sigmoid(value: float) -> float:
if not 0.0 < value < 1.0:
raise ValueError("sigmoid target must be strictly between zero and one")
return math.log(value / (1.0 - value))
def normalized_log_memory_scores(
memory_scores: torch.Tensor,
slot_valid_mask: torch.Tensor,
*,
std_floor: float = 0.08,
clip_value: float = 2.0,
) -> torch.Tensor:
"""Return per-sample normalized log scores with invalid slots at zero.
The normalization is causal and deployment-safe: it uses only the valid
memories in the current bank, never the sequence length or future frames.
A standard-deviation floor prevents a nearly uniform bank from amplifying
numerical noise, while clipping gives a strict bound to every downstream
score prior.
"""
if memory_scores.ndim != 2:
raise ValueError("memory_scores must have shape [B, M]")
valid = slot_valid_mask.to(
device=memory_scores.device,
dtype=torch.bool,
)
if valid.shape != memory_scores.shape:
raise ValueError("slot_valid_mask must match memory_scores")
if std_floor <= 0.0:
raise ValueError("std_floor must be positive")
if clip_value <= 0.0:
raise ValueError("clip_value must be positive")
scores = torch.where(
valid & torch.isfinite(memory_scores),
memory_scores.float().clamp(1e-6, 1.0),
torch.ones_like(memory_scores, dtype=torch.float32),
)
log_scores = scores.log()
weights = valid.to(log_scores.dtype)
count = weights.sum(dim=-1, keepdim=True).clamp_min(1.0)
mean = (log_scores * weights).sum(dim=-1, keepdim=True) / count
centered = (log_scores - mean) * weights
variance = centered.square().sum(dim=-1, keepdim=True) / count
# Clamp the variance before sqrt. ``sqrt(0).clamp_min(floor)`` has the
# right forward value but an undefined/infinite derivative at a uniform
# differentiable score vector (V9 initializes every gate at 0.5), which
# turns the first optimizer update into NaNs. This ordering is forward-
# equivalent and gives a zero, finite gradient below the variance floor.
denominator = variance.clamp_min(float(std_floor) ** 2).sqrt()
normalized = (centered / denominator).clamp(
-float(clip_value),
float(clip_value),
)
return normalized * weights
def relative_time_bucket(
frame_gaps: torch.Tensor,
num_buckets: int = 8,
) -> torch.Tensor:
"""Map non-negative frame gaps to causal logarithmic time buckets.
Bucket 0 is the same frame (normally unused for a causal memory). Positive
gaps use powers of two: 1, 2--3, 4--7, ..., with the tail clipped to the
final bucket. This has no dependency on the future sequence length.
"""
if num_buckets < 2:
raise ValueError("num_buckets must be at least 2")
gaps = torch.as_tensor(frame_gaps)
if torch.any(gaps < 0):