"""Raw-evidence/question recurrent reasoning on Qwen2.5-VL.
The oracle-selected ``raw_evidence_question`` path keeps the complete frozen
visual state as immutable evidence ``E=V*`` and initializes the recurrent state
with the text-only question scaffold ``R0=Q*``. One shared Cross-Attention
transition reads ``E`` for four residual steps, then parameter-free trajectory
aggregation preserves every recurrent depth::
R0 = Q*
Rk = R{k-1} + CrossAttn(R{k-1}, E)
R_agg = mean(R1, R2, R3, R4) -> y
During training only, ``R_agg`` is also restored toward the image-conditioned
question rows from the frozen multimodal lower pass. That target is already
available from the pass that builds ``V*`` and never enters the inference
decoder. Historical final-state, state-replacement, and question-anchor modes
remain loadable for checkpoint comparison.
The multimodal upper continuation and its KV cache are discarded. The lower
replacement cache contains only text-only ``Q*`` rows; the same positions are
overwritten by the selected decoder latent above ``ell_star``. Thus both
halves have length ``N_q``:
there are no workspace slots, placeholder tokens, learned aggregation, or
direct ``E``-to-decoder connection in this path.
The cross-attention uses Qwen's native 28-query/4-KV-head GQA projections and
starts from the first upper decoder layer's Q/K/V/O and RMSNorm weights. It is
still cross-attention, so the source layer's self-attention RoPE is not reused.
The current method has no workspace FFN. In residual mode, disabling a visual
read makes that step an identity; in replacement mode it writes zero on valid
question rows. A compatibility-only FFN can still be allocated for older
checkpoints.
Historical slot-workspace implementations remain below for loading and
evaluating old checkpoints. Configuration validation makes those modules
mutually exclusive with ``raw_evidence_question``.
Shapes: ``B`` batch, ``N_v`` visual tokens, ``N_q`` question tokens, and ``d``
the native backbone width (3584).
"""
from __future__ import annotations
from contextlib import contextmanager
from dataclasses import dataclass, replace
import math
import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint
from transformers.cache_utils import Cache, DynamicCache
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.modeling_outputs import ModelOutput
from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
Qwen2_5_VLForConditionalGeneration,
Qwen2_5_VLPreTrainedModel,
apply_multimodal_rotary_pos_emb,
)
from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
from .configuration_source_qwen25 import CloseQwen2_5_VLConfig
from .source_perceive import (
PerceiveDeliberateLayout,
build_answer_decode_mask,
build_perceive_deliberate_mask,
build_perceive_deliberate_positions,
)
from .source_spatial import SpatialVisualRecurrentCell
from .source_splitting import (
block_attention_edges,
embed_multimodal,
final_norm,
image_token_mask,
make_split_context,
run_layer_range,
select_tokens_padded,
vision_span_mask,
)
def _replace_selected_rows_from_padded(
full_state: torch.Tensor,
selection_mask: torch.Tensor,
padded_rows: torch.Tensor,
padding_mask: torch.Tensor,
) -> torch.Tensor:
"""Replace variable-count rows of ``full_state`` from a padded sequence.
This is the inverse layout operation of :func:`select_tokens_padded` for
valid rows. The small per-example loop is intentional: image resolution
makes the multimodal sequence lengths heterogeneous, while preserving each
sample's original physical/M-RoPE positions is part of the native-layer
counterfactual.
"""
if selection_mask.dtype != torch.bool or padding_mask.dtype != torch.bool:
raise TypeError("selection and padding masks must be boolean")
if full_state.ndim != 3 or padded_rows.ndim != 3:
raise ValueError("full_state and padded_rows must be [B, L, D]")
if selection_mask.shape != full_state.shape[:2]:
raise ValueError("selection_mask does not match full_state")
if padding_mask.shape != padded_rows.shape[:2]:
raise ValueError("padding_mask does not match padded_rows")
if full_state.shape[0] != padded_rows.shape[0] or (
full_state.shape[-1] != padded_rows.shape[-1]
):
raise ValueError("full and padded states have incompatible shapes")
selected_counts = selection_mask.sum(dim=-1)
padded_counts = (~padding_mask).sum(dim=-1)
if not torch.equal(selected_counts, padded_counts):
raise ValueError(
"selected-row counts disagree with padded valid counts: "
f"{selected_counts.tolist()} vs {padded_counts.tolist()}"
)
output = full_state.clone()
for batch_index, count_tensor in enumerate(selected_counts):
count = int(count_tensor)
output[batch_index, selection_mask[batch_index]] = padded_rows[
batch_index, :count
]
return output
def answer_cross_entropy(
logits: torch.Tensor,
labels: torch.Tensor,
*,
per_example: bool = False,
) -> torch.Tensor:
"""Answer CE with either token- or example-balanced reduction.
``per_example=False`` exactly matches PyTorch's historical global valid-token
mean. ``True`` gives every non-empty sample equal weight regardless of its
answer length, preventing long caption answers from overwhelming one-token
visual-decision examples.
"""
if not per_example:
return nn.functional.cross_entropy(
logits.reshape(-1, logits.shape[-1]).float(),
labels.reshape(-1),
ignore_index=-100,
)
bsz = labels.shape[0]
token_loss = nn.functional.cross_entropy(
logits.reshape(-1, logits.shape[-1]).float(),
labels.reshape(-1),
ignore_index=-100,
reduction="none",
).view(bsz, -1)
valid = labels.reshape(bsz, -1).ne(-100)
counts = valid.sum(dim=-1)
nonempty = counts > 0
if not bool(nonempty.any()):
return logits.float().sum() * 0.0
sample_loss = (token_loss * valid).sum(dim=-1) / counts.clamp_min(1)
return sample_loss[nonempty].mean()
def _answer_nll_per_example(
logits: torch.Tensor,
labels: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Return length-normalized answer NLL and a non-empty-row mask.
The CRR reliance intervention compares factual and swapped answers across
examples with very different target lengths. A per-token mean keeps that
score gap in the same units for one-token MC labels and caption answers.
"""
bsz = labels.shape[0]
token_loss = nn.functional.cross_entropy(
logits.reshape(-1, logits.shape[-1]).float(),
labels.reshape(-1),
ignore_index=-100,
reduction="none",
).view(bsz, -1)
valid = labels.reshape(bsz, -1).ne(-100)
counts = valid.sum(dim=-1)
nll = (token_loss * valid).sum(dim=-1) / counts.clamp_min(1)
return nll, counts > 0
def _different_answer_residual_donors(
labels: torch.Tensor,
question_attention_mask: torch.Tensor | None = None,
) -> tuple[torch.LongTensor, torch.BoolTensor]:
"""Choose a deterministic batch-local donor with a different answer.
Equal question lengths are preferred so that right-padded Q-shaped
residuals align exactly. Equal answer lengths are the second preference,
which usually keeps MC examples in the same formatting regime. Rows for
which the local batch contains no different answer are marked invalid and
do not contribute to the reliance loss.
"""
bsz = labels.shape[0]
rows = torch.arange(bsz, device=labels.device)
donors = rows.clone()
found = torch.zeros(bsz, dtype=torch.bool, device=labels.device)
if bsz < 2:
return donors, found
same_answer = labels[:, None, :].eq(labels[None, :, :]).all(dim=-1)
answer_lengths = labels.ne(-100).sum(dim=-1)
if question_attention_mask is None:
question_lengths = torch.ones_like(answer_lengths)
else:
question_lengths = question_attention_mask.gt(0).sum(dim=-1)
# Cyclic offsets avoid concentrating every negative on batch row zero.
# Compatibility is relaxed only after all offsets at the stricter level
# have been exhausted.
for compatibility in ("question", "answer", "any"):
for offset in range(1, bsz):
candidates = (rows + offset) % bsz
eligible = ~same_answer[rows, candidates]
if compatibility == "question":
eligible &= question_lengths.eq(
question_lengths.index_select(0, candidates)
)
elif compatibility == "answer":
eligible &= answer_lengths.eq(
answer_lengths.index_select(0, candidates)
)
choose = (~found) & eligible
donors[choose] = candidates[choose]
found |= choose
if bool(found.all()):
return donors, found
return donors, found
def _replacement_prefix_layout(
question_mask: torch.Tensor, num_slots: int
) -> tuple[torch.Tensor, torch.LongTensor, torch.LongTensor]:
"""Mask and logical positions for the physical ``[q_pad; slots]`` layout.
Questions are right-padded to a common physical width before the ``S`` slot
rows are appended. The slots must nevertheless occupy the *logical*
positions immediately after each item's last real question token. Using a
batch-wide ``arange(N_q + S)`` makes RoPE positions, and therefore answers,
depend on the longest question in the batch.
Returns ``(prefix_mask, position_ids, logical_lengths)`` with shapes
``[B,N_q+S]``, ``[3,B,N_q+S]``, and ``[B]``. ``logical_lengths`` is the
first answer-token position, i.e. ``question_length + S``.
"""
if question_mask.ndim != 2:
raise ValueError(
f"question_mask must have shape [B,N_q], got {tuple(question_mask.shape)}"
)
if num_slots < 1:
raise ValueError("num_slots must be >= 1")
keep = question_mask > 0
b = keep.shape[0]
q_positions = keep.long().cumsum(dim=-1) - 1
q_positions = q_positions.masked_fill(~keep, 0)
q_lengths = keep.sum(dim=-1).long() # [B]
slot_offsets = torch.arange(num_slots, device=question_mask.device)
slot_positions = q_lengths[:, None] + slot_offsets[None, :] # [B,S]
one_d = torch.cat([q_positions, slot_positions], dim=-1) # [B,N_q+S]
position_ids = one_d.unsqueeze(0).expand(3, b, -1)
slot_mask = torch.ones(
b, num_slots, dtype=question_mask.dtype, device=question_mask.device
)
prefix_mask = torch.cat([question_mask, slot_mask], dim=-1)
return prefix_mask, position_ids, q_lengths + num_slots
# ---------------------------------------------------------------------------
# workspace blocks
# ---------------------------------------------------------------------------
class _Block(nn.Module):
"""Pre-norm block: optional cross-attention over ``context``, then FFN."""
def __init__(self, d_w: int, n_heads: int, ffn_mult: int, dropout: float, cross: bool):
super().__init__()
self.self_norm = nn.LayerNorm(d_w)
self.self_attn = nn.MultiheadAttention(
d_w, n_heads, dropout=dropout, batch_first=True
)
self.cross = cross
if cross:
self.cross_norm_q = nn.LayerNorm(d_w)
self.cross_norm_kv = nn.LayerNorm(d_w)
self.cross_attn = nn.MultiheadAttention(
d_w, n_heads, dropout=dropout, batch_first=True
)
self.ffn_norm = nn.LayerNorm(d_w)
self.ffn = nn.Sequential(
nn.Linear(d_w, ffn_mult * d_w),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(ffn_mult * d_w, d_w),
)
def forward(
self,
z: torch.Tensor, # [B, S, d_w]
context: torch.Tensor | None = None, # [B, N, d_w]
context_padding_mask: torch.Tensor | None = None, # [B, N] True == ignore
cond: torch.Tensor | None = None, # [1|B, S, d_w] norm 입력 조건 (carrier 미반영)
) -> torch.Tensor:
# persistent slot-ID(advisor): cond는 attention/FFN을 조건화할 뿐
# residual carrier에는 절대 실리지 않는다 — z' = z + F(LN(z + E)).
# cond=None이면 비트동일.
zc = z if cond is None else z + cond
h = self.self_norm(zc)
z = z + self.self_attn(h, h, h, need_weights=False)[0]
if self.cross:
if context is None:
raise ValueError("cross-attention block called without context")
zc = z if cond is None else z + cond
q = self.cross_norm_q(zc)
kv = self.cross_norm_kv(context)
z = z + self.cross_attn(
q, kv, kv, key_padding_mask=context_padding_mask, need_weights=False
)[0]
zc = z if cond is None else z + cond
return z + self.ffn(self.ffn_norm(zc))
class _LowRankMHA(nn.Module):
"""rank-r 부분공간 attention: Up(MHA(Down(q), Down(kv))) — d 유지, 파라미터 O(d·r)."""
def __init__(self, d: int, r: int, n_heads: int):
super().__init__()
self.dq = nn.Linear(d, r, bias=False)
self.dkv = nn.Linear(d, r, bias=False)
self.attn = nn.MultiheadAttention(r, n_heads, dropout=0.0, batch_first=True)
self.up = nn.Linear(r, d, bias=False)
def forward(self, q, kv, key_padding_mask=None):
out, _ = self.attn(self.dq(q), self.dkv(kv), self.dkv(kv),
key_padding_mask=key_padding_mask, need_weights=False)
return self.up(out)
class _LowRankBlock(nn.Module):
"""_Block의 low-rank 동형: pre-norm + rank-r attention/FFN 잔차.
ER v2 (advisor): d_w=d에서 full dense workspace의 파라미터 폭증을 막는
"state width matching과 parameter explosion의 분리". cond 인터페이스는
_Block과 동일 (미사용 시 비트동일 경로).
"""
def __init__(self, d_w: int, n_heads: int, ffn_mult: int, dropout: float,
cross: bool, rank: int):
super().__init__()
self.self_norm = nn.LayerNorm(d_w)
self.self_attn = _LowRankMHA(d_w, rank, n_heads)
self.cross = cross
if cross:
self.cross_norm_q = nn.LayerNorm(d_w)
self.cross_norm_kv = nn.LayerNorm(d_w)
self.cross_attn = _LowRankMHA(d_w, rank, n_heads)
self.ffn_norm = nn.LayerNorm(d_w)
self.ffn = nn.Sequential(
nn.Linear(d_w, ffn_mult * rank), nn.GELU(), nn.Linear(ffn_mult * rank, d_w),
)
def forward(self, z, context=None, context_padding_mask=None, cond=None):
zc = z if cond is None else z + cond
h = self.self_norm(zc)
z = z + self.self_attn(h, h)
if self.cross:
if context is None:
raise ValueError("cross-attention block called without context")
zc = z if cond is None else z + cond
q = self.cross_norm_q(zc)
kv = self.cross_norm_kv(context)
z = z + self.cross_attn(q, kv, key_padding_mask=context_padding_mask)
zc = z if cond is None else z + cond
return z + self.ffn(self.ffn_norm(zc))
def _make_block(config, cross: bool):
r = int(getattr(config, "workspace_low_rank", 0) or 0)
if r > 0:
return _LowRankBlock(config.workspace_width, config.workspace_num_heads,
config.workspace_ffn_mult, config.workspace_dropout,
cross, r)
return _Block(config.workspace_width, config.workspace_num_heads,
config.workspace_ffn_mult, config.workspace_dropout, cross)
class WorkspaceRead(nn.Module):
"""``r_theta``: condition learnable slots on ``Q*``, then read ``V*``.
``Z^(0) = r_theta(Z_init, Q*, V*)``. ``V*`` is consumed only inside this
module: by :meth:`forward` (the initial read), and -- when
``config.read_gating`` is on -- by :meth:`reread`, which applies the SAME
cross-attention weights to the evolved state under a scalar gate. With
gating off the read is single by construction and the caller discards
``V*`` immediately after; either way no other module ever sees it.
"""
def __init__(self, config: CloseQwen2_5_VLConfig):
super().__init__()
d, d_w = config.backbone_width, config.workspace_width
self.slots = nn.Parameter(torch.randn(config.num_workspace_slots, d_w) * 0.02)
_eye = d == d_w # ER v2: 등폭이면 입사 projection(d→d) 자체가 불필요
self.q_in = nn.Identity() if _eye else nn.Linear(d, d_w, bias=False)
self.v_in = nn.Identity() if _eye else nn.Linear(d, d_w, bias=False)
self.condition = _make_block(config, cross=True)
self.read = nn.ModuleList(
_make_block(config, cross=True) for _ in range(config.read_num_blocks)
)
# v2: image-dependent from step 0. A random cross-attention starts out
# near-uniform, so its read is ~the mean of V* regardless of content and
# the optimizer can settle on ignoring vision (probe 20143). The
# shortcut injects an orthogonal projection of the LN'd pooled V* into
# every slot, so ignoring the image is never the zero-gradient default.
self.pool_in = None
if getattr(config, "read_shortcut", False):
self.pool_in = nn.Linear(d, d_w, bias=False)
# NOTE: post_init()'s HF sweep has overwritten this orthogonal init
# with N(0, initializer_range) in every run since v2 -- discovered
# while diagnosing the gate-init wipe (smoke 20552). Deliberately
# NOT flagged: v4/v5 trained under the overwritten init, and v6
# must differ from v5 by the gate alone. Revisit only in a fresh
# baseline generation.
nn.init.orthogonal_(self.pool_in.weight)
self.out_norm = nn.LayerNorm(d_w)
# v6 candidate (read_gating): the SAME reader may be consulted again
# between transitions, scaled by a scalar gate per item. b_gamma =
# logit(eps): near-identity start (training begins ~at the proven
# read-once regime; opening is earned by the losses). No new reader
# parameters -- re-reads share v_in / read blocks / out_norm, so "how
# it reads" is identical and only "when" is learned.
self.gate = None
if getattr(config, "read_gating", False):
self.gate = nn.Linear(2 * d_w, 1)
nn.init.zeros_(self.gate.weight)
eps = float(getattr(config, "gate_init_eps", 0.02))
eps = min(max(eps, 1e-6), 1 - 1e-6)
nn.init.constant_(self.gate.bias, math.log(eps / (1 - eps)))
# post_init()'s HF sweep re-inits every unflagged nn.Linear
# (weight ~ N(0, initializer_range), bias -> 0), which silently
# destroyed this init: with a random w over the unnormalised q̄
# scale the gate saturated OPEN from step 0 (smoke 20552: gamma
# 0.98 at the first log, 1.0 thereafter -- the exact opposite of
# near-identity). The flag is respected by PreTrainedModel.
# _initialize_weights and keeps the custom init authoritative.
self.gate._is_hf_initialized = True
# v7 (read_residual_pure): always-on PURE visual residual reread.
# A^(k) = XAttn(LN(U^(k)) -> queries, V* -> keys/values)
# Z^(k) = U^(k) + W_o A^(k)
# No self-attention, no FFN, no output norm, no biases anywhere in the
# branch, W_o zero-initialised (exact identity at init -- stronger than
# v6's near-identity). Every atom of the update is a value vector of
# V*: the battery on ckpt-5724 showed the gated reader's worth was
# extra DEPTH, not vision; this branch cannot provide content-free
# depth, so any benefit is visual by construction.
self.pure_ln = None
if (getattr(config, "read_residual_pure", False)
or getattr(config, "final_verify", False)
or getattr(config, "mid_read_step", 0)):
self.pure_ln = nn.LayerNorm(d_w, bias=False)
self.pure_attn = nn.MultiheadAttention(
d_w, config.workspace_num_heads, batch_first=True, bias=False,
kdim=d, vdim=d,
)
nn.init.zeros_(self.pure_attn.out_proj.weight)
if getattr(config, "pure_warm_start", False):
# copy visual-reading geometry from the initial reader: K/V
# from v_in ([d_w, d] both), Q from the first read block's
# query in-proj slice ([d_w, d_w]).
with torch.no_grad():
self.pure_attn.k_proj_weight.copy_(self.v_in.weight)
self.pure_attn.v_proj_weight.copy_(self.v_in.weight)
qw = self.read[0].cross_attn.in_proj_weight[: config.workspace_width]
self.pure_attn.q_proj_weight.copy_(qw)
# post_init()'s sweep visits every SUBMODULE; a flag on the parent
# MHA does not shield its out_proj Linear (the gate-wipe lesson,
# caught again by tests/test_pure_residual.py). Flag the leaf.
self.pure_attn._is_hf_initialized = True
self.pure_attn.out_proj._is_hf_initialized = True
self.pure_ln._is_hf_initialized = True
def forward(
self,
q_star: torch.Tensor, # [B, N_q, d]
v_star: torch.Tensor, # [B, N_v, d]
q_padding_mask: torch.Tensor | None = None, # [B, N_q] True == ignore
v_padding_mask: torch.Tensor | None = None, # [B, N_v] True == ignore
) -> torch.Tensor:
b = q_star.shape[0]
z = self.slots.unsqueeze(0).expand(b, -1, -1) # [B, S, d_w]
q = self.q_in(q_star) # [B, N_q, d_w]
z = self.condition(z, context=q, context_padding_mask=q_padding_mask)
if self.pool_in is not None:
if v_padding_mask is None:
pooled = v_star.mean(1) # [B, d]
else:
keep = (~v_padding_mask).to(v_star.dtype).unsqueeze(-1) # [B, N_v, 1]
pooled = (v_star * keep).sum(1) / keep.sum(1).clamp_min(1.0)
pooled = nn.functional.layer_norm(pooled, (pooled.shape[-1],))
z = z + self.pool_in(pooled).unsqueeze(1) # broadcast over slots
# Dynamic resolution means N_v varies across a batch; V* arrives padded
# to the batch maximum and the mask hides the filler.
v = self.v_in(v_star) # [B, N_v, d_w]
for block in self.read:
z = block(z, context=v, context_padding_mask=v_padding_mask)
return self.out_norm(z) # [B, S, d_w]
def visual_residual(
self,
query_state: torch.Tensor, # [B, S, d_w] -- query only; NOT added to
v_star: torch.Tensor, # [B, N_v, d]
v_padding_mask: torch.Tensor | None = None,
) -> torch.Tensor:
"""``W_o XAttn(LN(q), V*)`` -- the raw visual residual, base-free.
Separated from :meth:`reread_pure` so held-out interventions can pair
an arbitrary query source with an arbitrary decoder base (advisor's
query-only substitution)."""
if self.pure_ln is None:
raise RuntimeError("visual_residual() requires config.read_residual_pure=True")
a, _ = self.pure_attn(
self.pure_ln(query_state), v_star, v_star,
key_padding_mask=v_padding_mask, need_weights=False,
)
return a
def reread_pure(
self,
u: torch.Tensor, # [B, S, d_w] provisional state
v_star: torch.Tensor, # [B, N_v, d]
v_padding_mask: torch.Tensor | None = None,
query_state: torch.Tensor | None = None, # persistent slot-ID query
) -> torch.Tensor:
"""v7: ``Z = U + W_o XAttn(LN(U), V*)``. Update lives in span(V(V*)).
``query_state``가 주어지면 query만 그것으로 대체한다 (residual 기저는
여전히 ``u``) — persistent slot identity가 ``LN(U+E_slot)``을 query로
쓰되 state에는 E를 누적하지 않기 위한 통로. ``None``이면 비트동일."""
q = u if query_state is None else query_state
return u + self.visual_residual(q, v_star, v_padding_mask)
def reread(
self,
z: torch.Tensor, # [B, S, d_w] current workspace state
q_star: torch.Tensor, # [B, N_q, d]
v_star: torch.Tensor, # [B, N_v, d]
q_padding_mask: torch.Tensor | None = None,
v_padding_mask: torch.Tensor | None = None,
gate_override: float | torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Gated consultation of ``V*`` by the EVOLVED workspace state.
``z' = z + g * (read(z, V*) - z)`` with scalar ``g`` per item, so
``g = 0`` reproduces ``z`` exactly (bit-identical to the no-reread
path) and ``g = 1`` is a full re-read. The reader weights are shared
with the initial read; only the gate decides WHEN looking again is
worth it. Returns ``(z', g)`` with ``g`` of shape ``[B]``.
``gate_override`` substitutes the learned gate with a fixed value
(forced-open 1.0 / forced-closed 0.0 / anything between). Held-out
causal evaluation ONLY -- §3.2's gate interventions. Never set during
training; ``None`` (the default) is the learned-gate path, bitwise.
"""
if self.gate is None:
raise RuntimeError("reread() requires config.read_gating=True")
if q_padding_mask is None:
q_bar = self.q_in(q_star).mean(1) # [B, d_w]
else:
keep = (~q_padding_mask).to(q_star.dtype).unsqueeze(-1)
q_bar = self.q_in((q_star * keep).sum(1) / keep.sum(1).clamp_min(1.0))
# LN on the pooled state: the gate should react to the state's
# direction, not to scale drift across steps/runs.
z_bar = nn.functional.layer_norm(z.mean(1), (z.shape[-1],)) # [B, d_w]
g = torch.sigmoid(self.gate(torch.cat([z_bar, q_bar], dim=-1))).squeeze(-1) # [B]
if gate_override is not None:
g = torch.as_tensor(gate_override, dtype=g.dtype, device=g.device).expand_as(g)
if float(g.max()) == 0.0:
return z, g # exact noop, mirrors the closed-gate residual form
v = self.v_in(v_star) # [B, N_v, d_w]
upd = z
for block in self.read:
upd = block(upd, context=v, context_padding_mask=v_padding_mask)
upd = self.out_norm(upd)
return z + g.view(-1, 1, 1) * (upd - z), g
class WorkspaceTransition(nn.Module):
"""``f_theta``: shared recurrent step. Sees ``Z^(k-1)`` and ``Q*``. Never ``V*``.
Shared across all ``K`` steps, so the trainable parameter count is
independent of ``K`` (§3.2).
"""
def __init__(self, config: CloseQwen2_5_VLConfig):
super().__init__()
d, d_w = config.backbone_width, config.workspace_width
self.q_in = (nn.Identity() if d == d_w
else nn.Linear(d, d_w, bias=False))
self.blocks = nn.ModuleList(
_make_block(config, cross=True) for _ in range(config.transition_num_blocks)
)
self.out_norm = nn.LayerNorm(d_w)
def forward(
self,
z: torch.Tensor, # [B, S, d_w]
q_star: torch.Tensor, # [B, N_q, d]
q_padding_mask: torch.Tensor | None = None,
cond: torch.Tensor | None = None, # persistent slot-ID (norm 입력 전용)
) -> torch.Tensor:
q = self.q_in(q_star) # [B, N_q, d_w]
for block in self.blocks:
z = block(z, context=q, context_padding_mask=q_padding_mask, cond=cond)
return self.out_norm(z) # [B, S, d_w]
class RawEvidenceQuestionTransition(nn.Module):
"""Shared recurrent update for the slot-free ``E=V*; R0=Q*`` path.
The state has the same rows and width as the text-only question scaffold.
Each step reads the complete immutable visual memory with one pre-norm
cross-attention. The residual-compatible mode used by existing runs is::
R_next = R + CrossAttn(RMSNorm(R), RMSNorm(E), RMSNorm(E))
The state-replacement candidate removes that recurrent identity path::
H_next = CrossAttn(RMSNorm(H), RMSNorm(E), RMSNorm(E))
Its question-anchored form replaces the one-time decoder addition with a
fixed scaffold at every recurrent step::
Z_next = Q* + CrossAttn(RMSNorm(Z), RMSNorm(E), RMSNorm(E))
There is no self-attention or FFN in either current method: ``Q*`` is already
contextualized by the frozen lower backbone, and the upper decoder retains
its pretrained attention/MLP stack. Consequently every recurrent state
change is caused by a visual-memory read. The same module instance is
reused for every recurrent step.
"""
def __init__(self, config: CloseQwen2_5_VLConfig):
super().__init__()
text_cfg = config.text_config
d = config.backbone_width
self.num_heads = int(text_cfg.num_attention_heads)
self.num_key_value_heads = int(text_cfg.num_key_value_heads)
self.num_key_value_groups = self.num_heads // self.num_key_value_heads
self.head_dim = d // self.num_heads
self.dropout = float(config.workspace_dropout)
if self.num_heads % self.num_key_value_heads:
raise ValueError("query heads must be divisible by KV heads")
# Native Qwen GQA shapes. Q/K/V/O and the input RMSNorm are copied
# from L(ell_star+1) after post_init, before upper-layer LoRA injection.
self.read_norm_q = Qwen2RMSNorm(d, eps=text_cfg.rms_norm_eps)
self.read_norm_e = Qwen2RMSNorm(d, eps=text_cfg.rms_norm_eps)
self.q_proj = nn.Linear(d, self.num_heads * self.head_dim, bias=True)
self.k_proj = nn.Linear(
d, self.num_key_value_heads * self.head_dim, bias=True
)
self.v_proj = nn.Linear(
d, self.num_key_value_heads * self.head_dim, bias=True
)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, d, bias=False)
self.state_replacement = bool(
getattr(config, "raw_state_replacement", False)
)
self.question_anchored = bool(
getattr(config, "raw_question_anchor", False)
)
# Compatibility only: old raw-E/Q checkpoints used a random d->d->d
# FFN after every read and did not serialize a feature flag. Their
# configs resolve raw_transition_ffn=None to True; all new launches
# explicitly set it False and allocate neither module nor parameters.
self.use_ffn = bool(getattr(config, "raw_transition_ffn", True))
if self.use_ffn:
self.ffn_norm = Qwen2RMSNorm(d, eps=text_cfg.rms_norm_eps)
self.ffn = nn.Sequential(
nn.Linear(d, config.workspace_ffn_mult * d),
nn.GELU(),
nn.Dropout(config.workspace_dropout),
nn.Linear(config.workspace_ffn_mult * d, d),
)
def forward(
self,
r: torch.Tensor, # [B,N_q,d]
evidence: torch.Tensor | None = None, # [B,N_v,d] == raw V*
evidence_padding_mask: torch.Tensor | None = None,
question_padding_mask: torch.Tensor | None = None,
read_scale: float = 1.0, # held-out intervention only
projected_evidence: tuple[torch.Tensor, torch.Tensor] | None = None,
question_anchor: torch.Tensor | None = None,
) -> torch.Tensor:
"""Apply one recurrent update.
``projected_evidence`` contains head-split K/V tensors returned by
:meth:`project_evidence`. Since raw E is immutable across recurrent
steps, the normal rollout computes these once and reuses them. This is
algebraically identical to calling ``nn.MultiheadAttention`` four
times, but avoids repeating the dominant ``N_v x d x d`` K/V
projections for every step.
"""
if projected_evidence is None:
if evidence is None:
raise ValueError("evidence or projected_evidence is required")
projected_evidence = self.project_evidence(evidence)
k, v = projected_evidence # [B,H,N_v,D_h]
q_norm = self.read_norm_q(r)
d = q_norm.shape[-1]
q = self._split_heads(self.q_proj(q_norm), self.num_heads)
attn_mask = None
if evidence_padding_mask is not None:
# SDPA's additive mask uses -inf for keys that must be ignored.
attn_mask = torch.zeros(
evidence_padding_mask.shape[0],
1,
1,
evidence_padding_mask.shape[1],
dtype=q.dtype,
device=q.device,
).masked_fill(
evidence_padding_mask[:, None, None, :], float("-inf")
)
delta = nn.functional.scaled_dot_product_attention(
q,
k,
v,
attn_mask=attn_mask,
dropout_p=(self.dropout if self.training else 0.0),
is_causal=False,
enable_gqa=True,
)
delta = delta.transpose(1, 2).contiguous().view(r.shape[0], r.shape[1], d)
delta = self.o_proj(delta)
if self.state_replacement:
r_next = float(read_scale) * delta
if self.question_anchored:
if question_anchor is None:
raise ValueError(
"question_anchor is required by raw_question_anchor"
)
# Preserve Q* but replace the visual response. There is no
# R_{k-1} residual, so the old identity carrier stays removed.
r_next = question_anchor + r_next
else:
r_next = r + float(read_scale) * delta
if self.use_ffn:
r_next = r_next + self.ffn(self.ffn_norm(r_next))
if question_padding_mask is not None:
# Padding rows are physical cache plumbing only. Keeping them
# unchanged makes batched and per-item recurrent states agree.
padding_value = (
question_anchor
if self.question_anchored and question_anchor is not None
else r
)
r_next = torch.where(
question_padding_mask.unsqueeze(-1), padding_value, r_next
)
return r_next
def _split_heads(self, x: torch.Tensor, num_heads: int) -> torch.Tensor:
b, n, d = x.shape
return x.view(b, n, num_heads, d // num_heads).transpose(1, 2)
def project_evidence(
self, evidence: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""Normalize and project immutable E once for a recurrent rollout."""
e = self.read_norm_e(evidence)
k = self.k_proj(e)
v = self.v_proj(e)
return (
self._split_heads(k, self.num_key_value_heads),
self._split_heads(v, self.num_key_value_heads),
)
#: Backbone projections the post-``l*`` adapter attaches to.
ADAPTER_TARGET_MODULES = (
"q_proj",
"k_proj",
"v_proj",
"o_proj",
"gate_proj",
"up_proj",
"down_proj",
)
def build_adapter_config(config: CloseQwen2_5_VLConfig):
"""LoRA config restricted to the decoder layers above ``l*``.
peft rather than a hand-rolled adapter: a hook-attached delta holds a
closure over its target module, which is fragile under ``device_map``
sharding and FSDP -- exactly the settings this will train in. peft's
``layers_to_transform`` expresses the "above ``l*`` only" restriction
directly, and its save/load path is already Trainer-compatible.
"""
from peft import LoraConfig
upper = config.upper_slice
excl = set(getattr(config, "adapter_exclude_layers", []) or [])
if getattr(config, "recurrence_only_adapter", False):
target_layers = [int(config.ell_star) + 1]
else:
target_layers = [
li for li in range(upper.start, upper.stop) if li not in excl
]
return LoraConfig(
r=config.adapter_rank,
lora_alpha=config.adapter_alpha,
lora_dropout=config.adapter_dropout,
bias="none",
target_modules=list(ADAPTER_TARGET_MODULES),
layers_to_transform=target_layers,
layers_pattern="layers",
task_type=None,
)
# ---------------------------------------------------------------------------
# output type
# ---------------------------------------------------------------------------
@dataclass
class CloseCausalLMOutput(ModelOutput):
"""Adds the workspace trajectory, which the Trainer needs for ``L_traj``."""
loss: torch.FloatTensor | None = None
logits: torch.FloatTensor | None = None
past_key_values: Cache | None = None
workspace_states: tuple[torch.FloatTensor, ...] | None = None # K+1 x [B, S, d_w]
#: LN'd masked-mean of V*, detached -- the target side of the Trainer's
#: read-grounding InfoNCE. Never fed back into the model.
v_star_pooled: torch.FloatTensor | None = None # [B, d]
#: bridge-only interface distillation 항 (advisor): 1-cos(LN h_R, LN h_base).
#: 학습 시에만 채워짐. gradient는 workspace_out·splice_gain·l*+1층 adapter 한정.
iface_loss: torch.FloatTensor | None = None
#: Training-only normalized state-restoration loss. Its target is the
#: image-conditioned question state from the frozen multimodal lower pass;
#: inference never receives that state.
functional_state_loss: torch.FloatTensor | None = None
#: Training-only score gap for the matched residual-swap intervention:
#: normalized log p(y_i | B_i + C_i) - log p(y_i | B_i + C_donor).
residual_swap_score_gap: torch.FloatTensor | None = None # [B]
residual_swap_valid: torch.BoolTensor | None = None # [B]
#: 1-indexed CRR transition selected by the transition-aware intervention.
#: Aggregate-swap checkpoints and ordinary inference leave this unset.
residual_swap_transition: torch.LongTensor | None = None # scalar
#: Training-only decoder NLLs for the selected hard-case corrective
#: transition. ``current`` remains differentiable; ``previous`` and
#: ``swap`` are comparator values detached by the Trainer so the objective
#: cannot win by degrading either baseline.
corrective_current_nll: torch.FloatTensor | None = None # [B]
corrective_previous_nll: torch.FloatTensor | None = None # [B]
corrective_swap_nll: torch.FloatTensor | None = None # [B]
corrective_valid: torch.BoolTensor | None = None # [B]
#: Training-only gold-answer NLL at each decoded CRR prefix
#: Z_k = B + mean(C_1..C_k). Every column is scored by the same actual
#: upper decoder; the Trainer applies a one-sided no-regression objective.
step_answer_nll: torch.FloatTensor | None = None # [B, T]
step_answer_valid: torch.BoolTensor | None = None # [B, T]
#: RL-only greedy generations sampled from a stochastic CRR trajectory.
#: Ordinary SFT/inference calls leave every field below unset.
generated_ids: torch.LongTensor | None = None # [B, <=N_a]
#: Frozen-decoder, token-normalized gold-answer log likelihood for a
#: sampled CRR trajectory. This is an environment score only: it is
#: computed under no_grad and cannot form a pathwise decoder gradient.
latent_policy_gold_logprob: torch.FloatTensor | None = None # [B]
#: The same frozen gold-answer score at every decoded CRR prefix
#: Z_k = B + mean(C_1..C_k). This is populated only when latent RL asks
#: for transition-aware credit assignment; its last column is exactly
#: ``latent_policy_gold_logprob``.
latent_policy_step_gold_logprob: torch.FloatTensor | None = None # [B, T]
#: Dimension-normalized Gaussian score log-density for C2..CT.
latent_policy_log_probs: torch.FloatTensor | None = None # [B, T-1]
latent_policy_mean_update_rms: torch.FloatTensor | None = None # [B, T-1]
latent_policy_noise_rms: torch.FloatTensor | None = None # [B, T-1]
@dataclass
class LatentPolicyTrace:
"""Training-only score-function trace for stochastic CRR transitions.
The sampled recurrent states themselves are detached. Consequently the
only differentiable quantity returned to the RL Trainer is ``log_probs``;
answer generation and reward computation cannot create a pathwise gradient
through the frozen decoder.
"""
log_probs: torch.FloatTensor # [B, T-1]
mean_update_rms: torch.FloatTensor # [B, T-1], relative to C1 RMS
noise_rms: torch.FloatTensor # [B, T-1], relative to C1 RMS
class CRRLatentPolicyAdapter(nn.Module):
"""Zero-init low-rank correction used only by CRR states ``C2..CT``.
This is intentionally narrower than adapting the shared native L21: L21
also constructs B/C1 and processes generated answer tokens. A separate
latent-only adapter keeps those competence-critical paths frozen and makes
the score-function action policy the only behavior changed by RL.
"""
def __init__(self, width: int, rank: int, alpha: float):
super().__init__()
if rank < 1:
raise ValueError("CRRLatentPolicyAdapter rank must be >= 1")
self.down = nn.Linear(width, rank, bias=False)
self.up = nn.Linear(rank, width, bias=False)
self.scaling = float(alpha) / float(rank)
def forward(self, state: torch.Tensor) -> torch.Tensor:
return state + self.up(self.down(state)) * self.scaling
class CRRFeaturewiseAggregator(nn.Module):
"""Native-basis-preserving linear aggregation over recurrent depth.
This is the block-diagonal special case of ``Linear(T * d, d)``: block
``k`` is a learned diagonal matrix instead of a dense ``d x d`` matrix.
It can therefore select a different recurrent depth for every native
hidden feature without rotating the pretrained decoder basis or adding a
51M-parameter all-reduce at ``d=3584, T=4``.
"""
def __init__(self, steps: int, width: int):
super().__init__()
if steps < 1 or width < 1:
raise ValueError("steps and width must both be positive")
self.steps = int(steps)
self.width = int(width)
self.weight = nn.Parameter(torch.empty(self.steps, self.width))
self.reset_parameters()
def reset_parameters(self) -> None:
# Exactly reproduces the historical mean interface at initialization.
with torch.no_grad():
self.weight.fill_(1.0 / self.steps)
def forward(self, concatenated: torch.Tensor) -> torch.Tensor:
expected = self.steps * self.width
if concatenated.shape[-1] != expected:
raise ValueError(
f"expected concatenated width {expected}, got "
f"{concatenated.shape[-1]}"
)
states = concatenated.reshape(
*concatenated.shape[:-1], self.steps, self.width
)
return (states * self.weight).sum(dim=-2)
# ---------------------------------------------------------------------------
# the model
# ---------------------------------------------------------------------------
class CloseQwen2_5_VLForConditionalGeneration(Qwen2_5_VLPreTrainedModel):
def _init_weights(self, module):
super()._init_weights(module)
if isinstance(module, SpatialVisualRecurrentCell):
# Generic HF initialization visits the child Linears first. Zero
# the residual writer only after that sweep so T=8 is the exact
# base visual path at initialization and on missing-key loads.
module.zero_output()
if isinstance(module, CRRFeaturewiseAggregator):
# Also runs for missing-key initialization under from_pretrained.
module.reset_parameters()
# ER v2.2 C: aggregation 파라미터 init. post_init 스윕(from_backbone)과
# missing-key 초기화(from_pretrained) 모두 이 훅을 지나므로 두 로드
# 경로에서 동일 값이 보장된다 (__init__ 직접 대입은 meta 경유에서 유실).
if module is self and getattr(
self.config, "concat_aggregation", False
) and hasattr(self, "reasoning_aggregator"):
# 개정 init (advisor 2026-08-12): W = [I/T, ..., I/T] — step 0에서
# R_agg = mean(R1..RT) 정확히. (구판 [0,..,0,I]는 폐기; random 금지.)
# I/4 = 0.25는 bf16 정확 표현 — 반올림 오차 없음. 네 블록은 이후
# answer CE로 독립 분화.
_T = int(self.config.num_workspace_steps)
d_w = int(self.reasoning_aggregator.out_features)
with torch.no_grad():
w = self.reasoning_aggregator.weight # [D, T*D]
w.zero_()
_eye = torch.eye(d_w, dtype=w.dtype, device=w.device) / _T
for _i in range(_T):
w[:, _i * d_w:(_i + 1) * d_w].copy_(_eye)
if module is self and getattr(self.config, "temporal_aggregation", False) \
and hasattr(self, "temporal_logits"):
_T = int(self.config.num_workspace_steps)
_w0 = torch.tensor([0.4, 0.3, 0.2, 0.1][:_T]) if _T <= 4 \
else torch.full((_T,), 1.0 / _T)
_w0 = _w0 / _w0.sum()
with torch.no_grad():
self.temporal_logits.copy_(_w0.log().to(self.temporal_logits.dtype))
if module is self and hasattr(self, "crr_policy_adapter"):
# Exact identity at the SFT -> latent-RL handoff. Down follows the
# backbone initializer; zero Up is the standard LoRA-style start.
with torch.no_grad():
self.crr_policy_adapter.up.weight.zero_()
"""CLOSE: Qwen2.5-VL with its upper visual path replaced by a read-once workspace."""
config_class = CloseQwen2_5_VLConfig
_no_split_modules = ["Qwen2_5_VLDecoderLayer", "Qwen2_5_VLVisionBlock"]
def __init__(self, config: CloseQwen2_5_VLConfig, backbone=None):
super().__init__(config)
config.validate()
spatial_visual = bool(
getattr(config, "spatial_visual_recurrence", False)
)
perceive_deliberate = bool(
getattr(config, "perceive_deliberate_chain", False)
)
local_visual_chain = bool(
getattr(config, "local_visual_latent_chain", False)
)
if not spatial_visual:
config._require_ell_star()
# The host model. Pass an already-loaded one via :meth:`from_backbone`;
# `CloseQwen2_5_VLForConditionalGeneration.from_pretrained()`
# does NOT work, because base checkpoint keys are "model.*" / "lm_head.*"
# while this tree needs "backbone.model.*". transformers only warns about
# that, leaving all 8.3B parameters random and every forward NaN.
self.backbone = (
backbone
if backbone is not None
else Qwen2_5_VLForConditionalGeneration._from_config(config)
)
raw_eq = bool(getattr(config, "raw_evidence_question", False))
crr = bool(
getattr(config, "counterfactual_residual_recurrence", False)
)
q_state_interface = raw_eq or crr
if perceive_deliberate:
# LVLC uses one homogeneous latent type; historical PDLC uses two
# LOOK/THINK types. Distinct native M-RoPE positions carry order
# in both one-pass graphs.
self.latent_chain_embeddings = nn.Embedding(
1 if local_visual_chain else 2,
config.backbone_width,
)
self.workspace_out = nn.Identity()
elif spatial_visual:
self.spatial_visual_cell = SpatialVisualRecurrentCell(
config.backbone_width,
int(config.spatial_recurrence_inner_width),
int(config.spatial_recurrence_heads),
rms_norm_eps=float(config.text_config.rms_norm_eps),
residual_scale=float(config.counterfactual_beta),
)
self.workspace_out = nn.Identity()
elif raw_eq:
# Replacement, not an addition: the raw-evidence path allocates no
# slot reader, slot transition, evidence slots, reasoning slots,
# mid-read module, or learned trajectory aggregator.
self.raw_eq_transition = RawEvidenceQuestionTransition(config)
self.workspace_out = nn.Identity()
elif crr:
# CRR reuses the native L(ell_star+1) module itself. No copied
# transition, visual memory, slots, projection, or aggregation is
# allocated. RL checkpoints may add one shared low-rank correction
# that is called only for C2..CT, never for B/C1/token decoding.
self.workspace_out = nn.Identity()
if getattr(config, "crr_concat_aggregation", False):
width = int(config.backbone_width)
steps = int(config.num_workspace_steps)
self.reasoning_aggregator = CRRFeaturewiseAggregator(
steps, width
)
if int(getattr(config, "crr_policy_rank", 0)) > 0:
self.crr_policy_adapter = CRRLatentPolicyAdapter(
config.backbone_width,
int(config.crr_policy_rank),
float(config.crr_policy_alpha),
)
else:
self.workspace_read = WorkspaceRead(config)
self.workspace_transition = WorkspaceTransition(config)
self.workspace_out = nn.Linear(
config.workspace_width, config.backbone_width
)
# -- ER v2 (explicit evidence read; advisor 2026-08-11) --------------
# R이 E를 별도 memory로 읽는 pre-norm cross-attn 잔차 + 질문 조건화 R^0.
if getattr(config, "explicit_evidence_read", False):
# ER v2 canonical (§5): R이 persistent E를 읽는 명시적 cross-attention.
# query=LN(R), key/value=LN(E). 전 step 공유 (step별 복제 금지, §14).
assert config.workspace_width == config.backbone_width, (
"explicit_evidence_read는 native width(d_w == backbone d) 전제 (§2)")
d_w = config.workspace_width
self.evidence_read_norm_q = nn.LayerNorm(d_w)
self.evidence_read_norm_kv = nn.LayerNorm(d_w)
self.evidence_read = nn.MultiheadAttention(
d_w, config.workspace_num_heads, dropout=0.0, batch_first=True)
# §8: 차원-변환 전용 P_Z는 native width에서 Identity.
# interface 보정 수신자 = splice_gain + post-l* LoRA (§9.5에서 검증).
self.workspace_out = nn.Identity()
if getattr(config, "r_init_question", False):
d_w = config.workspace_width
# Legacy config name, current semantics: R^0 is learnable slots
# only. Q* first enters through the shared transition.
self.reasoning_slots = nn.Parameter(
torch.randn(config.num_workspace_slots, d_w) * 0.02
)
if getattr(config, "concat_aggregation", False):
# ER v2.2 C 확정판 (§6-8): R_agg = Linear(Concat(R1..RT, dim=hidden)).
# W ∈ [D, T·D], bias 없음. 초기 W=[I/T,...,I/T], so
# R_agg=mean(R1..RT). Temporal softmax/gates are absent.
assert getattr(config, "er_single_chain", False), (
"concat_aggregation은 er_single_chain=True 전제 (§12)")
_T = int(config.num_workspace_steps)
d_w = config.workspace_width
self.reasoning_aggregator = nn.Linear(_T * d_w, d_w, bias=False)
# 값(zero + 마지막 블록 Identity)은 _init_weights(self)에서 설정 —
# __init__ 직접 대입은 from_pretrained meta 경유에서 유실 (실증됨).
if getattr(config, "temporal_aggregation", False):
# (legacy — softmax 변형; C 확정판에서 미사용, §7 금지 목록)
assert getattr(config, "er_single_chain", False), (
"temporal_aggregation은 er_single_chain=True 전제 (§12)")
assert not getattr(config, "concat_aggregation", False), "동시 사용 금지"
_T = int(config.num_workspace_steps)
self.temporal_logits = nn.Parameter(torch.zeros(_T))
# Placeholder embedding for the S slot positions in the text-only prefill.
# Cache-shape plumbing below ell*: these rows are overwritten by R_agg
# before the upper decoder. They are a fixed text-only placeholder, not
# an extra learned prompt outside the declared workspace/LoRA method.
if not q_state_interface and not spatial_visual and not perceive_deliberate:
self.slot_embedding = nn.Parameter(
torch.randn(config.num_workspace_slots, config.backbone_width) * 0.02,
requires_grad=False,
)
# RMS splice gain is retained for historical interface-loss checkpoints.
# In the current answer-only method it is a fixed 1, not an undeclared
# trainable decoder-interface scalar (and bf16 updates near 1 would be
# numerically ineffective at the configured learning rate anyway).
self.splice_gain = (
nn.Parameter(torch.ones(1), requires_grad=bool(config.interface_loss))
if (
not q_state_interface
and not spatial_visual
and not perceive_deliberate
and getattr(config, "splice_scale_match", False)
)
else None
)
self.post_init()
if perceive_deliberate:
init_id = int(getattr(config, "latent_chain_init_token_id", -1))
if init_id < 0:
init_id = 0
with torch.no_grad():
source = self.vl.get_input_embeddings().weight[init_id]
self.latent_chain_embeddings.weight.copy_(
source.unsqueeze(0).expand_as(
self.latent_chain_embeddings.weight
)
)
if raw_eq:
self._initialize_raw_attention_from_upper_layer()
# peft wraps the targeted Linears in place, so this must run after
# post_init: the adapters start as an exact no-op (B initialised to 0)
# and must not be re-initialised by the backbone's weight init.
self.adapter_config = None
if config.adapter_rank > 0:
from peft import inject_adapter_in_model
self.adapter_config = build_adapter_config(config)
# Inject at `backbone`, not `text_model`: peft matches layer indices
# with `.*\.layers\.(\d+)\.`, which needs a dot before "layers".
# From text_model the keys start "layers.0..." and nothing matches;
# from backbone they are "model.language_model.layers.0...".
inject_adapter_in_model(self.adapter_config, self.backbone)
if getattr(config, "recurrence_only_adapter", False):
# Execution defaults to the frozen pretrained path. The
# recurrent loop temporarily enables these same trainable
# parameters for k>=2 without mutating requires_grad.
from peft.tuners.tuners_utils import BaseTunerLayer
for module in self.backbone.modules():
if isinstance(module, BaseTunerLayer):
module._disable_adapters = True
def _initialize_raw_attention_from_upper_layer(self) -> None:
"""Copy the first upper decoder layer's visual-read-compatible basis.
L(ell_star+1) is pretrained to consume layer-ell_star hidden states,
exactly the feature space occupied by both Q* and V*. Cross-attention
remains position-free, but its Q/K/V/O maps and adjacent RMSNorms start
from that decoder layer instead of an unrelated random basis.
This runs after ``post_init`` (so HF initialization cannot overwrite
the copy) and before LoRA injection (so source projections are plain
Linears and the raw transition receives only the frozen base weights).
"""
source_idx = int(getattr(self.config, "raw_attention_source_layer", -1))
if source_idx < 0:
source_idx = int(self.config.ell_star) + 1
source_layer = self.text_model.layers[source_idx]
source_attn = source_layer.self_attn
target = self.raw_eq_transition
pairs = (
(target.q_proj, source_attn.q_proj),
(target.k_proj, source_attn.k_proj),
(target.v_proj, source_attn.v_proj),
(target.o_proj, source_attn.o_proj),
)
with torch.no_grad():
for dst, src in pairs:
if dst.weight.shape != src.weight.shape:
raise ValueError(
f"raw GQA/source shape mismatch at L{source_idx}: "
f"{tuple(dst.weight.shape)} vs {tuple(src.weight.shape)}"
)
dst.weight.copy_(src.weight)
if dst.bias is not None:
if src.bias is None:
raise ValueError("raw GQA expects the source Q/K/V biases")
dst.bias.copy_(src.bias)
target.read_norm_q.weight.copy_(source_layer.input_layernorm.weight)
target.read_norm_e.weight.copy_(source_layer.input_layernorm.weight)
if target.use_ffn:
target.ffn_norm.weight.copy_(
source_layer.post_attention_layernorm.weight
)
self.raw_attention_source_layer = source_idx
@classmethod
def from_backbone(cls, pretrained: str, config: CloseQwen2_5_VLConfig, **kwargs):
"""Build CLOSE around a real Qwen2.5-VL checkpoint.
The backbone is loaded first and handed to ``__init__`` so its weights
are never routed through this class's key namespace. Use this, not
``from_pretrained``, when starting from the base model;
``from_pretrained`` remains correct for reloading a *saved CLOSE*
checkpoint, whose keys already match.
"""
backbone = Qwen2_5_VLForConditionalGeneration.from_pretrained(pretrained, **kwargs)
model = cls(config, backbone=backbone)
return model.to(backbone.dtype)
# -- convenience views -------------------------------------------------
@property
def vl(self):
"""``Qwen2_5_VLModel`` (visual + language_model)."""
return self.backbone.model
@property
def text_model(self):
"""``Qwen2_5_VLTextModel``."""
return self.backbone.model.language_model
def _uses_question_state_interface(self) -> bool:
"""Whether the replacement prefix has exactly the question rows."""
return bool(
getattr(self.config, "raw_evidence_question", False)
or getattr(
self.config, "counterfactual_residual_recurrence", False
)
)
def _upper_decoder_start(self) -> int:
"""First layer that consumes the image-dependent decoder latent.
CRR turns L(ell_star+1) into its shared recurrent operator. That layer
still processes answer tokens, but its prefix KV is populated only by
the text-only scaffold; the image-dependent prefix starts at the next
layer. Historical methods continue to splice at ell_star+1.
"""
return int(self.config.ell_star) + (
2
if getattr(
self.config, "counterfactual_residual_recurrence", False
)
else 1
)
# -- parameter freezing ------------------------------------------------
def freeze_backbone(self) -> dict[str, int]:
"""Freeze everything, then open exactly the declared method parameters.
§3.3: "Vision encoder와 기본 backbone은 동결하고 workspace module,
projection, post-l* low-rank adapter만 학습".
LoRA parameters live *inside* ``backbone`` after peft injection, so they
are re-enabled by name after the blanket freeze.
"""
for p in self.parameters():
p.requires_grad_(False)
local_visual_chain = bool(
getattr(self.config, "local_visual_latent_chain", False)
)
# LVLC executes solely through the native Transformer, its homogeneous
# latent embedding, and LoRA. Opening legacy recurrent workspace
# modules here creates dead trainable parameters that never participate
# in the one-pass graph.
method_modules = (
(getattr(self, "latent_chain_embeddings", None),)
if local_visual_chain
else (
getattr(self, "latent_chain_embeddings", None),
getattr(self, "spatial_visual_cell", None),
getattr(self, "workspace_read", None),
getattr(self, "workspace_transition", None),
getattr(self, "raw_eq_transition", None),
self.workspace_out,
)
)
for module in method_modules:
if module is None:
continue
for p in module.parameters():
p.requires_grad_(True)
if not local_visual_chain:
for name in (
"evidence_read_norm_q",
"evidence_read_norm_kv",
"evidence_read",
"reasoning_aggregator",
):
module = getattr(self, name, None)
if module is not None:
for p in module.parameters():
p.requires_grad_(True)
for name in ("reasoning_slots", "temporal_logits"):
p = getattr(self, name, None)
if p is not None:
p.requires_grad_(True)
if self.splice_gain is not None and getattr(
self.config, "interface_loss", False
):
self.splice_gain.requires_grad_(True)
n_lora = 0
for name, p in self.backbone.named_parameters():
if "lora_" in name:
p.requires_grad_(True)
n_lora += p.numel()
n_train = sum(p.numel() for p in self.parameters() if p.requires_grad)
n_total = sum(p.numel() for p in self.parameters())
if local_visual_chain:
latent_parameter = self.latent_chain_embeddings.weight
expected_trainable = n_lora + latent_parameter.numel()
if not latent_parameter.requires_grad or n_train != expected_trainable:
raise RuntimeError(
"LVLC must train exactly its latent embedding and LoRA: "
f"latent_requires_grad={latent_parameter.requires_grad}, "
f"latent_numel={latent_parameter.numel()}, lora={n_lora}, "
f"actual_trainable={n_train}, expected={expected_trainable}"
)
return {"trainable": n_train, "total": n_total, "lora": n_lora}
def _aggregate_crr_residuals(
self, residual_steps: list[torch.Tensor] | tuple[torch.Tensor, ...]
) -> torch.Tensor:
"""Aggregate one CRR residual prefix with a stable initial interface.
The learned full-horizon path is ``Linear(Concat(C1..CT))``. A shorter
evaluation/training prefix zero-pads its missing future blocks and is
multiplied by ``T/k``; at the prescribed ``[I/T,...,I/T]``
initialization this remains exactly ``mean(C1..Ck)``. Historical CRR
checkpoints keep their parameter-free mean because the flag defaults
to false.
"""
steps = list(residual_steps)
if not steps:
raise ValueError("CRR aggregation requires at least one residual")
horizon = int(self.config.num_workspace_steps)
if len(steps) > horizon:
raise ValueError(
f"received {len(steps)} CRR states for horizon {horizon}"
)
if not getattr(self.config, "crr_concat_aggregation", False):
return torch.stack(steps, dim=0).mean(dim=0)
if not hasattr(self, "reasoning_aggregator"):
raise RuntimeError(
"crr_concat_aggregation=True but reasoning_aggregator is missing"
)
if len(steps) < horizon:
steps.extend(torch.zeros_like(steps[0]) for _ in range(horizon - len(steps)))
aggregate = self.reasoning_aggregator(torch.cat(steps, dim=-1))
prefix_steps = len(residual_steps)
if prefix_steps < horizon:
aggregate = aggregate * (horizon / prefix_steps)
return aggregate
def decoder_latent_from_states(
self, states: list[torch.Tensor] | tuple[torch.Tensor, ...]
) -> torch.Tensor:
"""Return the exact pre-interface latent decoded by the current model.
Keeping this operation on the model prevents evaluation scripts from
silently targeting a stored state that the selected interface does not
actually decode. Residual Raw-E/Q returns ``R_T``; state replacement
returns ``Q* + H_T``; question-anchored replacement returns ``Z_T``;
historical concat models reconstruct their learned ``R_agg``.
"""
if not states:
raise ValueError("states must not be empty")
if getattr(
self.config, "counterfactual_residual_recurrence", False
):
# states = [B, C1, ..., CT], where B=F(Q_text) is strictly
# image-independent and every Ck is an image-conditioned residual.
if len(states) != self.config.num_workspace_steps + 1:
raise ValueError(
f"expected anchor plus {self.config.num_workspace_steps} "
f"counterfactual residuals, got {len(states)} tensors"
)
residual_steps = list(states[1:])
# Evaluation-only prefix sweep. Step k uses the same interface
# trained by L_step_keep: Z_k = B + mean(C_1..C_k), rather than
# decoding C_k alone. Normal training/inference has no override.
k_ov = int(getattr(self.config, "er_step_override", 0) or 0)
if k_ov:
if not 1 <= k_ov <= self.config.num_workspace_steps:
raise ValueError(
"er_step_override must be in "
f"[1,{self.config.num_workspace_steps}]"
)
residual_steps = residual_steps[:k_ov]
if getattr(
self.config, "visual_cumulative_recurrence", False
):
# Cumulative CVRR stores C_k = R_k - B and exposes the actual
# recurrent state R_k. There is no temporal mean that can
# dilute or cancel the accumulated visual corrections.
return states[0] + residual_steps[-1]
return states[0] + self._aggregate_crr_residuals(residual_steps)
if getattr(self.config, "raw_evidence_question", False):
# states = [R0, R1, ..., RT] in residual mode and
# [Q*, H1, ..., HT] in state-replacement mode.
if len(states) != self.config.num_workspace_steps + 1:
raise ValueError(
f"expected R0 plus {self.config.num_workspace_steps} recurrent "
f"states, got {len(states)} tensors"
)
if getattr(self.config, "raw_mean_aggregation", False):
return torch.stack(list(states[1:]), dim=0).mean(dim=0)
if getattr(self.config, "raw_state_replacement", False):
if getattr(self.config, "raw_question_anchor", False):
return states[-1]
return states[0] + states[-1]
return states[-1]
if not getattr(self.config, "evidence_reasoning", False):
return states[-1]
r_steps = list(states[1:]) # states[0] is persistent evidence E0
if getattr(self.config, "concat_aggregation", False):
if len(r_steps) != self.config.num_workspace_steps:
raise ValueError(
f"expected {self.config.num_workspace_steps} reasoning states, "
f"got {len(r_steps)}"
)
return self.reasoning_aggregator(torch.cat(r_steps, dim=-1))
if getattr(self.config, "temporal_aggregation", False):
a = torch.softmax(self.temporal_logits.float(), dim=0).to(r_steps[0].dtype)
return torch.stack(r_steps, dim=0).mul(
a.view(-1, *([1] * r_steps[0].ndim))
).sum(dim=0)
return states[-1]
def decoder_latent_at_depth(
self,
states: list[torch.Tensor] | tuple[torch.Tensor, ...],
depth: int | torch.Tensor,
) -> torch.Tensor:
"""Decode a CRR prefix at a training-selected recurrent depth.
``depth`` may be one scalar for the batch or one value per example.
The latter lets annotated GQA rows follow the program curriculum while
unannotated rows keep the normal full-horizon answer path. No new
readout module is introduced: this is the same native decoder interface
used by the existing early-stop evaluation.
"""
if not getattr(
self.config, "counterfactual_residual_recurrence", False
):
raise ValueError("depth-selected decoding is implemented only for CRR")
horizon = int(self.config.num_workspace_steps)
if len(states) != horizon + 1:
raise ValueError(
f"expected anchor plus {horizon} CRR states, got {len(states)}"
)
def latent_for(prefix_depth: int) -> torch.Tensor:
if not 1 <= prefix_depth <= horizon:
raise ValueError(
f"curriculum readout depth must be in [1,{horizon}], "
f"got {prefix_depth}"
)
residual_steps = list(states[1 : prefix_depth + 1])
if getattr(
self.config, "visual_cumulative_recurrence", False
):
return states[0] + residual_steps[-1]
return states[0] + self._aggregate_crr_residuals(residual_steps)
if not isinstance(depth, torch.Tensor):
return latent_for(int(depth))
if depth.ndim == 0:
return latent_for(int(depth.item()))
if depth.ndim != 1 or depth.shape[0] != states[0].shape[0]:
raise ValueError(
"per-example curriculum depth must have shape [B], got "
f"{tuple(depth.shape)} for batch {states[0].shape[0]}"
)
depth = depth.to(device=states[0].device, dtype=torch.long)
if bool(((depth < 1) | (depth > horizon)).any()):
raise ValueError(
f"curriculum readout depths must lie in [1,{horizon}]"
)
candidates = torch.stack(
[latent_for(prefix_depth) for prefix_depth in range(1, horizon + 1)],
dim=1,
)
rows = torch.arange(candidates.shape[0], device=candidates.device)
return candidates[rows, depth - 1]
# -- non-recurrent perceive/deliberate latent chain -------------------
def _perceive_deliberate_inputs(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
attention_mask: torch.Tensor | None,
question_attention_mask: torch.Tensor | None,
answer_input_ids: torch.LongTensor | None = None,
answer_input_attention_mask: torch.Tensor | None = None,
) -> tuple[
torch.Tensor,
torch.LongTensor,
torch.BoolTensor,
PerceiveDeliberateLayout,
torch.Tensor,
torch.Tensor,
]:
"""Materialize one native block-sparse LOOK/THINK sequence."""
if attention_mask is None:
attention_mask = torch.ones_like(input_ids, dtype=torch.long)
if question_attention_mask is None:
question_attention_mask = torch.ones_like(
question_ids, dtype=torch.long
)
if input_ids.shape[0] != question_ids.shape[0]:
raise ValueError("multimodal and clean-question batch sizes differ")
multimodal_embeddings, multimodal_positions = embed_multimodal(
self.vl,
input_ids,
pixel_values,
image_grid_thw,
attention_mask,
)
question_embeddings = self.vl.get_input_embeddings()(question_ids)
batch_size = input_ids.shape[0]
num_pairs = int(self.config.num_workspace_steps) // 2
if getattr(self.config, "local_visual_latent_chain", False):
latent_type_ids = torch.zeros(
2 * num_pairs,
device=input_ids.device,
dtype=torch.long,
)
else:
latent_type_ids = torch.tensor(
[0, 1] * num_pairs,
device=input_ids.device,
dtype=torch.long,
)
latent_embeddings = self.latent_chain_embeddings(
latent_type_ids
).unsqueeze(0).expand(batch_size, -1, -1)
latent_embeddings = latent_embeddings.to(multimodal_embeddings.dtype)
answer_input_length = (
0 if answer_input_ids is None else int(answer_input_ids.shape[1])
)
layout = PerceiveDeliberateLayout(
multimodal_length=input_ids.shape[1],
question_length=question_ids.shape[1],
num_pairs=num_pairs,
answer_input_length=answer_input_length,
)
pieces = [
multimodal_embeddings,
question_embeddings.to(multimodal_embeddings.dtype),
latent_embeddings,
]
if answer_input_ids is not None:
pieces.append(
self.vl.get_input_embeddings()(answer_input_ids).to(
multimodal_embeddings.dtype
)
)
embeddings = torch.cat(pieces, dim=1)
position_ids = build_perceive_deliberate_positions(
layout,
multimodal_positions,
attention_mask,
question_attention_mask,
answer_input_attention_mask,
)
visibility = build_perceive_deliberate_mask(
layout,
attention_mask,
question_attention_mask,
answer_input_attention_mask,
multimodal_visual_mask=(
image_token_mask(input_ids, int(self.config.image_token_id))
| image_token_mask(input_ids, int(self.config.video_token_id))
),
local_visual_chain=bool(
getattr(self.config, "local_visual_latent_chain", False)
),
transition_conditioned=bool(
getattr(self.config, "pdlc_transition_conditioned", False)
),
question_visible_through_pair=int(
getattr(
self,
"_pdlc_question_visible_through_pair_override",
getattr(
self.config,
"pdlc_question_visible_through_pair",
1,
),
)
),
disable_chain_links=bool(
getattr(self.config, "pdlc_disable_chain_links", False)
),
disable_late_visual=bool(
getattr(self.config, "pdlc_disable_late_visual", False)
),
final_pair_only=bool(
getattr(self.config, "pdlc_final_pair_only", False)
),
)
return (
embeddings,
position_ids,
visibility,
layout,
attention_mask,
question_attention_mask,
)
@staticmethod
def _with_attention_graph(ctx, visibility: torch.BoolTensor):
"""Replace every native attention type with one explicit graph."""
mappings = {
attention_type: visibility
for attention_type in ctx.causal_mask_mapping
}
return replace(ctx, causal_mask_mapping=mappings)
def _run_perceive_deliberate_training_layers(self, ctx) -> torch.Tensor:
"""Run the native stack with explicit per-layer checkpointing.
Calling decoder layers directly is required for the custom graph, so
the stock text-model wrapper cannot install its gradient-checkpoint
closures for us. Mirror that behavior here; otherwise gradients to
the LOOK/THINK embeddings retain all 28 layer activations and erase
most of the B200 batch-size advantage.
"""
hidden = ctx.hidden_states
checkpoint_layers = bool(
self.training
and getattr(self.text_model, "gradient_checkpointing", False)
)
for layer in self.text_model.layers:
attention_type = getattr(layer, "attention_type", "full_attention")
layer_mask = ctx.causal_mask_mapping[attention_type]
def layer_forward(states, *, current_layer=layer, mask=layer_mask):
output = current_layer(
states,
attention_mask=mask,
position_ids=ctx.text_position_ids,
past_key_values=None,
use_cache=False,
cache_position=ctx.cache_position,
position_embeddings=ctx.position_embeddings,
)
return output[0] if isinstance(output, tuple) else output
if checkpoint_layers:
hidden = checkpoint(
layer_forward,
hidden,
use_reentrant=False,
)
else:
hidden = layer_forward(hidden)
return hidden
def _forward_perceive_deliberate_chain(
self,
input_ids: torch.LongTensor,
attention_mask: torch.Tensor | None,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
question_attention_mask: torch.Tensor | None,
labels: torch.LongTensor,
answer_ids: torch.LongTensor,
) -> CloseCausalLMOutput:
"""Teacher-forced answer CE through one native LOOK/THINK graph."""
if labels.ndim != 2 or answer_ids.ndim != 2:
raise ValueError("labels and answer_ids must have shape [B, N_a]")
if labels.shape != answer_ids.shape or labels.shape[1] < 1:
raise ValueError(
"labels and answer_ids must share one non-empty answer shape"
)
answer_input_ids = answer_ids[:, :-1]
answer_input_mask = labels[:, :-1].ne(-100)
(
embeddings,
position_ids,
visibility,
layout,
_,
_,
) = self._perceive_deliberate_inputs(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
answer_input_ids,
answer_input_mask,
)
ctx = make_split_context(
self.text_model,
embeddings,
position_ids,
attention_mask=None,
)
ctx = self._with_attention_graph(ctx, visibility)
hidden = self._run_perceive_deliberate_training_layers(ctx)
hidden = final_norm(self.text_model, hidden)
# The final THINK predicts the first answer token. Teacher-forced
# answer rows predict all subsequent tokens, matching the repository's
# already-aligned labels convention without an internal label shift.
seed = hidden[:, layout.think_indices[-1] : layout.think_indices[-1] + 1]
if layout.answer_input_length:
prediction_hidden = torch.cat(
[seed, hidden[:, layout.answer_slice]], dim=1
)
else:
prediction_hidden = seed
logits = self.backbone.lm_head(prediction_hidden)
loss = answer_cross_entropy(
logits,
labels,
per_example=bool(
getattr(self.config, "per_example_answer_loss", False)
),
)
latent_hidden = hidden[:, layout.latent_start : layout.answer_start]
expected_latents = int(self.config.num_workspace_steps)
if latent_hidden.shape[1] != expected_latents:
raise RuntimeError(
"latent layout produced an invalid state count: "
f"expected {expected_latents}, got {latent_hidden.shape[1]}"
)
with torch.no_grad():
normalized = nn.functional.normalize(latent_hidden.float(), dim=-1)
adjacent = 1.0 - (normalized[:, 1:] * normalized[:, :-1]).sum(-1)
look = normalized[:, 0::2]
think = normalized[:, 1::2]
pair = 1.0 - (look * think).sum(-1)
think_progress = (
1.0 - (think[:, 1:] * think[:, :-1]).sum(-1)
if think.shape[1] > 1
else think.new_empty(think.shape[0], 0)
)
if getattr(self.config, "local_visual_latent_chain", False):
self.last_local_visual_latent_telemetry = {
"transition_drift": adjacent.mean(dim=0),
}
else:
self.last_perceive_deliberate_telemetry = {
"adjacent_drift": adjacent.mean(dim=0),
"pair_drift": pair.mean(dim=0),
"think_progress": think_progress.mean(dim=0),
}
states = tuple(
latent_hidden[:, index : index + 1]
for index in range(latent_hidden.shape[1])
)
return CloseCausalLMOutput(
loss=loss,
logits=logits,
past_key_values=None,
workspace_states=states,
v_star_pooled=None,
)
@torch.no_grad()
def _generate_perceive_deliberate_chain(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
attention_mask: torch.Tensor | None,
question_attention_mask: torch.Tensor | None,
*,
max_new_tokens: int,
eos_token_id: int | None,
) -> torch.LongTensor:
"""Greedy generation with the multimodal/LOOK cache rows masked."""
if max_new_tokens < 1:
raise ValueError("max_new_tokens must be >= 1")
(
embeddings,
position_ids,
visibility,
layout,
_,
question_attention_mask,
) = self._perceive_deliberate_inputs(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
)
cache = DynamicCache(config=self.text_model.config)
ctx = make_split_context(
self.text_model,
embeddings,
position_ids,
attention_mask=None,
past_key_values=cache,
)
ctx = self._with_attention_graph(ctx, visibility)
hidden = run_layer_range(
self.text_model, ctx, 0, None, use_cache=True
)
hidden = final_norm(self.text_model, hidden)
seed = hidden[:, layout.think_indices[-1] : layout.think_indices[-1] + 1]
nxt = self.backbone.lm_head(seed).argmax(dim=-1)
output = [nxt]
batch_size = input_ids.shape[0]
done = torch.zeros(
batch_size, dtype=torch.bool, device=input_ids.device
)
if eos_token_id is not None:
done |= nxt.squeeze(1).eq(eos_token_id)
final_prefix_position = position_ids[:, :, -1]
for step in range(1, max_new_tokens):
if bool(done.all()):
break
token_embeddings = self.vl.get_input_embeddings()(nxt)
cache_position = torch.tensor(
[layout.answer_start + step - 1], device=input_ids.device
)
token_positions = (
final_prefix_position + step
).unsqueeze(-1)
generated_mask = torch.ones(
batch_size,
step,
dtype=question_attention_mask.dtype,
device=input_ids.device,
)
answer_visibility = build_answer_decode_mask(
layout,
question_attention_mask,
generated_mask,
local_visual_chain=bool(
getattr(
self.config,
"local_visual_latent_chain",
False,
)
),
final_pair_only=bool(
getattr(self.config, "pdlc_final_pair_only", False)
),
transition_conditioned=bool(
getattr(
self.config,
"pdlc_transition_conditioned",
False,
)
),
)
token_ctx = make_split_context(
self.text_model,
token_embeddings,
token_positions,
attention_mask=None,
past_key_values=cache,
cache_position=cache_position,
)
token_ctx = self._with_attention_graph(
token_ctx, answer_visibility
)
token_hidden = run_layer_range(
self.text_model,
token_ctx,
0,
None,
use_cache=True,
)
token_hidden = final_norm(self.text_model, token_hidden)
nxt = self.backbone.lm_head(token_hidden).argmax(dim=-1)
if eos_token_id is not None:
nxt = torch.where(
done.unsqueeze(1),
nxt.new_full(nxt.shape, eos_token_id),
nxt,
)
done |= nxt.squeeze(1).eq(eos_token_id)
output.append(nxt)
return torch.cat(output, dim=1)
# -- spatial recurrent visual field -----------------------------------
def _spatial_horizon(self) -> int:
"""Return the configured train horizon or an explicit eval prefix."""
override = int(getattr(self.config, "spatial_step_override", -1))
if override >= 0:
return override
return int(self.config.spatial_recurrence_steps)
def _spatial_image_to_batch(
self,
input_ids: torch.LongTensor,
attention_mask: torch.Tensor | None,
num_images: int,
) -> torch.LongTensor:
"""Map Qwen's batch-major image feature list back to sample rows."""
if input_ids.ndim != 2:
raise ValueError("spatial visual recurrence expects [B,L] input_ids")
starts = input_ids[:, :-1].eq(self.config.vision_start_token_id)
images = input_ids[:, 1:].eq(self.config.image_token_id)
markers = starts & images
if attention_mask is not None:
markers &= attention_mask[:, :-1].gt(0)
markers &= attention_mask[:, 1:].gt(0)
counts = markers.sum(dim=-1).long()
if int(counts.sum()) != int(num_images):
raise ValueError(
"image_grid_thw rows do not match "
f"segments: grids={num_images}, segments={int(counts.sum())}"
)
return torch.repeat_interleave(
torch.arange(input_ids.shape[0], device=input_ids.device), counts
)
def _spatial_visual_prefix_embeddings(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
attention_mask: torch.Tensor | None,
question_attention_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Build the ordinary Qwen prefix with image rows replaced by V_T.
Vision and token embeddings are frozen inputs. Gradients start at the
shared spatial cell, continue through the newly built full decoder
pass, and never enter a cache from visual-state construction.
"""
if image_grid_thw is None or image_grid_thw.ndim != 2:
raise ValueError("image_grid_thw is required for spatial recurrence")
if question_ids.ndim != 2 or question_ids.shape[0] != input_ids.shape[0]:
raise ValueError("question_ids must be [B,Q] and match input_ids")
prefix_mask = attention_mask
if prefix_mask is None:
prefix_mask = torch.ones_like(input_ids, dtype=torch.long)
with torch.no_grad():
prefix_embeddings = self.vl.get_input_embeddings()(input_ids)
question_embeddings = self.vl.get_input_embeddings()(question_ids)
image_parts = self.vl.get_image_features(
pixel_values, image_grid_thw
)
image_state = torch.cat(tuple(image_parts), dim=0).to(
device=prefix_embeddings.device,
dtype=prefix_embeddings.dtype,
)
merge = int(self.config.vision_config.spatial_merge_size)
merged_grid = image_grid_thw.to(device=image_state.device).clone()
if bool((merged_grid[:, 1:] % merge).ne(0).any()):
raise ValueError(
"image grid height/width must be divisible by spatial_merge_size"
)
merged_grid[:, 1:] //= merge
image_to_batch = self._spatial_image_to_batch(
input_ids,
prefix_mask,
merged_grid.shape[0],
).to(image_state.device)
refined, telemetry = self.spatial_visual_cell(
image_state,
merged_grid,
image_to_batch,
question_embeddings,
question_attention_mask,
steps=self._spatial_horizon(),
)
self.last_spatial_visual_telemetry = telemetry
image_mask, _ = self.vl.get_placeholder_mask(
input_ids,
inputs_embeds=prefix_embeddings,
image_features=refined,
)
prefix_embeddings = prefix_embeddings.masked_scatter(
image_mask, refined.to(prefix_embeddings.dtype)
)
return prefix_embeddings, prefix_mask
def _forward_spatial_visual_field(
self,
input_ids: torch.LongTensor,
attention_mask: torch.Tensor | None,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
question_attention_mask: torch.Tensor | None,
labels: torch.LongTensor | None,
answer_ids: torch.LongTensor | None,
*,
use_cache: bool | None,
) -> CloseCausalLMOutput:
"""Teacher-forced SRVF pass through the unmodified full decoder."""
prefix_embeddings, prefix_mask = self._spatial_visual_prefix_embeddings(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
)
prefix_positions, rope_deltas = self.vl.get_rope_index(
input_ids,
image_grid_thw=image_grid_thw,
attention_mask=prefix_mask,
)
prefix_length = input_ids.shape[1]
batch_size = input_ids.shape[0]
if not bool(prefix_mask.gt(0).any(dim=-1).all()):
raise ValueError("every sample needs a non-empty multimodal prefix")
full_embeddings = prefix_embeddings
full_positions = prefix_positions
full_mask = prefix_mask
answer_length = 0
if answer_ids is not None:
if answer_ids.ndim != 2 or answer_ids.shape[0] != batch_size:
raise ValueError("answer_ids must be [B,A] and match input_ids")
answer_length = answer_ids.shape[1]
with torch.no_grad():
answer_embeddings = self.vl.get_input_embeddings()(answer_ids)
if labels is not None:
if labels.shape != answer_ids.shape:
raise ValueError("labels and answer_ids must have the same shape")
answer_mask = labels.ne(-100).to(prefix_mask.dtype)
else:
answer_mask = torch.ones_like(answer_ids, dtype=prefix_mask.dtype)
physical = torch.arange(
prefix_length,
prefix_length + answer_length,
device=input_ids.device,
)[None, :]
answer_positions_1d = physical + rope_deltas.to(physical.device)
answer_positions = answer_positions_1d.unsqueeze(0).expand(
3, -1, -1
)
full_embeddings = torch.cat(
[prefix_embeddings, answer_embeddings], dim=1
)
full_positions = torch.cat(
[prefix_positions, answer_positions], dim=-1
)
full_mask = torch.cat([prefix_mask, answer_mask], dim=-1)
elif labels is not None:
raise ValueError("labels require teacher-forced answer_ids")
outputs = self.text_model(
input_ids=None,
inputs_embeds=full_embeddings,
position_ids=full_positions,
attention_mask=full_mask,
use_cache=bool(use_cache),
output_attentions=False,
output_hidden_states=False,
return_dict=True,
)
last_prefix = (prefix_mask.gt(0).sum(dim=-1).long() - 1).clamp_min(0)
seed_hidden = outputs.last_hidden_state[
torch.arange(batch_size, device=input_ids.device), last_prefix
].unsqueeze(1)
hidden_for_logits = seed_hidden
if answer_length > 1:
hidden_for_logits = torch.cat(
[
seed_hidden,
outputs.last_hidden_state[
:, prefix_length : prefix_length + answer_length - 1
],
],
dim=1,
)
logits = self.backbone.lm_head(hidden_for_logits)
loss = None
if labels is not None:
loss = answer_cross_entropy(
logits,
labels,
per_example=bool(
getattr(self.config, "per_example_answer_loss", False)
),
)
return CloseCausalLMOutput(
loss=loss,
logits=logits,
past_key_values=(outputs.past_key_values if use_cache else None),
workspace_states=(),
)
@torch.no_grad()
def _generate_spatial_visual_field(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
attention_mask: torch.Tensor | None,
question_attention_mask: torch.Tensor | None,
*,
max_new_tokens: int,
eos_token_id: int | None,
) -> torch.LongTensor:
"""Greedy decode from a cache built only after V_T is complete."""
if max_new_tokens < 1:
raise ValueError("max_new_tokens must be >= 1")
prefix_embeddings, prefix_mask = self._spatial_visual_prefix_embeddings(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
)
prefix = self.vl(
input_ids=input_ids,
inputs_embeds=prefix_embeddings,
image_grid_thw=image_grid_thw,
attention_mask=prefix_mask,
use_cache=True,
output_attentions=False,
output_hidden_states=False,
return_dict=True,
)
cache = prefix.past_key_values
rope_deltas = prefix.rope_deltas
batch_size, prefix_length = input_ids.shape
last_prefix = (prefix_mask.gt(0).sum(dim=-1).long() - 1).clamp_min(0)
seed = prefix.last_hidden_state[
torch.arange(batch_size, device=input_ids.device), last_prefix
].unsqueeze(1)
next_token = self.backbone.lm_head(seed).argmax(dim=-1)
generated = [next_token]
done = torch.zeros(
batch_size, dtype=torch.bool, device=input_ids.device
)
if eos_token_id is not None:
done |= next_token.squeeze(1).eq(eos_token_id)
for step in range(1, max_new_tokens):
if bool(done.all()):
break
with torch.no_grad():
token_embeddings = self.vl.get_input_embeddings()(next_token)
cache_position = torch.tensor(
[prefix_length + step - 1], device=input_ids.device
)
one_d = cache_position.view(1, 1).expand(batch_size, 1)
one_d = one_d + rope_deltas.to(one_d.device)
position_ids = one_d.unsqueeze(0).expand(3, -1, -1)
generated_mask = torch.ones(
batch_size,
step,
dtype=prefix_mask.dtype,
device=input_ids.device,
)
decode_mask = torch.cat([prefix_mask, generated_mask], dim=-1)
decoded = self.text_model(
input_ids=None,
inputs_embeds=token_embeddings,
position_ids=position_ids,
attention_mask=decode_mask,
past_key_values=cache,
cache_position=cache_position,
use_cache=True,
output_attentions=False,
output_hidden_states=False,
return_dict=True,
)
cache = decoded.past_key_values
next_token = self.backbone.lm_head(
decoded.last_hidden_state
).argmax(dim=-1)
if eos_token_id is not None:
next_token = torch.where(
done.unsqueeze(1),
next_token.new_full(next_token.shape, eos_token_id),
next_token,
)
done |= next_token.squeeze(1).eq(eos_token_id)
generated.append(next_token)
return torch.cat(generated, dim=1)
# -- §3.2 pipeline -----------------------------------------------------
@contextmanager
def _recurrence_adapter_execution(self, enabled: bool):
"""Temporarily toggle recurrence-only LoRA execution.
PEFT's public ``enable_adapters`` helper also changes
``requires_grad``. That is appropriate for persistent train/eval
configuration, but not for routing one shared L21 through a frozen
base call and a trainable recurrent call in the same autograd graph.
The tuner layer's execution flag is therefore saved and restored
directly while its parameters remain trainable throughout.
"""
if not getattr(self.config, "recurrence_only_adapter", False):
yield
return
from peft.tuners.tuners_utils import BaseTunerLayer
modules = [
module
for module in self.backbone.modules()
if isinstance(module, BaseTunerLayer)
]
previous = [module.disable_adapters for module in modules]
for module in modules:
module._disable_adapters = not enabled
try:
yield
finally:
for module, was_disabled in zip(modules, previous):
module._disable_adapters = was_disabled
@torch.no_grad()
def _set_adapters_enabled(self, enabled: bool) -> None:
"""post-l* LoRA on/off (iface teacher는 frozen-base 동작점 필요)."""
from peft.tuners.tuners_utils import BaseTunerLayer
for m in self.backbone.modules():
if isinstance(m, BaseTunerLayer):
m.enable_adapters(enabled)
def _encode_visual_memory(
self,
input_ids: torch.LongTensor, # [B, L_mm] multimodal sequence
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
attention_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""``V* = Pi_img(F^mm_{<=l*}(I, q))``, right-padded. ``[B, N_v, d]``.
Runs under ``no_grad``: the vision encoder and lower backbone are frozen,
and nothing above needs gradients through them.
"""
embeds, pos = embed_multimodal(
self.vl, input_ids, pixel_values, image_grid_thw, attention_mask
)
ctx = make_split_context(self.text_model, embeds, pos, attention_mask)
lower = run_layer_range(self.text_model, ctx, 0, self.config.ell_star + 1)
crr = bool(
getattr(
self.config, "counterfactual_residual_recurrence", False
)
)
capture_functional = bool(
self.training
and getattr(self.config, "functional_state_loss", False)
)
if crr or capture_functional:
# Free target from the same lower pass that produced V*: normal
# multimodal question states at the measured read boundary. The
# oracle shows that these rows alone retain base V* competence in a
# Q-only upper context. Store detached, padded rows for the
# training-only split-state restoration objective.
text_mask = ~vision_span_mask(input_ids, self.config)
if attention_mask is not None:
text_mask &= attention_mask > 0
q_mm, q_mm_pad = select_tokens_padded(lower, text_mask)
q_mm_ids, _ = select_tokens_padded(input_ids.unsqueeze(-1), text_mask)
q_mm = q_mm.detach()
q_mm_ids = q_mm_ids.squeeze(-1)
if crr:
# Consume-once factual boundary state. It is the sole image
# input to CRR; raw visual rows are not materialized below.
self._counterfactual_q_state = q_mm
self._counterfactual_q_pad = q_mm_pad
self._counterfactual_q_ids = q_mm_ids
if getattr(
self.config, "visual_counterfactual_recurrence", False
):
visual_mask = vision_span_mask(input_ids, self.config)
if attention_mask is not None:
visual_mask &= attention_mask > 0
# Consume-once native multimodal scaffold. The lower
# state/context never reaches the answer decoder; it is
# used only by L(ell_star+1)'s visual-on/off recurrent
# evaluations and cleared immediately afterwards.
self._visual_counterfactual_bundle = (
lower.detach(),
ctx,
text_mask,
visual_mask,
)
if capture_functional:
self._functional_q_target = q_mm
self._functional_q_target_pad = q_mm_pad
self._functional_q_target_ids = q_mm_ids
if self.training and getattr(self.config, "interface_loss", False):
# iface teacher: frozen-base 정상(mm) 경로의 q_end hidden @ l*+1.
# adapter를 꺼 base 동작점을 재현 — 학습 target일 뿐 student
# forward에 입력되지 않으므로 새 causal bypass 없음 (advisor).
self._set_adapters_enabled(False)
try:
h1 = run_layer_range(
self.text_model, ctx, self.config.ell_star + 1,
self.config.ell_star + 2, hidden_states=lower)
finally:
self._set_adapters_enabled(True)
if attention_mask is not None:
idx = attention_mask.sum(1).long() - 1 # [B] 마지막 non-pad
else:
idx = torch.full(
(lower.shape[0],), lower.shape[1] - 1,
dtype=torch.long, device=lower.device)
self._iface_h_base = h1[
torch.arange(lower.shape[0], device=lower.device), idx
].detach() # [B, d]
if crr:
# The multimodal lower pass is needed for Q_mm, but CRR has no V*
# memory. Returning an empty tensor prevents an accidental raw
# visual path while preserving the historical helper signature.
return (
lower.new_empty(lower.shape[0], 0, lower.shape[-1]),
torch.empty(
lower.shape[0], 0, dtype=torch.bool, device=lower.device
),
)
mask = image_token_mask(input_ids, self.config.image_token_id)
return select_tokens_padded(lower, mask) # [B, N_v, d], [B, N_v]
@torch.no_grad()
def _encode_hierarchical_visual_memories(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
attention_mask: torch.Tensor | None,
) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
"""Extract one frozen visual memory per visual-read step in one pass.
Layer numbers are 0-indexed outputs: requesting L6 returns the hidden
state after decoder layer 6. The multimodal branch stops after the
deepest requested layer because no later multimodal state is consumed
by this method; the independent text-only branch still reaches
``ell_star`` and supplies the decoder scaffold.
"""
layers = tuple(
int(layer)
for layer in getattr(
self.config, "hierarchical_visual_layers", ()
)
)
if not layers:
raise RuntimeError(
"hierarchical visual extraction requires configured layers"
)
embeds, pos = embed_multimodal(
self.vl, input_ids, pixel_values, image_grid_thw, attention_mask
)
ctx = make_split_context(self.text_model, embeds, pos, attention_mask)
visual_mask = image_token_mask(input_ids, self.config.image_token_id)
hidden = embeds
start = 0
memories: list[tuple[torch.Tensor, torch.Tensor]] = []
reference_pad = None
for layer in layers:
hidden = run_layer_range(
self.text_model,
ctx,
start,
layer + 1,
hidden_states=hidden,
)
memory, padding_mask = select_tokens_padded(hidden, visual_mask)
if reference_pad is None:
reference_pad = padding_mask
elif not torch.equal(reference_pad, padding_mask):
raise RuntimeError(
"visual-token padding changed across frozen layer reads"
)
memories.append((memory, padding_mask))
start = layer + 1
return tuple(memories)
def _prefill_text_branch(
self,
question_ids: torch.LongTensor, # [B, N_q] no image tokens
question_mask: torch.Tensor | None,
):
"""Run the text-only lower branch through ``F_{<=l*}``, caching KV.
Returns ``(q_star, ctx, cache)`` where ``q_star`` is ``[B, N_q, d]`` and
``ctx.hidden_states`` holds the lower-branch state. The raw-evidence
path uses exactly the question rows; historical slot paths append their
fixed cache-shape placeholders.
"""
cfg = self.config
b, n_q = question_ids.shape
embeds = self.vl.get_input_embeddings()(question_ids) # [B, N_q, d]
if question_mask is None:
question_mask = torch.ones(b, n_q, dtype=torch.long, device=embeds.device)
if self._uses_question_state_interface():
# No placeholders and no duplicated Q prefix. Logical text RoPE
# positions remain invariant to right-padding in a batch.
keep = question_mask > 0
one_d = keep.long().cumsum(dim=-1) - 1
one_d = one_d.masked_fill(~keep, 0)
position_ids = one_d.unsqueeze(0).expand(3, -1, -1)
cache = DynamicCache(config=self.text_model.config)
ctx = make_split_context(
self.text_model,
embeds,
position_ids,
question_mask,
past_key_values=cache,
)
lower = run_layer_range(
self.text_model, ctx, 0, cfg.ell_star + 1, use_cache=True
)
ctx.hidden_states = lower
return lower, ctx, cache
s = cfg.num_workspace_slots
slots = self.slot_embedding.unsqueeze(0).expand(b, -1, -1).to(embeds.dtype)
embeds = torch.cat([embeds, slots], dim=1) # [B, N_q + S, d]
# Physical rows are [right-padded q; slots], but each item's slots get
# logical RoPE positions immediately after its own real question.
full_mask, position_ids, _ = _replacement_prefix_layout(question_mask, s)
cache = DynamicCache(config=self.text_model.config)
ctx = make_split_context(
self.text_model, embeds, position_ids, full_mask, past_key_values=cache
)
lower = run_layer_range(
self.text_model, ctx, 0, cfg.ell_star + 1, use_cache=True
) # [B, N_q + S, d]
ctx.hidden_states = lower
return lower[:, :n_q], ctx, cache
def _decode_token_block(
self,
embeds: torch.Tensor, # [B,N_a,d]
position_ids: torch.LongTensor, # [3,B,N_a]
lower_attention_mask: torch.Tensor, # [B,N_q+S+N_a], slots masked
upper_attention_mask: torch.Tensor, # [B,N_q+S+N_a], slots visible
cache: Cache,
cache_position: torch.LongTensor, # [N_a], physical cache positions
return_lower: bool = False,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
"""Decode answer rows while preserving the strict split interface.
Lower layers may attend only to the text question and preceding answer
tokens. The physical placeholder slot rows exist in their cache solely
so every layer has the same cache length, and are masked here. Upper
layers see those same positions as the overwritten ``R_agg`` rows.
"""
tm = self.text_model
lower_ctx = make_split_context(
tm,
embeds,
position_ids,
lower_attention_mask,
past_key_values=cache,
cache_position=cache_position,
)
upper_start = self._upper_decoder_start()
hidden = run_layer_range(
tm, lower_ctx, 0, upper_start, use_cache=True
)
upper_ctx = make_split_context(
tm,
embeds,
position_ids,
upper_attention_mask,
past_key_values=cache,
cache_position=cache_position,
)
upper = run_layer_range(
tm,
upper_ctx,
upper_start,
None,
hidden_states=hidden,
use_cache=True,
)
if return_lower:
return upper, hidden
return upper
def _competitive_read(self, q_in: torch.Tensor, kv_in: torch.Tensor) -> torch.Tensor:
"""ER v2.1 (advisor 2026-08-11): competitive normalization evidence read.
기존 ``evidence_read`` MHA의 W_qkv/W_o를 그대로 사용하되, 표준
row-softmax(evidence 축) 대신 ① reasoning-slot 축 softmax — 각 E_j를
두고 R slot들이 경쟁 (Slot Attention식) — ② evidence 축 재정규화.
새 파라미터/모듈/loss 없음. vanilla read에서는 전 slot이 같은 E_j를
100% 읽는 해가 무벌점이라 실제로 그 해로 붕괴했다 (smoke A/B §5).
q_in: [B, S, d] = LN(R), kv_in: [B, N, d] = LN(E) — norm은 호출부에서.
"""
mha = self.evidence_read
d, h = mha.embed_dim, mha.num_heads
dh = d // h
w, b = mha.in_proj_weight, mha.in_proj_bias
q = nn.functional.linear(q_in, w[:d], None if b is None else b[:d])
k = nn.functional.linear(kv_in, w[d:2 * d], None if b is None else b[d:2 * d])
v = nn.functional.linear(kv_in, w[2 * d:], None if b is None else b[2 * d:])
B, S, _ = q.shape
N = k.shape[1]
q = q.view(B, S, h, dh).transpose(1, 2) # [B, H, S, dh]
k = k.view(B, N, h, dh).transpose(1, 2) # [B, H, N, dh]
v = v.view(B, N, h, dh).transpose(1, 2) # [B, H, N, dh]
logits = q @ k.transpose(-1, -2) / math.sqrt(dh) # [B, H, S, N]
# 확정 구현 (advisor): log-domain competitive normalization.
# A = softmax_j(logsoftmax_s(L)) ≡ C/Σ_j C (수학적 동일; FP64 편차 6e-17)
# 이되, 확률 공간의 나눗셈이 없어 전패 slot에서도 forward/backward 유계
# (극단 경쟁 unit test: exact=NaN, logdom max|grad| 0.68, 행합 1 유지).
# 안정화 이력: ①분모-ε → 1/ε=1e6 증폭, grad 1e10 폭발 (x1, step 2.4k)
# ②ε-분자 → ∂A/∂à = 1/Σ(Ã+ε) ≈ 1/(Nε) 증폭, 재폭발 (step 3.9k)
# ③분모 clamp → 유계지만 전패 slot 읽기 차단 (의미론 변경이라 기각)
# ④log-domain — ε/clamp/hyperparameter 없음. 경쟁 의미론 정확 보존.
log_comp = nn.functional.log_softmax(logits.float(), dim=-2) # slot 축 경쟁
a = torch.softmax(log_comp, dim=-1) # evidence 축 재정규화 (log 공간)
a = a.to(v.dtype)
if getattr(self, "_er_attn_capture", None) is not None:
self._er_attn_capture.append(a.detach().float().cpu()) # [B, H, S, N]
out = (a @ v).transpose(1, 2).reshape(B, S, d) # [B, S, d]
return mha.out_proj(out)
def _encode_workspace_er(
self,
e0: torch.Tensor, # [B, S, d_w] E^(0) = InitialRead(V*)
q_star: torch.Tensor,
v_star: torch.Tensor,
q_padding_mask: torch.Tensor | None,
v_padding_mask: torch.Tensor | None,
gate_overrides: list | None,
v_star_reread: torch.Tensor | None,
v_reread_padding_mask: torch.Tensor | None,
reread_zero_memory: bool,
reread_memory_overrides: list | None,
mid_step: int,
) -> list[torch.Tensor]:
"""Evidence–Reasoning 분리 — strict 버전 (advisor 2026-08-06 수정).
decoder 입력은 **R^(T) 단독** — E 직행 우회(E→decoder)를 구조적으로
제거해 모든 시각 정보가 reasoning을 통과하게 강제한다:
I → V* → E → R^(1:T) → answer.
recurrence 입력은 detach 없는 R + E (answer CE가 f_theta를 경유해
E·reader까지 학습). gradient routing 표(advisor)는 쌍둥이-그래프로
구현한다: traj가 소비할 R-체인만 sg[E]로 병렬 전개 —
값은 비트동일(dropout=0 필수), 그래프만 분리되어
- L_ans → 전부 (reader/E/transition/adapter) [answer 체인]
- L_anchor(k=0) → initial reader·evidence proj만 [E^(0) 그래프]
- L_reason(k>=1) → transition·d_psi만 [traj 체인, sg[E]]
이 되고 trainer/optimizer 변경이 필요 없다.
반환 states = [E^(0), R_t^(1..T)] (traj 체인). answer 체인의 R^(T)는
self._er_answer_final(consume-once)로 prefill에 전달. E의 detached
사본은 self._er_last_evidence에 남겨 probe가 쓴다.
"""
assert mid_step > 0, "ER 모드는 mid_read_step > 0 필요"
assert float(getattr(self.config, "workspace_dropout", 0.0)) == 0.0, (
"ER 쌍둥이-체인은 dropout=0 필요 (두 체인의 값 동일성 전제)")
self._er_answer_final = None # stale-state 방지: forward 끝에서만 설정
E = e0
v2 = bool(getattr(self.config, "explicit_evidence_read", False))
if getattr(self, "_debug_er_capture", False):
self._dbg_R_ans, self._dbg_R_traj = [], []
if getattr(self.config, "r_init_question", False):
# ER v2 canonical (§4): R^0 = R_init (learnable slots, batch broadcast).
# Q* 조건화 블록/Pool 없음 — Q*는 매 transition에서 입력되므로 불필요.
# traj 체인은 stopgrad(R_init) 값 공유 (§9.3: traj -X-> reasoning_slots).
r0 = self.reasoning_slots.unsqueeze(0).expand(E.shape[0], -1, -1)
R, R_t = r0, r0.detach()
else:
R = torch.zeros_like(E) # answer 체인 (E gradient 통과)
R_t = torch.zeros_like(E) # traj 체인 (sg[E] — evidence/reader 차단)
states = [E]
for k in range(self.config.num_workspace_steps):
if v2:
# ER v2 canonical (§5-6): R̄ = R + CrossAttn(LN(R), LN(E), LN(E));
# R = f_theta(R̄, Q*). transition은 잔차 블록이므로 R̄를 외부에서
# 재가산하지 않는다 (double-residual 금지, §6).
# traj 체인: evidence_read(R_t, sg(E)) — E 값은 읽되 E 생성부 차단 (§9.3).
def _read_E(r, mem):
q = self.evidence_read_norm_q(r)
kv = self.evidence_read_norm_kv(mem)
if getattr(self.config, "competitive_evidence_read", False):
return r + self._competitive_read(q, kv)
out, _ = self.evidence_read(q, kv, kv, need_weights=False)
return r + out
R = self.workspace_transition(_read_E(R, E), q_star, q_padding_mask)
if getattr(self.config, "er_single_chain", False):
# ER v2.2 C (§12): k>=1 traj 감독이 없으므로 twin 불필요 —
# states가 answer 체인을 직접 담아 L_ans가 R1..RT 전체에
# (aggregation 경유) gradient를 보낸다 (§11).
R_t = R
else:
R_t = self.workspace_transition(
_read_E(R_t, E.detach()), q_star, q_padding_mask)
if getattr(self, "_debug_er_capture", False):
self._dbg_R_ans.append(R.detach().clone())
self._dbg_R_traj.append(R_t.detach().clone())
else:
x = R + E
total = self.workspace_transition(x, q_star, q_padding_mask)
R = total - E
x_t = R_t + E.detach()
total_t = self.workspace_transition(x_t, q_star, q_padding_mask)
R_t = total_t - E.detach()
if (k + 1) == mid_step:
ov = gate_overrides[k] if gate_overrides is not None else None
off_mid = ov == 0.0
if not off_mid:
v_r, vp_r = v_star, v_padding_mask
if v_star_reread is not None:
v_r, vp_r = v_star_reread, v_reread_padding_mask
if reread_zero_memory:
v_r = torch.zeros_like(v_r)
elif (reread_memory_overrides is not None
and reread_memory_overrides[k] is not None):
v_r, vp_r = reread_memory_overrides[k]
# ER v2 (advisor): 재읽기 query는 R-only — "무엇을 다시 볼지"는
# reasoning state가 정하고, E는 저장소 역할만. (v1 경로는 E+R 유지.)
u_q = R if v2 else (E + R)
woa = self.workspace_read.visual_residual(u_q, v_r, vp_r)
a_eff = float(getattr(self.config, "mid_read_alpha", 1.0))
if ov is not None:
a_eff = float(ov) # eval sweep은 절대값으로 대체
E_new = E + a_eff * woa
with torch.no_grad():
uf = (E + R).float()
dv = (a_eff * woa).float()
den = uf.norm(dim=(-2, -1)).clamp_min(1e-6)
self.last_mid_stats_full = {
"u_mid": dv.norm(dim=(-2, -1)) / den,
"a_mid": nn.functional.cosine_similarity(
uf.flatten(1), dv.flatten(1), dim=1),
"r_mid": 1 - nn.functional.cosine_similarity(
uf.flatten(1), (E_new + R).float().flatten(1), dim=1),
"q_mid": (E_new + R).float().norm(dim=(-2, -1)) / den,
}
if E.shape[0] > 1:
def _div(xx):
pm = nn.functional.normalize(
xx.float().mean(1), dim=-1)
g = pm @ pm.T
b_ = g.shape[0]
off = g.masked_select(~torch.eye(
b_, dtype=torch.bool, device=g.device))
return 1 - off.mean()
self.last_mid_stats_full["div_z0"] = _div(
states[0]).expand(E.shape[0]).clone()
self.last_mid_stats_full["div_z1"] = _div(
states[1] if len(states) > 1 else states[0]
).expand(E.shape[0]).clone()
self.last_mid_stats_full["div_umid"] = _div(
E + R).expand(E.shape[0]).clone()
E = E_new
states.append(R_t)
# D_{k*->j}: read를 끈 shadow R-chain(E = E^(0) 유지) 대비 발산 —
# read 갱신이 최종 reasoning까지 남는지의 학습-중 감시.
with torch.no_grad():
Rs = states[mid_step - 1] if mid_step > 1 else torch.zeros_like(E)
prop = []
for j in range(mid_step, len(states)):
xs = Rs + e0.detach()
Rs = self.workspace_transition(xs, q_star, q_padding_mask) - e0.detach()
prop.append(1 - nn.functional.cosine_similarity(
states[j].float().flatten(1), Rs.float().flatten(1), dim=1))
self.last_mid_prop_full = torch.stack(prop, dim=1)
self.last_reread_gates = []
self.last_reread_update_mags = []
if getattr(self.config, "concat_aggregation", False):
# ER v2.2 C 확정판 (§6/§10): R_agg = W_agg · Concat(R1..RT).
# R0/E/V* 미포함 (§10); states는 single-chain이라 answer 체인 =
# grad 경로 → L_ans가 aggregator를 거쳐 R1..RT 각각에 직접 도달 (§14).
r_steps = states[1:] # R1..RT, 각 [B, S, d_w]
r_cat = torch.cat(r_steps, dim=-1) # [B, S, T*d_w]
R_final = self.reasoning_aggregator(r_cat) # [B, S, d_w]
with torch.no_grad():
# §18: step별 block norm ||W_k|| + cos(R_agg, R_k)
d_w = R.shape[-1]
w = self.reasoning_aggregator.weight.float() # [D, T*D]
self._er_agg_stats = {
**{f"aggW{_i + 1}": float(w[:, _i * d_w:(_i + 1) * d_w].norm())
for _i in range(len(r_steps))},
"agg_norm": float(R_final.float().norm(dim=(-2, -1)).mean()),
**{f"agg_cos_R{_i + 1}": float(nn.functional.cosine_similarity(
r_steps[_i].float().flatten(1),
R_final.float().flatten(1), dim=1).mean())
for _i in range(len(r_steps))},
}
self._er_answer_final = R_final
elif getattr(self.config, "temporal_aggregation", False):
# (legacy — softmax 변형)
a_w = torch.softmax(self.temporal_logits.float(), dim=0).to(R.dtype) # [T]
r_steps = states[1:] # R1..RT
R_final = r_steps[0] * a_w[0]
for _i in range(1, len(r_steps)):
R_final = R_final + a_w[_i] * r_steps[_i]
self._er_answer_final = R_final
else:
self._er_answer_final = R
self._er_last_evidence = E.detach()
return states
def _encode_counterfactual_residual_recurrence(
self,
q_text: torch.Tensor,
q_mm: torch.Tensor,
q_padding_mask: torch.Tensor | None,
ctx,
visual_bundle=None,
component_ablation_mode: str | None = None,
) -> list[torch.Tensor]:
"""Run the native-layer counterfactual residual recurrence.
``F=L(ell_star+1)`` is called as the one shared recurrent operator. Its
text-only call populates F's prefix cache, so answer tokens may traverse
F without receiving an image-conditioned KV bypass. Every other F call
uses a cache-free context and therefore changes only the returned
question-shaped state.
Returned states are ``[B, C1, ..., CT]``:
``B=F(Q_text)``, ``C1=F(Q_mm)-B``,
``G_B(C)=F(B+C)-F(B)``, and
``Ck=(1-beta)C{k-1}+beta*G_B(C{k-1})``.
"""
if getattr(
self.config, "visual_counterfactual_recurrence", False
):
if visual_bundle is None:
raise RuntimeError(
"visual counterfactual recurrence is missing its multimodal "
"boundary scaffold"
)
return self._encode_visual_counterfactual_recurrence(
q_text,
q_padding_mask,
ctx,
visual_bundle,
component_ablation_mode=component_ablation_mode,
)
if component_ablation_mode is not None:
raise ValueError(
"component ablations require visual counterfactual recurrence"
)
steps = int(self.config.num_workspace_steps)
beta = float(self.config.counterfactual_beta)
cell_index = int(self.config.ell_star) + 1
tm = self.text_model
# This is the only call that writes L21 prefix KV. Its input is
# strictly text-only, preserving the no-image lower-cache invariant.
base_anchor = run_layer_range(
tm,
ctx,
cell_index,
cell_index + 1,
hidden_states=q_text,
use_cache=True,
)
# Recurrent evaluations must neither read nor mutate the replacement
# cache. The causal mask and RoPE tables remain the same Q-only ones.
recurrent_ctx = replace(ctx, past_key_values=None)
def cell(hidden: torch.Tensor) -> torch.Tensor:
return run_layer_range(
tm,
recurrent_ctx,
cell_index,
cell_index + 1,
hidden_states=hidden,
use_cache=False,
)
factual_t1 = cell(q_mm)
residual = factual_t1 - base_anchor
if q_padding_mask is not None:
residual = residual.masked_fill(q_padding_mask.unsqueeze(-1), 0)
states = [base_anchor, residual]
if steps > 1:
fixed_anchor_next = cell(base_anchor)
for _ in range(1, steps):
proposed = cell(base_anchor + residual) - fixed_anchor_next
if beta == 1.0:
residual = proposed
elif beta == 0.0:
residual = residual
else:
residual = torch.lerp(residual, proposed, beta)
if hasattr(self, "crr_policy_adapter"):
residual = self.crr_policy_adapter(residual)
if q_padding_mask is not None:
residual = residual.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
states.append(residual)
return states
def _encode_visual_counterfactual_recurrence(
self,
q_text: torch.Tensor,
q_padding_mask: torch.Tensor | None,
text_ctx,
visual_bundle,
*,
component_ablation_mode: str | None = None,
) -> list[torch.Tensor]:
"""Run Causal Visual Residual Recurrence with one native layer.
The first state is the ordinary multimodal output of
``F=L(ell_star+1)``. For every later step, the exact same physical
multimodal scaffold is evaluated twice from the current question
state: once normally and once with only text-query -> visual-key/value
edges blocked. Their difference is therefore the update attributable
to a fresh visual read through ``F``:
``R1 = Pi_Q F(H_mm*)``
Historical CVRR uses the intervention difference
``Delta_k^V = Pi_Q(F_on([E;R{k-1}]) - F_off([E;R{k-1}]))``.
Full-state CVRR instead keeps the complete native computation:
``R_tilde_k = Pi_Q F_on([E;R{k-1}])`` and
``R_k = R_{k-1} + beta (R_tilde_k - R_{k-1})``. The latter removes
the visual-off branch from train/inference; it remains an evaluation
intervention only. The visual-preserving adapter-correction subtype
instead uses
``Delta_k=F_{base+LoRA,on}([E;R{k-1}])-F_{base,off}([E;R{k-1}])``:
this cancels frozen non-visual drift without cancelling the pretrained
image-caused transition.
``E`` is the persistent visual portion of the first native output.
It is never cached for, concatenated into, or exposed to L22+; only the
selected question-shaped recurrent state reaches the answer decoder
through CRR's strict interface.
"""
steps = int(self.config.num_workspace_steps)
beta = float(self.config.counterfactual_beta)
cumulative = bool(
getattr(self.config, "visual_cumulative_recurrence", False)
)
full_state = bool(
getattr(self.config, "visual_full_state_recurrence", False)
)
disable_visual_reread = (
component_ablation_mode == "without_visual_reread"
)
if disable_visual_reread and not full_state:
raise ValueError(
"without_visual_reread is defined for full-state CVRR only"
)
boundary_reentry = bool(
getattr(self.config, "visual_boundary_reentry", False)
)
source_centered_reentry = bool(
getattr(
self.config,
"visual_source_centered_reentry",
False,
)
)
recurrence_only_adapter = bool(
getattr(self.config, "recurrence_only_adapter", False)
)
visual_preserving_correction = bool(
getattr(
self.config,
"visual_preserving_adapter_correction",
False,
)
)
cell_index = int(self.config.ell_star) + 1
tm = self.text_model
mm_lower, mm_ctx, text_rows, visual_rows = visual_bundle
# Populate L21's answer-time prefix cache with text-only states. This
# is the sole cached call at the recurrent layer.
with self._recurrence_adapter_execution(False):
if recurrence_only_adapter:
with torch.no_grad():
base_anchor = run_layer_range(
tm,
text_ctx,
cell_index,
cell_index + 1,
hidden_states=q_text,
use_cache=True,
)
else:
base_anchor = run_layer_range(
tm,
text_ctx,
cell_index,
cell_index + 1,
hidden_states=q_text,
use_cache=True,
)
normal_ctx = replace(mm_ctx, past_key_values=None)
def initial_native_read(hidden: torch.Tensor) -> torch.Tensor:
return run_layer_range(
tm,
normal_ctx,
cell_index,
cell_index + 1,
hidden_states=hidden,
use_cache=False,
)
# This is the only CVRR pass that needs outputs for every visual row
# (to construct persistent E). Saving its 5x-width MLP activations is
# the dominant long-image memory cost; non-reentrant checkpointing
# recomputes exactly this one native layer during backward instead.
if recurrence_only_adapter:
# R1 is an exact pretrained multimodal anchor. Adapter gradients
# begin only at the R1->R2 transition.
with self._recurrence_adapter_execution(False), torch.no_grad():
first_full = initial_native_read(mm_lower)
elif self.training and torch.is_grad_enabled():
first_full = checkpoint(
initial_native_read,
mm_lower,
use_reentrant=False,
preserve_rng_state=True,
)
else:
first_full = initial_native_read(mm_lower)
r1, r1_padding = select_tokens_padded(first_full, text_rows)
if q_padding_mask is not None and not torch.equal(
r1_padding, q_padding_mask
):
raise ValueError(
"multimodal and text-only padding disagree in visual recurrence"
)
if component_ablation_mode == "without_entire_visual_path":
# Matched T-step text-only recurrence. The multimodal branch above
# is computed only to preserve the common evaluation setup; none
# of its states enter this returned decoder prefix.
text_recurrent_ctx = replace(text_ctx, past_key_values=None)
r_current = base_anchor
residual = torch.zeros_like(base_anchor)
states = [base_anchor, residual]
for _ in range(1, steps):
with self._recurrence_adapter_execution(True):
proposed = run_layer_range(
tm,
text_recurrent_ctx,
cell_index,
cell_index + 1,
hidden_states=r_current,
use_cache=False,
)
r_current = r_current + beta * (proposed - r_current)
residual = r_current - base_anchor
if q_padding_mask is not None:
residual = residual.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
states.append(residual)
return states
if component_ablation_mode == "without_native_mm_initialization":
r_current = base_anchor
residual = torch.zeros_like(base_anchor)
else:
r_current = r1
residual = r1 - base_anchor
if q_padding_mask is not None:
residual = residual.masked_fill(q_padding_mask.unsqueeze(-1), 0)
states = [base_anchor, residual]
if component_ablation_mode == "without_recurrence":
# Keep the serialized state-list contract while decoding the same
# native R1 at every nominal slot; no recurrent layer is called.
states.extend(residual for _ in range(1, steps))
return states
if component_ablation_mode == "without_visual_rows_keep_h1":
# Retain the exact native multimodal R1, then physically remove the
# visual rows for every later transition. The learned recurrent
# layer, T, beta, and strict decoder stay unchanged; recurrence is
# evaluated on the depth-matched text-only sequence geometry.
text_recurrent_ctx = replace(text_ctx, past_key_values=None)
for _ in range(1, steps):
with self._recurrence_adapter_execution(True):
proposed = run_layer_range(
tm,
text_recurrent_ctx,
cell_index,
cell_index + 1,
hidden_states=r_current,
use_cache=False,
)
r_current = r_current + beta * (proposed - r_current)
residual = r_current - base_anchor
if q_padding_mask is not None:
residual = residual.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
states.append(residual)
return states
update_rms: list[torch.Tensor] = []
update_relative: list[torch.Tensor] = []
boundary_q = None
if boundary_reentry:
boundary_q, boundary_padding = select_tokens_padded(
mm_lower, text_rows
)
if not torch.equal(boundary_padding, r1_padding):
raise RuntimeError(
"post-L20 and post-L21 question row layouts disagree"
)
# The source-centered control separates the fixed text-only anchor
# from the native image-dependent initial state:
# A = F_text(H20), Delta_1 = R1 - A.
# It also evaluates the recurrent cell once at zero deviation so every
# later update can subtract that response. Consequently G(0)=0 even
# after the recurrence-only adapter has learned a nonzero mapping.
source_anchor_q = None
if source_centered_reentry and steps > 1:
def source_anchor(
hidden: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
with self._recurrence_adapter_execution(True):
output, _, padding = self._visual_on_off_native_text_rows(
hidden,
normal_ctx,
text_rows,
visual_rows,
cell_index,
compute_off=False,
)
if output is None:
raise RuntimeError(
"source-centered anchor returned no output"
)
return output, padding
if self.training and torch.is_grad_enabled():
source_anchor_q, source_anchor_padding = checkpoint(
source_anchor,
mm_lower,
use_reentrant=False,
preserve_rng_state=True,
)
else:
source_anchor_q, source_anchor_padding = source_anchor(
mm_lower
)
if not torch.equal(source_anchor_padding, r1_padding):
raise RuntimeError(
"source-centered anchor row layout changed"
)
# Historical full-state recurrence reuses post-L21 visual rows and
# replaces its text rows with the current R. The boundary-reentry
# control instead rebuilds every call from the native post-L20 input
# and carries only R-R1 into those question rows. Both preserve the
# original multimodal ordering, masks, and M-RoPE geometry.
for _ in range(1, steps):
state_before = r_current
if boundary_reentry:
if boundary_q is None:
raise RuntimeError("boundary re-entry state was not initialized")
if source_centered_reentry:
# The initial deviation is the native visual effect
# R1-A, not zero. Re-enter L21 in its native H20
# coordinates while carrying only this evolving residual.
recurrent_q = boundary_q + (r_current - base_anchor)
else:
recurrent_q = boundary_q + (r_current - r1)
recurrent_input = _replace_selected_rows_from_padded(
mm_lower,
text_rows,
recurrent_q,
r1_padding,
)
else:
recurrent_input = _replace_selected_rows_from_padded(
first_full,
text_rows,
r_current,
r1_padding,
)
if visual_preserving_correction:
# The two executions intentionally use different adapter
# routing. Their difference retains the pretrained visual
# effect while cancelling the frozen layer's non-visual drift:
# F_{base+LoRA,on} - F_{base,off}
# = (F_{base,on} - F_{base,off})
# + (F_{base+LoRA,on} - F_{base,on}).
# Checkpoint both branches because retaining two native-layer
# activation sets would otherwise roughly double the long-image
# recurrent memory footprint.
def adapter_on(
hidden: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
with self._recurrence_adapter_execution(True):
output, _, padding = (
self._visual_on_off_native_text_rows(
hidden,
normal_ctx,
text_rows,
visual_rows,
cell_index,
compute_off=False,
)
)
if output is None:
raise RuntimeError("visual-on branch returned no output")
return output, padding
def frozen_visual_off(
hidden: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
with self._recurrence_adapter_execution(False):
_, output, padding = (
self._visual_on_off_native_text_rows(
hidden,
normal_ctx,
text_rows,
visual_rows,
cell_index,
compute_on=False,
)
)
if output is None:
raise RuntimeError("visual-off branch returned no output")
return output, padding
if self.training and torch.is_grad_enabled():
on_q, on_padding = checkpoint(
adapter_on,
recurrent_input,
use_reentrant=False,
preserve_rng_state=True,
)
off_q, off_padding = checkpoint(
frozen_visual_off,
recurrent_input,
use_reentrant=False,
preserve_rng_state=True,
)
else:
on_q, on_padding = adapter_on(recurrent_input)
off_q, off_padding = frozen_visual_off(recurrent_input)
if not torch.equal(on_padding, off_padding):
raise RuntimeError(
"visual-on and frozen visual-off row layouts disagree"
)
else:
with self._recurrence_adapter_execution(
component_ablation_mode != "without_learned_transition"
):
if full_state:
if disable_visual_reread:
_, off_q, on_padding = (
self._visual_on_off_native_text_rows(
recurrent_input,
normal_ctx,
text_rows,
visual_rows,
cell_index,
compute_on=False,
compute_off=True,
)
)
if off_q is None:
raise RuntimeError(
"visual-off recurrent branch returned no output"
)
# Preserve the complete learned recurrent layer and
# MLP while removing only Q->visual K/V edges.
on_q = off_q
else:
on_q, off_q, on_padding = (
self._visual_on_off_native_text_rows(
recurrent_input,
normal_ctx,
text_rows,
visual_rows,
cell_index,
compute_off=False,
)
)
else:
on_q, off_q, on_padding = (
self._visual_on_off_native_text_rows(
recurrent_input,
normal_ctx,
text_rows,
visual_rows,
cell_index,
)
)
if on_q is None:
raise RuntimeError("visual-on branch returned no output")
if visual_preserving_correction:
transition_delta = on_q - off_q
elif full_state:
if source_centered_reentry:
if source_anchor_q is None:
raise RuntimeError(
"source-centered anchor was not initialized"
)
# G(Delta) = F(H20 + Delta) - F(H20), hence G(0)=0.
# The state update below gives
# Delta_k = Delta_{k-1} + beta * G(Delta_{k-1}).
transition_delta = on_q - source_anchor_q
else:
transition_delta = on_q - state_before
else:
if off_q is None:
raise RuntimeError("visual-off branch was not computed")
transition_delta = on_q - off_q
if q_padding_mask is not None:
transition_delta = transition_delta.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
if not torch.equal(on_padding, r1_padding):
raise RuntimeError("recurrent text-row layout changed")
if full_state or cumulative:
r_current = state_before + beta * transition_delta
else:
r_current = r1 + beta * transition_delta
residual = r_current - base_anchor
if q_padding_mask is not None:
residual = residual.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
states.append(residual)
with torch.no_grad():
metric_delta = (
beta * transition_delta if full_state else transition_delta
)
valid = (
torch.ones_like(r1_padding, dtype=torch.bool)
if q_padding_mask is None
else ~q_padding_mask
)
valid_f = valid.unsqueeze(-1).float()
denom = valid_f.sum(dim=(1, 2)).clamp_min(1.0)
delta_rms = (
(metric_delta.float().square() * valid_f).sum(dim=(1, 2))
/ (denom * metric_delta.shape[-1])
).sqrt()
state_rms = (
(state_before.float().square() * valid_f).sum(dim=(1, 2))
/ (denom * state_before.shape[-1])
).sqrt()
update_rms.append(delta_rms)
update_relative.append(delta_rms / state_rms.clamp_min(1e-8))
update_rms_tensor = torch.stack(update_rms, dim=1) if update_rms else None
update_relative_tensor = (
torch.stack(update_relative, dim=1) if update_relative else None
)
if full_state:
self.last_recurrent_update_rms = update_rms_tensor
self.last_recurrent_update_relative = update_relative_tensor
self.last_visual_counterfactual_update_rms = None
self.last_visual_counterfactual_update_relative = None
else:
self.last_visual_counterfactual_update_rms = update_rms_tensor
self.last_visual_counterfactual_update_relative = update_relative_tensor
self.last_recurrent_update_rms = None
self.last_recurrent_update_relative = None
return states
def _visual_on_off_native_text_rows(
self,
full_input: torch.Tensor,
full_ctx,
text_rows: torch.Tensor,
visual_rows: torch.Tensor,
cell_index: int,
compute_on: bool = True,
compute_off: bool = True,
) -> tuple[
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor,
]:
"""Exact text-row execution of one native visual-on/off layer pair.
Decoder-layer attention and MLP are pointwise in their query/output
rows. Consequently, computing Q and the post-attention MLP only for
text rows while computing K/V for the complete multimodal sequence is
algebraically identical to two full layer calls followed by ``Pi_Q``.
Sharing the common Q/K/V projections between the on/off branches also
preserves gradients: both masks consume the same projected tensors and
autograd sums their two uses exactly as it would for tied weights.
This is an execution optimization only. It removes six discarded
visual-row MLP activations per T=4 forward, which otherwise dominate
memory on long 8k-token images.
"""
if not compute_on and not compute_off:
raise ValueError("at least one visual on/off branch must be computed")
layer = self.text_model.layers[cell_index]
attention = layer.self_attn
if getattr(layer, "attention_type", "full_attention") != "full_attention":
raise NotImplementedError(
"selective CVRR execution currently requires full attention"
)
normalized_full = layer.input_layernorm(full_input)
normalized_q, q_padding = select_tokens_padded(
normalized_full, text_rows
)
residual_q, residual_padding = select_tokens_padded(full_input, text_rows)
if not torch.equal(q_padding, residual_padding):
raise RuntimeError("native selective-query padding changed")
batch_size, query_length, _ = normalized_q.shape
full_length = normalized_full.shape[1]
head_dim = attention.head_dim
query = attention.q_proj(normalized_q).view(
batch_size, query_length, -1, head_dim
).transpose(1, 2)
key = attention.k_proj(normalized_full).view(
batch_size, full_length, -1, head_dim
).transpose(1, 2)
value = attention.v_proj(normalized_full).view(
batch_size, full_length, -1, head_dim
).transpose(1, 2)
cos_full, sin_full = full_ctx.position_embeddings
def select_rope_rows(
rope: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
# Qwen2.5 multimodal RoPE keeps the t/h/w axis explicitly as
# [3,B,L,D]; text-only variants may already be [B,L,D].
if rope.ndim == 3:
return select_tokens_padded(rope, text_rows)
if rope.ndim != 4 or rope.shape[0] != 3:
raise ValueError(
f"unexpected M-RoPE tensor shape {tuple(rope.shape)}"
)
selected = []
selected_padding = None
for axis in range(3):
rows, padding = select_tokens_padded(rope[axis], text_rows)
selected.append(rows)
if selected_padding is None:
selected_padding = padding
elif not torch.equal(selected_padding, padding):
raise RuntimeError("M-RoPE axis padding changed")
return torch.stack(selected, dim=0), selected_padding
cos_q, cos_padding = select_rope_rows(cos_full)
sin_q, sin_padding = select_rope_rows(sin_full)
if not torch.equal(cos_padding, q_padding) or not torch.equal(
sin_padding, q_padding
):
raise RuntimeError("native selective-query RoPE padding changed")
mrope_section = attention.rope_scaling["mrope_section"]
query, _ = apply_multimodal_rotary_pos_emb(
query,
query,
cos_q,
sin_q,
mrope_section,
)
_, key = apply_multimodal_rotary_pos_emb(
key,
key,
cos_full,
sin_full,
mrope_section,
)
# Rectangular causal mask at each text row's original physical index.
# Padded gathered rows get one harmless fallback key and are zeroed
# after the layer, avoiding all-masked-row kernel corner cases.
key_indices = torch.arange(full_length, device=full_input.device)
query_indices = torch.zeros(
batch_size,
query_length,
dtype=torch.long,
device=full_input.device,
)
for batch_index in range(batch_size):
physical = text_rows[batch_index].nonzero(as_tuple=False).squeeze(-1)
query_indices[batch_index, : physical.numel()] = physical
visible = key_indices.view(1, 1, full_length) <= query_indices.unsqueeze(-1)
if full_ctx.attention_mask is not None:
if full_ctx.attention_mask.shape != (batch_size, full_length):
raise ValueError(
"selective native execution requires a 2-D [B,L] mask"
)
visible &= full_ctx.attention_mask[:, None, :].bool()
valid_query = ~q_padding
visible &= valid_query.unsqueeze(-1)
for batch_index in range(batch_size):
if bool(q_padding[batch_index].any()):
visible[batch_index, q_padding[batch_index], 0] = True
normal_mask = visible.unsqueeze(1)
visual_off_mask = (
normal_mask & ~visual_rows[:, None, None, :]
if compute_off
else None
)
def finish(mask: torch.Tensor) -> torch.Tensor:
attended, _ = sdpa_attention_forward(
attention,
query,
key,
value,
mask,
dropout=(
attention.attention_dropout if self.training else 0.0
),
scaling=attention.scaling,
)
attended = attention.o_proj(
attended.reshape(batch_size, query_length, -1).contiguous()
)
hidden = residual_q + attended
hidden = hidden + layer.mlp(layer.post_attention_layernorm(hidden))
return hidden.masked_fill(q_padding.unsqueeze(-1), 0)
normal_output = finish(normal_mask) if compute_on else None
off_output = finish(visual_off_mask) if visual_off_mask is not None else None
return normal_output, off_output, q_padding
def _counterfactual_transition_rollout_prefix(
self,
states: list[torch.Tensor] | tuple[torch.Tensor, ...],
donors: torch.LongTensor,
transition_index: int,
question_attention_mask: torch.Tensor | None,
) -> torch.Tensor:
"""Intervene on one CRR transition and roll its causal tail forward.
``states`` is the factual ``[B, C1, ..., CT]`` trajectory. At the
selected 1-indexed transition ``k``, ``Ck`` is replaced by a detached
different-answer donor state. Earlier factual states remain fixed and
later states are recomputed with the target example's anchor and the
same shared native recurrent cell. The returned tensor is the exact
final decoder prefix ``B + Aggregate(C1..CT)`` under that intervention.
Detaching the donor and factual history prevents the negative branch
from making another example adversarial or rewriting pre-intervention
history. Gradients still pass through the counterfactual tail, the
shared aggregator (when present), and the upper decoder replay.
"""
horizon = int(self.config.num_workspace_steps)
if len(states) != horizon + 1:
raise ValueError(
f"expected B plus {horizon} CRR states, got {len(states)}"
)
if not 1 <= int(transition_index) <= horizon:
raise ValueError(
f"reliance_transition_index must be in [1,{horizon}], "
f"got {transition_index}"
)
batch_size, question_length = states[0].shape[:2]
if donors.shape != (batch_size,):
raise ValueError(
f"donors must have shape {(batch_size,)}, got {tuple(donors.shape)}"
)
if question_attention_mask is None:
q_mask = torch.ones(
batch_size,
question_length,
dtype=torch.long,
device=states[0].device,
)
else:
q_mask = question_attention_mask
q_keep = q_mask.gt(0)
donor_keep = q_keep.index_select(0, donors)
aligned = q_keep & donor_keep
k = int(transition_index)
base_anchor = states[0].detach()
residual = states[k].index_select(0, donors).detach()
residual = residual.masked_fill(~aligned.unsqueeze(-1), 0)
hybrid_residuals = [state.detach() for state in states[1:k]]
hybrid_residuals.append(residual)
if k < horizon:
q_pos_1d = q_keep.long().cumsum(dim=-1) - 1
q_pos_1d = q_pos_1d.masked_fill(~q_keep, 0)
q_positions = q_pos_1d.unsqueeze(0).expand(3, -1, -1)
recurrent_ctx = make_split_context(
self.text_model,
base_anchor,
q_positions,
q_mask,
)
cell_index = int(self.config.ell_star) + 1
def cell(hidden: torch.Tensor) -> torch.Tensor:
return run_layer_range(
self.text_model,
recurrent_ctx,
cell_index,
cell_index + 1,
hidden_states=hidden,
use_cache=False,
)
fixed_anchor_next = cell(base_anchor)
beta = float(self.config.counterfactual_beta)
q_padding_mask = ~q_keep
for _ in range(k, horizon):
proposed = cell(base_anchor + residual) - fixed_anchor_next
if beta == 1.0:
residual = proposed
elif beta == 0.0:
residual = residual
else:
residual = torch.lerp(residual, proposed, beta)
if hasattr(self, "crr_policy_adapter"):
residual = self.crr_policy_adapter(residual)
residual = residual.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
hybrid_residuals.append(residual)
if len(hybrid_residuals) != horizon:
raise RuntimeError(
"transition rollout did not reconstruct the full CRR horizon"
)
return base_anchor + self._aggregate_crr_residuals(hybrid_residuals)
def _encode_counterfactual_residual_policy(
self,
q_text: torch.Tensor,
q_mm: torch.Tensor,
q_padding_mask: torch.Tensor | None,
ctx,
*,
policy_std: float,
sample: bool,
sample_mask: torch.Tensor | None = None,
) -> tuple[list[torch.Tensor], LatentPolicyTrace]:
"""Sample ``C2..CT`` as an amortized Gaussian latent policy.
``B`` and ``C1`` remain the deterministic CRR visual anchor. For each
later step, the shared native cell predicts the Gaussian mean
``mu_k = (1-beta) C{k-1} + beta G_B(C{k-1})``.
The action is the next complete residual state
``Ck ~ N(mu_k, (policy_std * RMS(C1))^2 I)``. Sampled states and the
state supplied to the following transition are detached, as required by
score-function policy gradients. This blocks accidental pathwise
gradients while retaining ``d log pi(Ck|C{k-1}) / d theta`` through the
policy mean. Gaussian constants are omitted and the elementwise score
is averaged over valid latent dimensions to keep its scale independent
of question length.
The normal deterministic CRR path is implemented separately above and
remains untouched. ``sample=False`` is a structural test/inference
control: it executes the exact policy mean while still returning a
trace.
"""
if not math.isfinite(policy_std) or policy_std <= 0.0:
raise ValueError("policy_std must be finite and > 0")
steps = int(self.config.num_workspace_steps)
if steps < 2:
raise ValueError("latent policy requires at least two CRR steps")
beta = float(self.config.counterfactual_beta)
cell_index = int(self.config.ell_star) + 1
tm = self.text_model
recurrent_ctx = replace(ctx, past_key_values=None)
def cell(hidden: torch.Tensor) -> torch.Tensor:
return run_layer_range(
tm,
recurrent_ctx,
cell_index,
cell_index + 1,
hidden_states=hidden,
use_cache=False,
)
# C1 defines the policy's local visual region and is deliberately not an
# RL action. It also fills L21's prefix cache from the text-only branch,
# exactly as deterministic CRR does. No graph is retained for B/C1.
with torch.no_grad():
base_anchor = run_layer_range(
tm,
ctx,
cell_index,
cell_index + 1,
hidden_states=q_text,
use_cache=True,
)
residual = cell(q_mm) - base_anchor
if q_padding_mask is not None:
residual = residual.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
base_observation = base_anchor.detach()
residual = residual.detach()
states = [base_observation, residual]
if q_padding_mask is None:
valid = torch.ones(
residual.shape[:2],
dtype=torch.bool,
device=residual.device,
)
else:
valid = ~q_padding_mask
valid_3d = valid.unsqueeze(-1)
n_valid_dims = (
valid.sum(dim=1).to(torch.float32) * residual.shape[-1]
).clamp_min(1.0) # [B]
c1_square_sum = (
residual.float().square() * valid_3d
).sum(dim=(1, 2))
c1_rms = (c1_square_sum / n_valid_dims).sqrt().clamp_min(1e-6)
sigma = (float(policy_std) * c1_rms).view(-1, 1, 1).detach()
# F(B) is shared by every conditional policy mean. B is an observation,
# but the shared cell weights remain differentiable.
fixed_anchor_next = cell(base_observation)
log_probs: list[torch.Tensor] = []
mean_update_rms: list[torch.Tensor] = []
noise_rms: list[torch.Tensor] = []
for _ in range(1, steps):
observation = residual.detach()
proposed = cell(base_observation + observation) - fixed_anchor_next
if beta == 1.0:
policy_mean = proposed
elif beta == 0.0:
policy_mean = observation
else:
policy_mean = torch.lerp(observation, proposed, beta)
if hasattr(self, "crr_policy_adapter"):
policy_mean = self.crr_policy_adapter(policy_mean)
if q_padding_mask is not None:
policy_mean = policy_mean.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
if sample:
noise = torch.randn_like(policy_mean)
if sample_mask is not None:
if sample_mask.shape != (policy_mean.shape[0],):
raise ValueError(
"latent policy sample mask must have shape "
f"{(policy_mean.shape[0],)}, got "
f"{tuple(sample_mask.shape)}"
)
noise = noise * sample_mask.to(
device=noise.device, dtype=noise.dtype
).view(-1, 1, 1)
sampled_state = (
policy_mean.detach()
+ sigma.to(policy_mean.dtype) * noise
)
else:
sampled_state = policy_mean.detach()
if q_padding_mask is not None:
sampled_state = sampled_state.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
# Evaluate the sampled (detached) action under the differentiable
# policy mean. Omitting constants/log sigma is exact for gradients
# because sigma is detached and not a learned parameter.
standardized = (
sampled_state.float() - policy_mean.float()
) / sigma
score = -0.5 * standardized.square()
log_prob = (score * valid_3d).sum(dim=(1, 2)) / n_valid_dims
log_probs.append(log_prob)
with torch.no_grad():
update = (policy_mean.detach() - observation).float()
perturbation = (
sampled_state - policy_mean.detach()
).float()
mean_update_rms.append(
((update.square() * valid_3d).sum(dim=(1, 2)) / n_valid_dims)
.sqrt()
.div(c1_rms)
)
noise_rms.append(
(
(perturbation.square() * valid_3d).sum(dim=(1, 2))
/ n_valid_dims
)
.sqrt()
.div(c1_rms)
)
residual = sampled_state.detach()
states.append(residual)
return states, LatentPolicyTrace(
log_probs=torch.stack(log_probs, dim=1),
mean_update_rms=torch.stack(mean_update_rms, dim=1),
noise_rms=torch.stack(noise_rms, dim=1),
)
def _encode_visual_counterfactual_policy(
self,
q_text: torch.Tensor,
q_padding_mask: torch.Tensor | None,
text_ctx,
visual_bundle,
*,
policy_std: float,
sample: bool,
sample_mask: torch.Tensor | None = None,
) -> tuple[list[torch.Tensor], LatentPolicyTrace]:
"""Sample the current full-state visual recurrence itself.
This is the score-function counterpart of
:meth:`_encode_visual_counterfactual_recurrence`, not the historical
text-only CRR policy above. Persistent visual rows ``E`` come from the
exact pretrained first multimodal read. ``R1`` is deterministic, and
for every later transition the existing recurrence-only L21 LoRA
parameterizes
``mu_k = R{k-1} + beta * (F_on([E; R{k-1}])_Q - R{k-1})``.
The complete next residual state is sampled around that mean. Sampled
states are detached before the next transition, so answer rewards can
reach the recurrent LoRA only through the Gaussian score. With
``sample=False`` this path is exactly the ordinary deterministic
full-state recurrence.
"""
if not math.isfinite(policy_std) or policy_std <= 0.0:
raise ValueError("policy_std must be finite and > 0")
if not getattr(self.config, "visual_full_state_recurrence", False):
raise ValueError(
"visual latent policy requires "
"visual_full_state_recurrence=True"
)
if not getattr(self.config, "recurrence_only_adapter", False):
raise ValueError(
"visual latent policy trains the existing recurrence-only "
"adapter"
)
steps = int(self.config.num_workspace_steps)
if steps < 2:
raise ValueError("latent policy requires at least two CRR steps")
beta = float(self.config.counterfactual_beta)
cell_index = int(self.config.ell_star) + 1
tm = self.text_model
mm_lower, mm_ctx, text_rows, visual_rows = visual_bundle
# B, R1, and persistent E are fixed observations. In particular, the
# recurrence LoRA is disabled for these competence-critical reads.
with self._recurrence_adapter_execution(False), torch.no_grad():
base_anchor = run_layer_range(
tm,
text_ctx,
cell_index,
cell_index + 1,
hidden_states=q_text,
use_cache=True,
)
first_full = run_layer_range(
tm,
replace(mm_ctx, past_key_values=None),
cell_index,
cell_index + 1,
hidden_states=mm_lower,
use_cache=False,
)
r1, r1_padding = select_tokens_padded(first_full, text_rows)
if q_padding_mask is not None and not torch.equal(
r1_padding, q_padding_mask
):
raise ValueError(
"multimodal and text-only padding disagree in visual policy"
)
residual = r1 - base_anchor
if q_padding_mask is not None:
residual = residual.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
base_anchor = base_anchor.detach()
residual = residual.detach()
states = [base_anchor, residual]
valid = (
torch.ones(
residual.shape[:2],
dtype=torch.bool,
device=residual.device,
)
if q_padding_mask is None
else ~q_padding_mask
)
valid_3d = valid.unsqueeze(-1)
n_valid_dims = (
valid.sum(dim=1).to(torch.float32) * residual.shape[-1]
).clamp_min(1.0)
c1_square_sum = (
residual.float().square() * valid_3d
).sum(dim=(1, 2))
c1_rms = (c1_square_sum / n_valid_dims).sqrt().clamp_min(1e-6)
sigma = (float(policy_std) * c1_rms).view(-1, 1, 1).detach()
normal_ctx = replace(mm_ctx, past_key_values=None)
log_probs: list[torch.Tensor] = []
mean_update_rms: list[torch.Tensor] = []
noise_rms: list[torch.Tensor] = []
for _ in range(1, steps):
observation = residual.detach()
state_before = base_anchor + observation
recurrent_input = _replace_selected_rows_from_padded(
first_full,
text_rows,
state_before,
r1_padding,
)
with self._recurrence_adapter_execution(True):
on_q, _, on_padding = self._visual_on_off_native_text_rows(
recurrent_input,
normal_ctx,
text_rows,
visual_rows,
cell_index,
compute_off=False,
)
if not torch.equal(on_padding, r1_padding):
raise RuntimeError("recurrent visual-policy row layout changed")
policy_mean_full = torch.lerp(state_before, on_q, beta)
policy_mean = policy_mean_full - base_anchor
if q_padding_mask is not None:
policy_mean = policy_mean.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
if sample:
noise = torch.randn_like(policy_mean)
if sample_mask is not None:
if sample_mask.shape != (policy_mean.shape[0],):
raise ValueError(
"latent policy sample mask must have shape "
f"{(policy_mean.shape[0],)}, got "
f"{tuple(sample_mask.shape)}"
)
noise = noise * sample_mask.to(
device=noise.device, dtype=noise.dtype
).view(-1, 1, 1)
sampled_state = (
policy_mean.detach()
+ sigma.to(policy_mean.dtype) * noise
)
else:
sampled_state = policy_mean.detach()
if q_padding_mask is not None:
sampled_state = sampled_state.masked_fill(
q_padding_mask.unsqueeze(-1), 0
)
standardized = (
sampled_state.float() - policy_mean.float()
) / sigma
score = -0.5 * standardized.square()
log_probs.append(
(score * valid_3d).sum(dim=(1, 2)) / n_valid_dims
)
with torch.no_grad():
update = (policy_mean.detach() - observation).float()
perturbation = (
sampled_state - policy_mean.detach()
).float()
mean_update_rms.append(
(
(update.square() * valid_3d).sum(dim=(1, 2))
/ n_valid_dims
).sqrt().div(c1_rms)
)
noise_rms.append(
(
(perturbation.square() * valid_3d).sum(dim=(1, 2))
/ n_valid_dims
).sqrt().div(c1_rms)
)
residual = sampled_state.detach()
states.append(residual)
return states, LatentPolicyTrace(
log_probs=torch.stack(log_probs, dim=1),
mean_update_rms=torch.stack(mean_update_rms, dim=1),
noise_rms=torch.stack(noise_rms, dim=1),
)
def _encode_raw_evidence_question(
self,
q_star: torch.Tensor,
v_star: torch.Tensor,
q_padding_mask: torch.Tensor | None,
v_padding_mask: torch.Tensor | None,
gate_overrides: list | None,
v_star_reread: torch.Tensor | None,
v_reread_padding_mask: torch.Tensor | None,
reread_memory_overrides: list | None,
reread_zero_memory: bool,
hierarchical_visual_memories: (
tuple[tuple[torch.Tensor, torch.Tensor], ...] | None
) = None,
) -> list[torch.Tensor]:
"""Slot-free recurrent read with immutable raw visual evidence.
Residual mode stores ``[R0=Q*, R1, ..., RT]``. State-replacement mode
stores ``[Q*, H1, ..., HT]`` and the decoder later receives ``Q*+HT``.
Question-anchored replacement stores ``[Z0=Q*, Z1, ..., ZT]`` where
``Zk=Q*+CrossAttn(Z{k-1},V*)`` and directly decodes ``ZT``.
In hierarchical mode the first T-1 steps read their prescribed frozen
visual memories and the final step reads Q* with the same CrossAttn:
``R_T = R_{T-1} + CA(R_{T-1}, Q*)``. Thus R_T integrates the visual
trajectory without another image access or another module. Visual
memories are never updated or exposed directly to the decoder.
"""
r = q_star
states = [r]
hierarchy = tuple(hierarchical_visual_memories or ())
if hierarchy:
expected_visual_reads = self.config.num_workspace_steps - 1
if len(hierarchy) != expected_visual_reads:
raise ValueError(
"hierarchical visual memories must cover T-1 reads: "
f"got {len(hierarchy)} for T={self.config.num_workspace_steps} "
f"(expected {expected_visual_reads})"
)
if v_star_reread is not None:
raise ValueError(
"hierarchical swaps must use per-step memory overrides"
)
base_e = base_pad = base_projected = None
question_projected = self.raw_eq_transition.project_evidence(q_star)
else:
base_e, base_pad = v_star, v_padding_mask
if v_star_reread is not None:
base_e, base_pad = v_star_reread, v_reread_padding_mask
if reread_zero_memory:
base_e = torch.zeros_like(base_e)
base_projected = self.raw_eq_transition.project_evidence(base_e)
for k in range(self.config.num_workspace_steps):
if hierarchy and k < len(hierarchy):
e_k, pad_k = hierarchy[k]
if reread_zero_memory:
e_k = torch.zeros_like(e_k)
projected_k = self.raw_eq_transition.project_evidence(e_k)
elif hierarchy:
# Final image-free integration: the evolved state queries the
# immutable text-only scaffold through the same shared module.
e_k, pad_k = q_star, q_padding_mask
projected_k = question_projected
else:
e_k, pad_k = base_e, base_pad
projected_k = base_projected
if reread_memory_overrides is not None and reread_memory_overrides[k] is not None:
e_k, pad_k = reread_memory_overrides[k]
if reread_zero_memory:
e_k = torch.zeros_like(e_k)
projected_k = self.raw_eq_transition.project_evidence(e_k)
scale = 1.0
if gate_overrides is not None and gate_overrides[k] is not None:
# Evaluation-only read ablation/scale; there is no learned gate.
scale = float(gate_overrides[k])
r = self.raw_eq_transition(
r,
e_k,
evidence_padding_mask=pad_k,
question_padding_mask=q_padding_mask,
read_scale=scale,
projected_evidence=projected_k,
question_anchor=(
q_star
if getattr(self.config, "raw_question_anchor", False)
else None
),
)
states.append(r)
self.last_reread_gates = []
self.last_reread_update_mags = []
return states
def encode_workspace(
self,
q_star: torch.Tensor, # [B, N_q, d]
v_star: torch.Tensor, # [B, N_v, d]
q_padding_mask: torch.Tensor | None = None,
v_padding_mask: torch.Tensor | None = None,
gate_overrides: list | None = None, # eval-only: per-step None/float
v_star_reread: torch.Tensor | None = None, # eval-only: swap rereads' V*
v_reread_padding_mask: torch.Tensor | None = None,
reread_memory_overrides: list | None = None, # eval-only: per-step (v, pad) | None
reread_zero_memory: bool = False, # eval-only: V* -> 0 at EVERY reread (spec 10.3)
final_read_spec: dict | None = None, # eval-only: latent-necessity battery
hierarchical_visual_memories: (
tuple[tuple[torch.Tensor, torch.Tensor], ...] | None
) = None,
) -> list[torch.Tensor]:
"""``Z^(0..K)``. ``V*`` is consumed here and nowhere else.
With ``config.read_gating`` the workspace may re-consult ``V*`` between
transitions through the shared reader (gated, see
:meth:`WorkspaceRead.reread`); the decoder-facing contract is unchanged
either way. Returns ``K + 1`` tensors of shape ``[B, S, d_w]``.
The three eval-only knobs implement §3.2's held-out gate battery:
``gate_overrides[k]`` forces step ``k``'s gate (forced-open/closed/
step-specific), ``v_star_reread`` feeds the REREADS a different visual
memory while the initial read keeps the true one (mid-recurrence swap).
Defaults leave the learned path bitwise intact. They are reachable only
through :meth:`prefill` / :meth:`generate_answer`, never ``forward``,
so the training loop cannot consume them even by accident.
"""
if getattr(self.config, "raw_evidence_question", False):
if final_read_spec is not None:
raise ValueError("raw_evidence_question does not use final_read_spec")
return self._encode_raw_evidence_question(
q_star,
v_star,
q_padding_mask,
v_padding_mask,
gate_overrides,
v_star_reread,
v_reread_padding_mask,
reread_memory_overrides,
reread_zero_memory,
hierarchical_visual_memories,
)
z = self.workspace_read(q_star, v_star, q_padding_mask, v_padding_mask)
gating = getattr(self.config, "read_gating", False)
pure = getattr(self.config, "read_residual_pure", False)
verify = getattr(self.config, "final_verify", False)
mid_step = int(getattr(self.config, "mid_read_step", 0) or 0)
sid = None
if getattr(self.config, "persistent_slot_id", False):
# E_slot을 매 step 입력에 일시 제공 (누적 금지). 저장 state는
# f_theta 출력 그대로다.
sid = self.workspace_read.slots.unsqueeze(0).to(z.dtype) # [1, S, d_w]
if getattr(self.config, "evidence_reasoning", False):
# advisor 최종안: Evidence–Reasoning 분리. 전용 경로로 위임 —
# flag off면 이 아래 기존 코드가 비트동일하게 돈다.
assert final_read_spec is None, "ER 모드는 final_read_spec 미지원"
return self._encode_workspace_er(
z, q_star, v_star, q_padding_mask, v_padding_mask,
gate_overrides, v_star_reread, v_reread_padding_mask,
reread_zero_memory, reread_memory_overrides, mid_step,
)
if final_read_spec is not None:
# Latent-necessity battery (advisor): run the ZERO-residual chain
# U^(0..T) (U^(0) := Z^(0)), then compose one final visual read
# with independently chosen query source / decoder base:
# Z_out = U^(base_t) + alpha * W_o XAttn(LN(query_src), V*).
# query_override (a tensor) replaces the query source entirely
# (matched query swap). final_read=False skips the read
# (latent-only / initial-read-only).
if not pure and not verify:
raise RuntimeError("final_read_spec requires a pure verifier branch")
spec = final_read_spec
chain = [z]
for _ in range(self.config.num_workspace_steps):
chain.append(self.workspace_transition(
chain[-1], q_star, q_padding_mask, cond=sid))
base = chain[int(spec.get("base_t", self.config.num_workspace_steps))]
if spec.get("final_read", True):
q_src = spec.get("query_override")
if q_src is None:
q_src = chain[int(spec.get("query_t", self.config.num_workspace_steps))]
a = self.workspace_read.visual_residual(q_src, v_star, v_padding_mask)
z_out = base + float(spec.get("alpha", 1.0)) * a
else:
z_out = base
self.last_reread_gates = []
self.last_reread_update_mags = []
self.last_reread_gates_full = None
self.last_reread_update_mags_full = None
self.last_chain_states = chain # U^(0..T), query-swap donor 추출용
return chain + [z_out]
if not gating and not pure and not verify and not mid_step:
del v_star # make the read-once contract visible at the call site
states = [z]
gates: list[torch.Tensor] = []
upd_mags: list[float] = []
upd_full: list[torch.Tensor] = []
for k in range(self.config.num_workspace_steps):
u = self.workspace_transition(z, q_star, q_padding_mask, cond=sid)
z = u
if mid_step and (k + 1) == mid_step:
# MAIN: single intermediate pure read at LOGICAL step k*
# (python index k = k*-1; the +1 here is the off-by-one guard
# the directive calls out). Eval hooks share the step-wise
# keying: gate_overrides[k]==0.0 -> mid-off,
# reread_memory_overrides[k] -> blank/swap, zero flag -> zero.
off_mid = gate_overrides is not None and gate_overrides[k] == 0.0
if not off_mid:
v_r, vp_r = v_star, v_padding_mask
if v_star_reread is not None:
v_r, vp_r = v_star_reread, v_reread_padding_mask
if reread_zero_memory:
v_r = torch.zeros_like(v_r)
elif (reread_memory_overrides is not None
and reread_memory_overrides[k] is not None):
v_r, vp_r = reread_memory_overrides[k]
z = self.workspace_read.reread_pure(
u, v_r, vp_r,
query_state=None if sid is None else u + sid)
# read 스케일: 학습은 config.mid_read_alpha, eval은
# gate_overrides가 절대값으로 대체 (sweep 의미 보존).
# reread_pure가 z = u + WoA이므로 (z-u) = WoA 정확.
a_eff = float(getattr(self.config, "mid_read_alpha", 1.0))
if gate_overrides is not None and gate_overrides[k] is not None:
a_eff = float(gate_overrides[k])
if a_eff != 1.0:
z = u + a_eff * (z - u)
with torch.no_grad():
dv = (z - u).float()
uf = u.float()
den = uf.norm(dim=(-2, -1)).clamp_min(1e-6)
self.last_mid_stats_full = {
"u_mid": dv.norm(dim=(-2, -1)) / den,
"a_mid": nn.functional.cosine_similarity(
uf.flatten(1), dv.flatten(1), dim=1),
"r_mid": 1 - nn.functional.cosine_similarity(
uf.flatten(1), z.float().flatten(1), dim=1),
"q_mid": z.float().norm(dim=(-2, -1)) / den,
}
# 시간축 diversity 감시 (advisor): 배치 내 항목쌍은 무관
# 입력이므로, pooled state의 평균 pairwise cos-dist가
# d_unrelated의 학습 중 근사가 된다. pre-read 수축이
# 스케줄 교정으로 완화되는지 이 곡선으로 판정.
if u.shape[0] > 1:
def _div(x): # [B, S, d_w] -> 평균 pairwise 1-cos
pm = nn.functional.normalize(
x.float().mean(1), dim=-1) # [B, d_w]
g = pm @ pm.T
b = g.shape[0]
off = g.masked_select(
~torch.eye(b, dtype=torch.bool, device=g.device))
return 1 - off.mean()
self.last_mid_stats_full["div_z0"] = _div(
states[0]).expand(u.shape[0]).clone()
self.last_mid_stats_full["div_z1"] = _div(
states[1] if len(states) > 1 else states[0]
).expand(u.shape[0]).clone()
self.last_mid_stats_full["div_umid"] = _div(u).expand(
u.shape[0]).clone()
if pure:
# Eval-only residual scale (advisor alpha-sweep): a non-None
# gate_overrides[k] acts as alpha in Z = U + alpha*W_o A.
# alpha = 0.0 skips the branch entirely (== read_once), so the
# sweep's zero point and reread-off are one code path.
alpha_k = None
if gate_overrides is not None and gate_overrides[k] is not None:
alpha_k = float(gate_overrides[k])
if alpha_k == 0.0:
states.append(z)
continue
# all-steps substitute (global null/swap) -- the v7-ep1 run
# exposed that this was wired only into the gated branch:
# null_mem came back bit-identical to native (0.0% changes),
# which is impossible with a real substitution.
v_r, vp_r = v_star, v_padding_mask
if v_star_reread is not None:
v_r, vp_r = v_star_reread, v_reread_padding_mask
if reread_zero_memory:
v_r, vp_r = torch.zeros_like(v_star), v_padding_mask
elif reread_memory_overrides is not None and reread_memory_overrides[k] is not None:
v_r, vp_r = reread_memory_overrides[k]
z = self.workspace_read.reread_pure(u, v_r, vp_r)
rho = float(getattr(self.config, "residual_norm_cap", 0.0) or 0.0)
if rho > 0.0:
dv = z - u
s_cap = (
rho * u.float().pow(2).mean((-2, -1)).sqrt()
/ dv.float().pow(2).mean((-2, -1)).sqrt().clamp_min(1e-6)
).clamp(max=1.0).to(dv.dtype) # [B]
z = u + s_cap.view(-1, 1, 1) * dv
if alpha_k is not None and alpha_k != 1.0:
z = u + alpha_k * (z - u)
with torch.no_grad():
num = (z - u).float().norm(dim=(-2, -1))
den = u.float().norm(dim=(-2, -1)).clamp_min(1e-6)
u_item = num / den
upd_mags.append(float(u_item.mean()))
upd_full.append(u_item)
if gating:
# Consultation happens through the same reader and lands in Z;
# f_theta itself still never receives V* (read-mediation holds,
# the "once" becomes "as often as the gate earns").
v_r = v_star if v_star_reread is None else v_star_reread
vp_r = (
v_padding_mask if v_star_reread is None else v_reread_padding_mask
)
if reread_memory_overrides is not None and reread_memory_overrides[k] is not None:
# step-specific visual ablation: THIS reread consults a
# substitute memory (blank / donor), all other steps keep
# the true V*. The decisive §3.2 battery test.
v_r, vp_r = reread_memory_overrides[k]
z, g = self.workspace_read.reread(
u, q_star, v_r, q_padding_mask, vp_r,
gate_override=None if gate_overrides is None else gate_overrides[k],
)
gates.append(g)
# u_k: realized update magnitude ||Z-U||_F / ||U||_F -- the
# gate says how much the model WANTED to look, u_k what the
# look actually changed (draft §3.2).
with torch.no_grad():
num = (z - u).float().norm(dim=(-2, -1))
den = u.float().norm(dim=(-2, -1)).clamp_min(1e-6)
u_item = num / den # [B]
upd_mags.append(float(u_item.mean()))
upd_full.append(u_item)
states.append(z)
if verify:
# MAIN: single final visual verification. Mid-loop above ran pure
# f_theta only (no reader calls -- tests/test_final_verify.py
# counts them). Eval hooks: gate_overrides[-1]==0.0 -> branch off;
# reread_memory_overrides[-1] -> substitute verifier memory;
# reread_zero_memory -> zero tensor.
off = gate_overrides is not None and gate_overrides[-1] == 0.0
if not off:
v_r, vp_r = v_star, v_padding_mask
if reread_zero_memory:
v_r = torch.zeros_like(v_star)
elif (reread_memory_overrides is not None
and reread_memory_overrides[-1] is not None):
v_r, vp_r = reread_memory_overrides[-1]
states.append(self.workspace_read.reread_pure(states[-1], v_r, vp_r))
else:
states.append(states[-1])
if mid_step and len(states) == self.config.num_workspace_steps + 1:
# D_{k*->j} propagation telemetry (advisor: "step 4 update가 빠르게
# 소거되면 구조는 아름답지만 답에 기여하지 않는다"). Shadow chain =
# mid-null (read skipped); identical to clean bitwise before k*, so
# start it at Z^(k*-1). workspace_dropout=0.0 -> deterministic.
with torch.no_grad():
zs = states[mid_step - 1]
prop = []
for j in range(mid_step, len(states)):
zs = self.workspace_transition(
zs, q_star, q_padding_mask, cond=sid)
prop.append(1 - nn.functional.cosine_similarity(
states[j].float().flatten(1), zs.float().flatten(1), dim=1))
# [B, T-k*+1]; columns j = k*, ..., T
self.last_mid_prop_full = torch.stack(prop, dim=1)
# Telemetry for the trainer/analysis: per-step means, [K] each or [].
self.last_reread_gates = [float(g.detach().mean()) for g in gates]
self.last_reread_update_mags = upd_mags
# Full per-item resolution for offline probes (compact logs only carry
# the means). [B, K] or None when gating is off.
self.last_reread_gates_full = (
torch.stack([g.detach().float() for g in gates], dim=1) if gates else None
)
self.last_reread_update_mags_full = (
torch.stack(upd_full, dim=1) if upd_full else None
) # [B, K]
return states
# -- prefill -----------------------------------------------------------
def prefill(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
attention_mask: torch.Tensor | None = None,
question_attention_mask: torch.Tensor | None = None,
workspace_override: torch.Tensor | None = None,
gate_overrides: list | None = None,
reread_inputs: dict | None = None,
reread_step_inputs: dict | None = None,
reread_zero_memory: bool = False,
final_read_spec: dict | None = None,
return_decoder_latent: bool = False,
latent_policy_std: float | None = None,
latent_policy_sample: bool = True,
latent_policy_sample_mask: torch.Tensor | None = None,
return_latent_policy_trace: bool = False,
curriculum_readout_depth: int | torch.Tensor | None = None,
component_ablation_mode: str | None = None,
):
"""Run §3.2 end to end and leave a complete replacement cache behind.
Shared by training and generation, so the two cannot drift apart.
Args:
workspace_override: substitute ``Z^(K)`` (``[B, S, d_w]``). Held-out
causal evaluation only -- content replacement, matched swap,
path-leakage. Never set during training.
gate_overrides: per-step forced gate values (``None`` entries keep
the learned gate). Held-out gate battery only.
reread_inputs: multimodal inputs (``input_ids``/``pixel_values``/
``image_grid_thw``/``attention_mask``) of a DIFFERENT item; its
``V*`` is served to the re-reads while the initial read keeps
the true one (mid-recurrence swap). Held-out battery only.
reread_step_inputs: ``{step_k: inputs-dict}`` -- like
``reread_inputs`` but applied ONLY at the given steps' rereads
(step-specific blank/swap ablation). Held-out battery only.
Returns:
Normally ``(upper, cache, states, v_pooled)``. With
``return_decoder_latent=True``, appends the exact latent supplied to
the decoder (``R_T`` in the raw-E/Q method). This opt-in fifth
return value keeps old callers compatible while giving causal
evaluations a non-ambiguous intervention target.
"""
cfg = self.config
allowed_component_modes = {
None,
"without_recurrence",
"without_native_mm_initialization",
"without_learned_transition",
"without_entire_visual_path",
"without_visual_reread",
"without_visual_rows_keep_h1",
}
if component_ablation_mode not in allowed_component_modes:
raise ValueError(
f"unknown component_ablation_mode={component_ablation_mode!r}"
)
if component_ablation_mode is not None and self.training:
raise ValueError("component ablations are evaluation-only")
policy_enabled = latent_policy_std is not None
if return_latent_policy_trace and not policy_enabled:
raise ValueError(
"return_latent_policy_trace requires latent_policy_std"
)
if policy_enabled and not getattr(
cfg, "counterfactual_residual_recurrence", False
):
raise ValueError("latent policy is implemented only for CRR")
if curriculum_readout_depth is not None and policy_enabled:
raise ValueError(
"program curriculum and latent-policy sampling are mutually exclusive"
)
if curriculum_readout_depth is not None and workspace_override is not None:
raise ValueError(
"program curriculum and workspace_override are mutually exclusive"
)
tm = self.text_model
n_q = question_ids.shape[1]
# Consume-once training buffers must never survive into a later batch.
self._functional_q_target = None
self._functional_q_target_pad = None
self._functional_q_target_ids = None
self._functional_state_loss = None
self._counterfactual_q_state = None
self._counterfactual_q_pad = None
self._counterfactual_q_ids = None
self._visual_counterfactual_bundle = None
q_pad = (
None if question_attention_mask is None else (question_attention_mask == 0)
) # [B, N_q] True == ignore
# 1. Frozen visual memory. The hierarchical raw path captures one
# memory at each prescribed depth during a single multimodal pass.
hierarchical_visual_memories = None
if getattr(cfg, "hierarchical_visual_layers", ()):
hierarchical_visual_memories = (
self._encode_hierarchical_visual_memories(
input_ids,
pixel_values,
image_grid_thw,
attention_mask,
)
)
# Preserve the historical helper plumbing; the raw encoder below
# consumes the full tuple rather than this final-memory alias.
v_star, v_pad = hierarchical_visual_memories[-1]
else:
v_star, v_pad = self._encode_visual_memory(
input_ids, pixel_values, image_grid_thw, attention_mask
) # [B, N_v, d], [B, N_v]
# 2. Text-only lower prefill. Raw-E/Q uses Q rows only; historical slot
# methods append fixed cache-shape placeholders.
q_star, ctx, cache = self._prefill_text_branch(
question_ids, question_attention_mask
) # q_star [B, N_q, d]
if getattr(cfg, "counterfactual_residual_recurrence", False):
unsupported = {
"gate_overrides": gate_overrides is not None,
"reread_inputs": reread_inputs is not None,
"reread_step_inputs": reread_step_inputs is not None,
"reread_zero_memory": bool(reread_zero_memory),
"final_read_spec": final_read_spec is not None,
}
active = [name for name, enabled in unsupported.items() if enabled]
if active:
raise ValueError(
"CRR has no raw-evidence reread controls; unsupported: "
+ ", ".join(active)
)
q_mm = self._counterfactual_q_state
q_mm_pad = self._counterfactual_q_pad
q_mm_ids = self._counterfactual_q_ids
visual_bundle = self._visual_counterfactual_bundle
self._counterfactual_q_state = None
self._counterfactual_q_pad = None
self._counterfactual_q_ids = None
self._visual_counterfactual_bundle = None
if q_mm is None or q_mm_pad is None or q_mm_ids is None:
raise RuntimeError("CRR factual question state was not captured")
if q_mm.shape != q_star.shape:
raise ValueError(
"multimodal/text-only question rows disagree after removing "
f"the vision span: {tuple(q_mm.shape)} vs {tuple(q_star.shape)}"
)
valid = ~q_mm_pad
if q_pad is not None:
if q_pad.shape != q_mm_pad.shape or not torch.equal(
q_pad, q_mm_pad
):
raise ValueError(
"multimodal and text-only question padding masks disagree"
)
valid &= ~q_pad
if not torch.equal(q_mm_ids[valid], question_ids[valid]):
raise ValueError(
"text-only token ids do not equal multimodal prompt text rows"
)
policy_trace = None
if policy_enabled:
if getattr(
cfg, "visual_counterfactual_recurrence", False
):
z_states, policy_trace = (
self._encode_visual_counterfactual_policy(
q_star,
q_pad,
ctx,
visual_bundle,
policy_std=float(latent_policy_std),
sample=bool(latent_policy_sample),
sample_mask=latent_policy_sample_mask,
)
)
else:
z_states, policy_trace = (
self._encode_counterfactual_residual_policy(
q_star,
q_mm,
q_pad,
ctx,
policy_std=float(latent_policy_std),
sample=bool(latent_policy_sample),
sample_mask=latent_policy_sample_mask,
)
)
else:
z_states = self._encode_counterfactual_residual_recurrence(
q_star,
q_mm,
q_pad,
ctx,
visual_bundle=visual_bundle,
component_ablation_mode=component_ablation_mode,
)
if curriculum_readout_depth is not None:
z_final = self.decoder_latent_at_depth(
z_states, curriculum_readout_depth
)
else:
z_final = (
self.decoder_latent_from_states(z_states)
if workspace_override is None
else workspace_override
)
del v_star
lower = ctx.hidden_states
if z_final.shape != lower.shape:
raise ValueError(
"CRR decoder latent must match Q* shape, got "
f"{tuple(z_final.shape)} vs {tuple(lower.shape)}"
)
# In latent-policy mode the generated answer is a non-differentiable
# environment outcome. Detaching here makes the score trace above
# the sole gradient path even if a caller accidentally leaves an
# upper-decoder parameter trainable.
decoder_input = (
z_final.detach() if policy_enabled else z_final
).to(lower.dtype)
if policy_enabled:
with torch.no_grad():
upper = run_layer_range(
tm,
ctx,
self._upper_decoder_start(),
None,
hidden_states=decoder_input,
use_cache=True,
)
else:
upper = run_layer_range(
tm,
ctx,
self._upper_decoder_start(),
None,
hidden_states=decoder_input,
use_cache=True,
)
if return_latent_policy_trace:
if return_decoder_latent:
return upper, cache, z_states, None, z_final, policy_trace
return upper, cache, z_states, None, policy_trace
if return_decoder_latent:
return upper, cache, z_states, None, z_final
return upper, cache, z_states, None
# The removed visual-anchor loss must have zero runtime footprint in
# the current answer-only method. Historical anchor runs opt in through
# config.read_anchor.
v_pooled = None
if getattr(cfg, "read_anchor", False):
if v_pad is None:
v_pooled = v_star.float().mean(1)
else:
keep = (~v_pad).to(v_star.dtype).unsqueeze(-1)
v_pooled = (
(v_star * keep).sum(1).float()
/ keep.sum(1).float().clamp_min(1.0)
)
v_pooled = torch.nn.functional.layer_norm(
v_pooled, (v_pooled.shape[-1],)
).detach()
# 3. Z^(0..K). V* dies here.
v_swap, v_swap_pad = None, None
step_overrides = None
if (
hierarchical_visual_memories is not None
and reread_inputs is not None
and "__v_star__" in reread_inputs
):
raise ValueError(
"hierarchical recurrence requires one substituted memory per "
"depth; a single precomputed __v_star__ is ambiguous"
)
if reread_inputs is not None and "__v_star__" in reread_inputs:
# eval-only: precomputed substitute memory tensor (e.g. corpus-mean
# dictionary), bypassing the image encoder.
v_swap = reread_inputs["__v_star__"]
v_swap_pad = reread_inputs.get("__v_pad__")
elif (
reread_inputs is not None
and hierarchical_visual_memories is not None
):
swapped_memories = self._encode_hierarchical_visual_memories(
reread_inputs["input_ids"],
reread_inputs["pixel_values"],
reread_inputs["image_grid_thw"],
reread_inputs.get("attention_mask"),
)
step_overrides = [None] * self.config.num_workspace_steps
step_overrides[: len(swapped_memories)] = swapped_memories
elif reread_inputs is not None:
v_swap, v_swap_pad = self._encode_visual_memory(
reread_inputs["input_ids"],
reread_inputs["pixel_values"],
reread_inputs["image_grid_thw"],
reread_inputs.get("attention_mask"),
)
if reread_step_inputs is not None:
if step_overrides is None:
step_overrides = [None] * self.config.num_workspace_steps
for k_, inp in reread_step_inputs.items():
step_index = int(k_)
if not 0 <= step_index < self.config.num_workspace_steps:
raise ValueError(
"reread step index must be in "
f"[0,{self.config.num_workspace_steps - 1}], got "
f"{step_index}"
)
if "__v_star__" in inp:
# eval-only: direct tensor (e.g. true zero memory) -- the
# blank IMAGE still yields nonzero V*, so "zero" must be
# injected as a tensor, not an image (advisor §4).
step_overrides[step_index] = (
inp["__v_star__"], inp.get("__v_pad__")
)
elif hierarchical_visual_memories is not None:
if step_index >= len(hierarchical_visual_memories):
raise ValueError(
"the final hierarchical step reads Q*, not an image; "
"use a direct tensor override only for an explicit "
"integration-step intervention"
)
replacement_memories = (
self._encode_hierarchical_visual_memories(
inp["input_ids"],
inp["pixel_values"],
inp["image_grid_thw"],
inp.get("attention_mask"),
)
)
step_overrides[step_index] = replacement_memories[step_index]
else:
step_overrides[step_index] = self._encode_visual_memory(
inp["input_ids"], inp["pixel_values"],
inp["image_grid_thw"], inp.get("attention_mask"),
)
z_states = self.encode_workspace(
q_star, v_star, q_pad, v_pad,
gate_overrides=gate_overrides,
v_star_reread=v_swap,
v_reread_padding_mask=v_swap_pad,
reread_memory_overrides=step_overrides,
reread_zero_memory=reread_zero_memory,
final_read_spec=final_read_spec,
hierarchical_visual_memories=hierarchical_visual_memories,
)
del v_star
if getattr(cfg, "raw_evidence_question", False):
# Residual: states=[R0=Q*, R1..RT], decode RT or mean(R1..RT).
# Replacement: states=[Q*, H1..HT], decode Q*+HT.
# Anchored replacement: states=[Q*, Z1..ZT], decode ZT.
# The selected mean aggregation is parameter-free; no learned
# temporal interface is allocated.
z_final = (
self.decoder_latent_from_states(z_states)
if workspace_override is None
else workspace_override
)
if (
self.training
and getattr(cfg, "functional_state_loss", False)
and workspace_override is None
):
target = self._functional_q_target
target_pad = self._functional_q_target_pad
target_ids = self._functional_q_target_ids
self._functional_q_target = None
self._functional_q_target_pad = None
self._functional_q_target_ids = None
if target is None or target_pad is None or target_ids is None:
raise RuntimeError("functional-state target was not captured")
if target.shape != q_star.shape:
raise ValueError(
"multimodal/text-only question rows disagree after the "
f"vision span is removed: {tuple(target.shape)} vs "
f"{tuple(q_star.shape)}"
)
valid = ~target_pad
if q_pad is not None:
if q_pad.shape != target_pad.shape or not torch.equal(q_pad, target_pad):
raise ValueError(
"multimodal and text-only question padding masks disagree"
)
valid &= ~q_pad
if not torch.equal(target_ids[valid], question_ids[valid]):
raise ValueError(
"text-only token ids do not equal the multimodal prompt "
"with its vision span removed"
)
# Scale-free visual-state restoration:
# ||R_agg-Q*_mm||^2 / sg(||Q*_mm-Q*_text||^2).
# The denominator removes the large shared text scaffold and
# defines 1.0 as the image-independent Q* floor. Per-example
# reduction prevents long prompts from dominating.
mask = valid.unsqueeze(-1).to(torch.float32)
err = (z_final.float() - target.float()).square() * mask
visual = (target.float() - q_star.float()).square() * mask
err_sum = err.sum(dim=(1, 2))
visual_sum = visual.sum(dim=(1, 2)).detach()
# The normal multimodal branch always has a nonzero image
# residual, but retain a finite floor for degenerate fixtures.
floor = mask.sum(dim=(1, 2)).clamp_min(1.0) * 1e-8
self._functional_state_loss = (
err_sum / visual_sum.clamp_min(floor)
).mean()
k_ov = int(getattr(cfg, "er_step_override", 0) or 0)
if k_ov and workspace_override is None:
if not 1 <= k_ov <= cfg.num_workspace_steps:
raise ValueError(
f"er_step_override must be in [1,{cfg.num_workspace_steps}]"
)
z_final = (
z_states[0] + z_states[k_ov]
if (
getattr(cfg, "raw_state_replacement", False)
and not getattr(cfg, "raw_question_anchor", False)
)
else z_states[k_ov]
)
lower = ctx.hidden_states # [B,N_q,d], lower-cache values stay text-only
if z_final.shape != lower.shape:
raise ValueError(
"raw-evidence decoder latent must match Q* shape, got "
f"{tuple(z_final.shape)} vs {tuple(lower.shape)}"
)
upper = run_layer_range(
tm,
ctx,
cfg.ell_star + 1,
None,
hidden_states=z_final.to(lower.dtype),
use_cache=True,
)
if return_decoder_latent:
return upper, cache, z_states, v_pooled, z_final
return upper, cache, z_states, v_pooled
# 4. overwrite the slot positions with P_Z(Z^(K)). Everything below l*
# was computed on a sequence with no image tokens, so this is the
# first and only point at which visual information enters.
lower = ctx.hidden_states # [B, N_q + S, d]
z_final = z_states[-1] if workspace_override is None else workspace_override
if (getattr(self.config, "evidence_reasoning", False)
and workspace_override is None):
# strict ER: decoder 입력 = answer 체인의 R^(T) 단독 —
# E 직행 우회 없음. consume-once로 stale-state 차단.
r_fin = self._er_answer_final
assert r_fin is not None, "ER: encode 없이 prefill 시도 (stale/미설정)"
self._er_answer_final = None
z_final = r_fin
# 진단 전용 (checkpoint eval): R^(k)로 디코드하는 step-sweep.
# ER states = [E^0, R_t^1..R_t^T]이고 dropout=0에서 R_t == R 비트동일.
# 학습은 이 config 속성이 없으므로 경로 불변 (advisor: 훈련 무접촉).
k_ov = int(getattr(self.config, "er_step_override", 0) or 0)
if k_ov:
T = self.config.num_workspace_steps
assert 1 <= k_ov <= T, f"er_step_override must be in [1,{T}]"
z_final = z_states[k_ov]
z_proj = self.workspace_out(z_final).to(lower.dtype) # [B, S, d]
if self.splice_gain is not None:
# RMS-match the slot rows to the (non-pad) question rows. Fresh
# linear output is O(1) while layer-20 hiddens run into the
# hundreds; unmatched, the decoder's attention treats the slots as
# near-zero vectors and ignoring them is free (probe 20143).
if question_attention_mask is not None:
qm = (question_attention_mask > 0).to(lower.dtype) # [B, N_q]
else:
qm = torch.ones(
lower.shape[0], n_q, dtype=lower.dtype, device=lower.device
)
q_rms = (
(lower[:, :n_q].float().pow(2).mean(-1).sqrt() * qm).sum(1)
/ qm.sum(1).clamp_min(1.0)
) # [B]
z_rms = z_proj.float().pow(2).mean(-1).sqrt().mean(1) # [B]
scale = (q_rms / z_rms.clamp_min(1e-6)).view(-1, 1, 1).to(z_proj.dtype)
z_proj = z_proj * scale * self.splice_gain.to(z_proj.dtype)
spliced = torch.cat([lower[:, :n_q], z_proj], dim=1) # [B, N_q + S, d]
if (self.training and getattr(cfg, "interface_loss", False)
and workspace_override is None
and getattr(self, "_iface_h_base", None) is not None):
# iface student: R^(T).detach() → 동일 splice 변환(bridge 재적용)
# → layer l*+1 한 층만 재계산 → 마지막 slot 위치 hidden.
# gradient는 workspace_out·splice_gain·해당 층 adapter에만 흐른다
# (R/f_theta/reader/E 차단 — advisor bridge-only 규정).
zd = self.workspace_out(z_final.detach()).to(lower.dtype)
if self.splice_gain is not None:
zd_rms = zd.float().pow(2).mean(-1).sqrt().mean(1) # [B]
sc = (q_rms / zd_rms.clamp_min(1e-6)).view(-1, 1, 1).to(zd.dtype)
zd = zd * sc * self.splice_gain.to(zd.dtype)
spliced_d = torch.cat([lower[:, :n_q].detach(), zd], dim=1)
# cache 없는 ctx 사본 필수: 원본 ctx로 돌리면 이 패스가 l*+1층
# cache에 slot 키를 저장해 본 upper 실행의 키가 2배로 중복된다
# (게이트에서 [*, 111, 222] shape 오류로 검출).
import dataclasses as _dc
ctx_if = _dc.replace(ctx, past_key_values=None)
h_if = run_layer_range(
tm, ctx_if, cfg.ell_star + 1, cfg.ell_star + 2,
hidden_states=spliced_d) # [B, N_q+S, d]
h_r = h_if[:, -1].float() # 마지막 prefix 위치 (= 마지막 slot)
hb = self._iface_h_base.float()
self._iface_h_base = None
d_ = h_r.shape[-1]
self._iface_loss = (1 - nn.functional.cosine_similarity(
nn.functional.layer_norm(h_r, (d_,)),
nn.functional.layer_norm(hb, (d_,)), dim=-1)).mean()
# 5. F_{>l*} over M* = [Q*; P_Z(Z^(K))], completing the cache.
upper = run_layer_range(
tm, ctx, cfg.ell_star + 1, None, hidden_states=spliced, use_cache=True
) # [B, N_q + S, d]
if return_decoder_latent:
return upper, cache, z_states, v_pooled, z_final
return upper, cache, z_states, v_pooled
@torch.no_grad()
def _greedy_decode_from_prefill(
self,
upper: torch.Tensor,
cache: Cache,
question_ids: torch.LongTensor,
question_attention_mask: torch.Tensor | None,
*,
max_new_tokens: int,
eos_token_id: int | None,
) -> torch.LongTensor:
"""Greedily decode from an already-built replacement prefix cache."""
if max_new_tokens < 1:
raise ValueError("max_new_tokens must be >= 1")
tm = self.text_model
b = question_ids.shape[0]
q_state_interface = self._uses_question_state_interface()
n_ctx = question_ids.shape[1] + (
0 if q_state_interface else self.config.num_workspace_slots
)
device = question_ids.device
q_mask = question_attention_mask
if q_mask is None:
q_mask = torch.ones_like(question_ids, dtype=torch.long)
if q_state_interface:
prefix_mask = q_mask
logical_prefix_lengths = (q_mask > 0).sum(dim=-1).long()
last_q = (logical_prefix_lengths - 1).clamp_min(0)
seed_hidden = upper[
torch.arange(b, device=device), last_q
].unsqueeze(1)
else:
prefix_mask, _, logical_prefix_lengths = _replacement_prefix_layout(
q_mask, self.config.num_workspace_slots
)
seed_hidden = upper[:, -1:]
nxt = self.backbone.lm_head(final_norm(tm, seed_hidden)).argmax(-1) # [B, 1]
out = [nxt]
done = torch.zeros(b, dtype=torch.bool, device=device)
if eos_token_id is not None:
done |= nxt.squeeze(1) == eos_token_id
for step in range(1, max_new_tokens):
if bool(done.all()):
break
pos = n_ctx + step - 1
embeds = self.vl.get_input_embeddings()(nxt) # [B, 1, d]
cache_position = torch.tensor([pos], device=device)
logical_pos = logical_prefix_lengths + step - 1 # [B]
position_ids = logical_pos.view(1, b, 1).expand(3, -1, -1)
generated_mask = torch.ones(
b, step, dtype=prefix_mask.dtype, device=device
)
upper_mask = torch.cat([prefix_mask, generated_mask], dim=-1)
if q_state_interface:
lower_mask = upper_mask
else:
lower_prefix_mask = prefix_mask.clone()
lower_prefix_mask[:, -self.config.num_workspace_slots:] = 0
lower_mask = torch.cat([lower_prefix_mask, generated_mask], dim=-1)
h = self._decode_token_block(
embeds,
position_ids,
lower_mask,
upper_mask,
cache,
cache_position,
)
nxt = self.backbone.lm_head(final_norm(tm, h)).argmax(-1) # [B, 1]
if eos_token_id is not None:
nxt = torch.where(done.unsqueeze(1), nxt.new_full(nxt.shape, eos_token_id), nxt)
done |= nxt.squeeze(1) == eos_token_id
out.append(nxt)
return torch.cat(out, dim=1)
@torch.no_grad()
def _gold_answer_logprob_from_prefill(
self,
upper: torch.Tensor,
cache: Cache,
question_ids: torch.LongTensor,
question_attention_mask: torch.Tensor | None,
answer_ids: torch.LongTensor,
labels: torch.LongTensor,
) -> torch.FloatTensor:
"""Score a gold answer against an already sampled latent prefix.
The score is the mean teacher-forced log probability over valid answer
tokens. Running this helper under ``no_grad`` is intentional: the
decoder acts only as a frozen reward environment, while policy gradient
reaches the latent adapter solely through ``LatentPolicyTrace``.
"""
if answer_ids.ndim != 2 or labels.shape != answer_ids.shape:
raise ValueError(
"answer_ids and labels must have the same [B, N_a] shape"
)
if answer_ids.shape[0] != question_ids.shape[0]:
raise ValueError("answer and question batch sizes disagree")
if answer_ids.shape[1] < 1:
raise ValueError("gold-answer scoring requires at least one token")
tm = self.text_model
batch_size, answer_length = answer_ids.shape
q_state_interface = self._uses_question_state_interface()
q_mask = question_attention_mask
if q_mask is None:
q_mask = torch.ones_like(question_ids, dtype=torch.long)
if q_state_interface:
prefix_mask = q_mask
logical_prefix_lengths = (q_mask > 0).sum(dim=-1).long()
last_q = (logical_prefix_lengths - 1).clamp_min(0)
hidden_for_logits = upper[
torch.arange(batch_size, device=upper.device), last_q
].unsqueeze(1)
else:
prefix_mask, _, logical_prefix_lengths = _replacement_prefix_layout(
q_mask, self.config.num_workspace_slots
)
hidden_for_logits = upper[:, -1:]
# The final prefix state predicts y_0. Consuming gold y_0..y_{N-2}
# supplies the remaining N-1 teacher-forced prediction states; there is
# no reason to decode y_{N-1}, whose output would predict beyond labels.
if answer_length > 1:
teacher_ids = answer_ids[:, :-1]
teacher_mask = labels[:, :-1].ne(-100).to(prefix_mask.dtype)
n_ctx = question_ids.shape[1] + (
0 if q_state_interface else self.config.num_workspace_slots
)
teacher_embeds = self.vl.get_input_embeddings()(teacher_ids)
cache_position = torch.arange(
n_ctx,
n_ctx + teacher_ids.shape[1],
device=teacher_embeds.device,
)
logical_answer_pos = (
logical_prefix_lengths[:, None]
+ torch.arange(
teacher_ids.shape[1], device=teacher_embeds.device
)[None, :]
)
answer_positions = logical_answer_pos.unsqueeze(0).expand(
3, -1, -1
)
upper_mask = torch.cat([prefix_mask, teacher_mask], dim=-1)
if q_state_interface:
lower_mask = upper_mask
else:
lower_prefix_mask = prefix_mask.clone()
lower_prefix_mask[:, -self.config.num_workspace_slots:] = 0
lower_mask = torch.cat(
[lower_prefix_mask, teacher_mask], dim=-1
)
answer_hidden = self._decode_token_block(
teacher_embeds,
answer_positions,
lower_mask,
upper_mask,
cache,
cache_position,
)
hidden_for_logits = torch.cat(
[hidden_for_logits, answer_hidden], dim=1
)
logits = self.backbone.lm_head(final_norm(tm, hidden_for_logits))
nll, valid = _answer_nll_per_example(logits, labels)
if not bool(valid.all()):
raise ValueError("every latent-policy reward row needs a gold token")
return -nll
@torch.no_grad()
def _gold_answer_logprobs_by_crr_prefix(
self,
upper: torch.Tensor,
cache: Cache,
z_states: list[torch.Tensor] | tuple[torch.Tensor, ...],
question_ids: torch.LongTensor,
question_attention_mask: torch.Tensor | None,
answer_ids: torch.LongTensor,
labels: torch.LongTensor,
) -> torch.FloatTensor:
"""Score ``Z_1..Z_T`` for transition-aware latent RL.
The final prefix reuses the factual cached decode. Earlier prefixes
are stacked into one frozen L22--L27 replay and share the same
text-only lower answer states, exactly preserving the strict CRR
interface. No reward tensor returned here carries a decoder gradient.
"""
if not getattr(
self.config, "counterfactual_residual_recurrence", False
):
raise ValueError("prefix reward scoring requires CRR")
expected_states = int(self.config.num_workspace_steps) + 1
if len(z_states) != expected_states:
raise ValueError(
f"expected B plus {expected_states - 1} CRR states, "
f"got {len(z_states)}"
)
if answer_ids.ndim != 2 or labels.shape != answer_ids.shape:
raise ValueError(
"answer_ids and labels must have the same [B, N_a] shape"
)
batch_size, answer_length = answer_ids.shape
if batch_size != question_ids.shape[0] or answer_length < 1:
raise ValueError("prefix rewards require a non-empty aligned batch")
tm = self.text_model
num_steps = expected_states - 1
n_q = question_ids.shape[1]
q_mask = question_attention_mask
if q_mask is None:
q_mask = torch.ones_like(question_ids, dtype=torch.long)
q_keep = q_mask.gt(0)
logical_prefix_lengths = q_keep.sum(dim=-1).long()
last_q = (logical_prefix_lengths - 1).clamp_min(0)
prefixes: list[torch.Tensor] = []
residual_prefix: list[torch.Tensor] = []
for residual in z_states[1:]:
residual_prefix.append(residual)
if getattr(
self.config, "visual_cumulative_recurrence", False
):
# Full-state CVRR decodes the current recurrent state, not the
# historical mean of all residual prefixes.
prefixes.append(z_states[0] + residual)
else:
prefixes.append(
z_states[0]
+ self._aggregate_crr_residuals(residual_prefix)
)
final_hidden_for_logits = upper[
torch.arange(batch_size, device=upper.device), last_q
].unsqueeze(1)
answer_lower_hidden = None
answer_positions = None
combined_mask = None
if answer_length > 1:
teacher_ids = answer_ids[:, :-1]
teacher_mask = labels[:, :-1].ne(-100).to(q_mask.dtype)
teacher_embeds = self.vl.get_input_embeddings()(teacher_ids)
cache_position = torch.arange(
n_q,
n_q + teacher_ids.shape[1],
device=teacher_embeds.device,
)
logical_answer_pos = (
logical_prefix_lengths[:, None]
+ torch.arange(
teacher_ids.shape[1], device=teacher_embeds.device
)[None, :]
)
answer_positions = logical_answer_pos.unsqueeze(0).expand(
3, -1, -1
)
combined_mask = torch.cat([q_mask, teacher_mask], dim=-1)
final_answer_hidden, answer_lower_hidden = self._decode_token_block(
teacher_embeds,
answer_positions,
combined_mask,
combined_mask,
cache,
cache_position,
return_lower=True,
)
final_hidden_for_logits = torch.cat(
[final_hidden_for_logits, final_answer_hidden], dim=1
)
final_logits = self.backbone.lm_head(
final_norm(tm, final_hidden_for_logits)
)
final_nll, final_valid = _answer_nll_per_example(final_logits, labels)
if not bool(final_valid.all()):
raise ValueError("every latent-policy reward row needs a gold token")
if num_steps == 1:
return -final_nll.unsqueeze(1)
num_partial = num_steps - 1
partial_prefix = torch.cat(prefixes[:-1], dim=0)
q_pos_1d = q_keep.long().cumsum(dim=-1) - 1
q_pos_1d = q_pos_1d.masked_fill(~q_keep, 0)
q_positions = q_pos_1d.unsqueeze(0).expand(3, -1, -1)
partial_q_positions = q_positions.repeat(1, num_partial, 1)
partial_q_mask = q_mask.repeat(num_partial, 1)
if answer_lower_hidden is not None:
partial_input = torch.cat(
[
partial_prefix.to(answer_lower_hidden.dtype),
answer_lower_hidden.repeat(num_partial, 1, 1),
],
dim=1,
)
partial_positions = torch.cat(
[
partial_q_positions,
answer_positions.repeat(1, num_partial, 1),
],
dim=-1,
)
partial_attention_mask = combined_mask.repeat(num_partial, 1)
else:
partial_input = partial_prefix
partial_positions = partial_q_positions
partial_attention_mask = partial_q_mask
partial_ctx = make_split_context(
tm,
partial_input,
partial_positions,
partial_attention_mask,
)
partial_upper = run_layer_range(
tm,
partial_ctx,
self._upper_decoder_start(),
None,
hidden_states=partial_input,
use_cache=False,
)
partial_last_q = last_q.repeat(num_partial)
partial_hidden_for_logits = partial_upper[
torch.arange(partial_upper.shape[0], device=partial_upper.device),
partial_last_q,
].unsqueeze(1)
if answer_lower_hidden is not None:
partial_hidden_for_logits = torch.cat(
[
partial_hidden_for_logits,
partial_upper[
:, n_q : n_q + answer_lower_hidden.shape[1]
],
],
dim=1,
)
partial_logits = self.backbone.lm_head(
final_norm(tm, partial_hidden_for_logits)
)
partial_nll, partial_valid = _answer_nll_per_example(
partial_logits,
labels.repeat(num_partial, 1),
)
if not bool(partial_valid.all()):
raise ValueError("every CRR prefix needs a valid gold-answer score")
partial_scores = -partial_nll.view(num_partial, batch_size).transpose(
0, 1
)
return torch.cat([partial_scores, -final_nll.unsqueeze(1)], dim=1)
@torch.no_grad()
def generate_answer(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
attention_mask: torch.Tensor | None = None,
question_attention_mask: torch.Tensor | None = None,
max_new_tokens: int = 8,
eos_token_id: int | None = None,
workspace_override: torch.Tensor | None = None,
gate_overrides: list | None = None,
reread_inputs: dict | None = None,
reread_zero_memory: bool = False,
) -> torch.LongTensor:
"""Greedy decode against the replacement cache. ``[B, <=N_a]``.
A purpose-built loop rather than ``generate()``: the prefill is not a
token sequence the generation utilities can reconstruct, and the targets
here are short VQA answers. Correctness of the read-once contract is
easier to guarantee with the loop visible.
"""
if getattr(self.config, "perceive_deliberate_chain", False):
if (
workspace_override is not None
or gate_overrides is not None
or reread_inputs is not None
or reread_zero_memory
):
raise ValueError(
"historical workspace/reread interventions do not apply "
"to perceive_deliberate_chain"
)
return self._generate_perceive_deliberate_chain(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
max_new_tokens=max_new_tokens,
eos_token_id=eos_token_id,
)
if getattr(self.config, "spatial_visual_recurrence", False):
if (
workspace_override is not None
or gate_overrides is not None
or reread_inputs is not None
or reread_zero_memory
):
raise ValueError(
"historical workspace/reread interventions do not apply "
"to spatial_visual_recurrence"
)
return self._generate_spatial_visual_field(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
max_new_tokens=max_new_tokens,
eos_token_id=eos_token_id,
)
upper, cache, _, _ = self.prefill(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
workspace_override=workspace_override,
gate_overrides=gate_overrides,
reread_inputs=reread_inputs,
reread_zero_memory=reread_zero_memory,
)
return self._greedy_decode_from_prefill(
upper,
cache,
question_ids,
question_attention_mask,
max_new_tokens=max_new_tokens,
eos_token_id=eos_token_id,
)
def generate_answer_with_latent_policy(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
attention_mask: torch.Tensor | None = None,
question_attention_mask: torch.Tensor | None = None,
*,
policy_std: float,
max_new_tokens: int = 8,
eos_token_id: int | None = None,
sample: bool = True,
sample_mask: torch.Tensor | None = None,
) -> tuple[torch.LongTensor, LatentPolicyTrace]:
"""Sample CRR latent actions, then greedily decode a terminal answer.
The returned token ids are an environment outcome and carry no graph.
Only ``trace.log_probs`` is differentiable. Deterministic inference
continues to use :meth:`generate_answer` and therefore the policy mean.
"""
if not getattr(
self.config, "counterfactual_residual_recurrence", False
):
raise ValueError("latent-policy generation requires a CRR checkpoint")
upper, cache, _, _, trace = self.prefill(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
latent_policy_std=float(policy_std),
latent_policy_sample=bool(sample),
latent_policy_sample_mask=sample_mask,
return_latent_policy_trace=True,
)
generated = self._greedy_decode_from_prefill(
upper,
cache,
question_ids,
question_attention_mask,
max_new_tokens=max_new_tokens,
eos_token_id=eos_token_id,
)
return generated, trace
def score_gold_answer_with_latent_policy(
self,
input_ids: torch.LongTensor,
pixel_values: torch.Tensor,
image_grid_thw: torch.LongTensor,
question_ids: torch.LongTensor,
answer_ids: torch.LongTensor,
labels: torch.LongTensor,
attention_mask: torch.Tensor | None = None,
question_attention_mask: torch.Tensor | None = None,
*,
policy_std: float,
sample: bool = True,
sample_mask: torch.Tensor | None = None,
score_prefixes: bool = False,
) -> tuple[
torch.FloatTensor,
torch.FloatTensor | None,
LatentPolicyTrace,
]:
"""Sample latent actions and return frozen gold-answer scores.
``score_prefixes=False`` preserves the original terminal-only reward
path. The opt-in prefix path scores ``Z_1..Z_T`` with the same frozen
decoder and gold tokens, while keeping the Gaussian score trace as the
only differentiable path.
"""
if not getattr(
self.config, "counterfactual_residual_recurrence", False
):
raise ValueError("latent-policy scoring requires a CRR checkpoint")
upper, cache, z_states, _, trace = self.prefill(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
latent_policy_std=float(policy_std),
latent_policy_sample=bool(sample),
latent_policy_sample_mask=sample_mask,
return_latent_policy_trace=True,
)
if score_prefixes:
step_scores = self._gold_answer_logprobs_by_crr_prefix(
upper,
cache,
z_states,
question_ids,
question_attention_mask,
answer_ids,
labels,
)
score = step_scores[:, -1]
else:
score = self._gold_answer_logprob_from_prefill(
upper,
cache,
question_ids,
question_attention_mask,
answer_ids,
labels,
)
step_scores = None
return score, step_scores, trace
# -- forward -----------------------------------------------------------
def forward(
self,
input_ids: torch.LongTensor | None = None,
attention_mask: torch.Tensor | None = None,
pixel_values: torch.Tensor | None = None,
image_grid_thw: torch.LongTensor | None = None,
question_ids: torch.LongTensor | None = None,
question_attention_mask: torch.Tensor | None = None,
labels: torch.LongTensor | None = None,
answer_ids: torch.LongTensor | None = None,
past_key_values: Cache | None = None,
use_cache: bool | None = None,
return_dict: bool | None = None,
workspace_override: torch.Tensor | None = None,
latent_policy_std: float | None = None,
latent_policy_sample: bool = True,
latent_policy_sample_mask: torch.Tensor | None = None,
latent_policy_max_new_tokens: int = 8,
latent_policy_eos_token_id: int | None = None,
latent_policy_score_gold: bool = False,
latent_policy_score_prefixes: bool = False,
reliance_transition_index: int | torch.Tensor | None = None,
curriculum_readout_depth: int | torch.Tensor | None = None,
**kwargs,
) -> CloseCausalLMOutput:
"""Prefill through the workspace, then decode the answer.
Args:
input_ids: the *multimodal* sequence (image placeholders included),
used only to build ``V*``.
question_ids: the *text-only* question, used for ``Q*`` and for the
replacement cache. Required -- there is deliberately no fallback
that strips the vision span implicitly, because getting that
wrong would silently reintroduce image tokens below ``l*``.
answer_ids: teacher-forced answer tokens ``[B, N_a]``.
labels: ``[B, N_a]`` targets **already aligned** to ``answer_ids`` --
``labels[:, i]`` is the target at answer position ``i``, and no
internal shift is applied. The shift lives on the hidden-state
side: position ``i``'s prediction comes from
``[last replacement-cache position] + answer[:i]``. ``-100`` ignores.
past_key_values: not accepted yet -- see below.
return_dict: must be ``True``/``None``; tuple returns are not supported.
"""
cfg = self.config
if kwargs:
raise TypeError(f"unexpected forward kwargs: {sorted(kwargs)}")
if getattr(cfg, "perceive_deliberate_chain", False):
if past_key_values is not None:
raise NotImplementedError(
"perceive/deliberate forward constructs one fresh graph"
)
if return_dict is False:
raise ValueError(
"perceive/deliberate chain returns CloseCausalLMOutput only"
)
if workspace_override is not None:
raise ValueError(
"workspace_override does not apply to latent chain tokens"
)
if latent_policy_std is not None or latent_policy_sample_mask is not None:
raise ValueError(
"the historical CRR latent policy is incompatible with "
"the perceive/deliberate chain"
)
if reliance_transition_index is not None:
raise ValueError(
"CRR transition interventions do not apply to the chain"
)
if curriculum_readout_depth is not None:
raise ValueError(
"program-curriculum readout does not apply to the chain"
)
if input_ids is None or pixel_values is None or image_grid_thw is None:
raise ValueError(
"perceive/deliberate chain requires multimodal inputs"
)
if question_ids is None:
raise ValueError(
"perceive/deliberate chain requires clean question_ids"
)
if labels is None or answer_ids is None:
raise ValueError(
"teacher-forced chain forward requires labels and answer_ids"
)
return self._forward_perceive_deliberate_chain(
input_ids,
attention_mask,
pixel_values,
image_grid_thw,
question_ids,
question_attention_mask,
labels,
answer_ids,
)
if getattr(cfg, "spatial_visual_recurrence", False):
if past_key_values is not None:
raise NotImplementedError(
"spatial forward always constructs a fresh final-field prefix"
)
if return_dict is False:
raise ValueError(
"spatial visual recurrence returns CloseCausalLMOutput only"
)
if workspace_override is not None:
raise ValueError(
"workspace_override does not apply to the visual field"
)
if latent_policy_std is not None or latent_policy_sample_mask is not None:
raise ValueError(
"the historical CRR latent policy is incompatible with "
"spatial visual recurrence"
)
if reliance_transition_index is not None:
raise ValueError(
"transition reliance interventions are not part of SRVF"
)
if curriculum_readout_depth is not None:
raise ValueError(
"program-curriculum readout does not apply to SRVF"
)
if input_ids is None or pixel_values is None or image_grid_thw is None:
raise ValueError(
"spatial recurrence requires input_ids, pixel_values, and "
"image_grid_thw"
)
if question_ids is None:
raise ValueError(
"spatial recurrence requires text-only question_ids for "
"visual-to-question cross-attention"
)
return self._forward_spatial_visual_field(
input_ids,
attention_mask,
pixel_values,
image_grid_thw,
question_ids,
question_attention_mask,
labels,
answer_ids,
use_cache=use_cache,
)
if latent_policy_std is None and latent_policy_sample_mask is not None:
raise ValueError(
"latent_policy_sample_mask requires latent_policy_std"
)
if latent_policy_std is not None:
if workspace_override is not None:
raise ValueError(
"latent policy and workspace_override are mutually exclusive"
)
if past_key_values is not None:
raise ValueError("latent policy always builds a fresh prefix cache")
if return_dict is False:
raise ValueError("latent policy requires return_dict=True")
if input_ids is None or pixel_values is None or image_grid_thw is None:
raise ValueError(
"latent policy requires multimodal input_ids, pixel_values, "
"and image_grid_thw"
)
if question_ids is None:
raise ValueError("latent policy requires text-only question_ids")
generated = None
gold_logprob = None
step_gold_logprob = None
if latent_policy_score_gold:
if answer_ids is None or labels is None:
raise ValueError(
"latent_policy_score_gold requires answer_ids and labels"
)
gold_logprob, step_gold_logprob, trace = (
self.score_gold_answer_with_latent_policy(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
answer_ids,
labels,
attention_mask,
question_attention_mask,
policy_std=float(latent_policy_std),
sample=bool(latent_policy_sample),
sample_mask=latent_policy_sample_mask,
score_prefixes=bool(latent_policy_score_prefixes),
)
)
else:
if latent_policy_score_prefixes:
raise ValueError(
"latent_policy_score_prefixes requires "
"latent_policy_score_gold"
)
generated, trace = self.generate_answer_with_latent_policy(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
policy_std=float(latent_policy_std),
max_new_tokens=int(latent_policy_max_new_tokens),
eos_token_id=latent_policy_eos_token_id,
sample=bool(latent_policy_sample),
sample_mask=latent_policy_sample_mask,
)
return CloseCausalLMOutput(
generated_ids=generated,
latent_policy_gold_logprob=gold_logprob,
latent_policy_step_gold_logprob=step_gold_logprob,
latent_policy_log_probs=trace.log_probs,
latent_policy_mean_update_rms=trace.mean_update_rms,
latent_policy_noise_rms=trace.noise_rms,
)
compute_residual_swap = bool(
self.training
and getattr(cfg, "counterfactual_reliance_loss", False)
and labels is not None
and workspace_override is None
)
compute_step_retention = bool(
self.training
and getattr(cfg, "counterfactual_step_retention_loss", False)
and labels is not None
and workspace_override is None
)
need_answer_lower = compute_residual_swap or compute_step_retention
# Every prefill rebuilds the replacement cache from scratch. Accepting a
# caller-supplied cache would silently discard it and redo the visual
# read, so refuse instead of pretending to support incremental decode.
if past_key_values is not None:
raise NotImplementedError(
"past_key_values is not supported: forward() always rebuilds the "
"replacement cache. Incremental decoding needs prepare_inputs_for_"
"generation, which is not implemented yet."
)
if return_dict is False:
raise ValueError(
"return_dict=False is not supported; CLOSE always returns a "
"CloseCausalLMOutput (the workspace trajectory has no tuple slot)."
)
if question_ids is None:
raise ValueError(
"question_ids (text-only, no vision span) is required. Strip the "
"vision span explicitly with utils.splitting.vision_span_mask; an "
"implicit fallback risks leaking image tokens below l*."
)
if pixel_values is None or input_ids is None:
raise ValueError("input_ids and pixel_values are required to build V*.")
tm = self.text_model
n_q = question_ids.shape[1]
upper, cache, z_states, v_pooled = self.prefill(
input_ids,
pixel_values,
image_grid_thw,
question_ids,
attention_mask,
question_attention_mask,
workspace_override=workspace_override,
curriculum_readout_depth=curriculum_readout_depth,
)
q_state_interface = self._uses_question_state_interface()
# The last real prefix position predicts answer token 0. Historical
# slot paths always end in a valid slot; raw Q* batches may end in
# right-padding and therefore require a per-item gather.
if q_state_interface:
q_mask_for_seed = question_attention_mask
if q_mask_for_seed is None:
q_mask_for_seed = torch.ones_like(question_ids, dtype=torch.long)
last_q = ((q_mask_for_seed > 0).sum(dim=-1).long() - 1).clamp_min(0)
hidden_for_logits = upper[
torch.arange(upper.shape[0], device=upper.device), last_q
].unsqueeze(1)
else:
hidden_for_logits = upper[:, -1:] # [B, 1, d]
answer_lower_hidden = None
ans_positions = None
upper_mask = None
if answer_ids is not None and answer_ids.shape[1] > 1:
n_ctx = n_q + (
0 if q_state_interface else cfg.num_workspace_slots
)
n_a = answer_ids.shape[1]
ans_embeds = self.vl.get_input_embeddings()(answer_ids) # [B, N_a, d]
cache_position = torch.arange(n_ctx, n_ctx + n_a, device=ans_embeds.device)
q_mask = question_attention_mask
if q_mask is None:
q_mask = torch.ones_like(question_ids, dtype=torch.long)
if q_state_interface:
prefix_mask = q_mask
logical_prefix_lengths = (q_mask > 0).sum(dim=-1).long()
else:
prefix_mask, _, logical_prefix_lengths = _replacement_prefix_layout(
q_mask, cfg.num_workspace_slots
)
logical_answer_pos = (
logical_prefix_lengths[:, None]
+ torch.arange(n_a, device=ans_embeds.device)[None, :]
) # [B,N_a]
ans_positions = logical_answer_pos.unsqueeze(0).expand(3, -1, -1)
if labels is not None and labels.shape == answer_ids.shape:
answer_mask = labels.ne(-100).to(prefix_mask.dtype)
else:
answer_mask = torch.ones(
answer_ids.shape, dtype=prefix_mask.dtype, device=answer_ids.device
)
upper_mask = torch.cat([prefix_mask, answer_mask], dim=-1)
if q_state_interface:
lower_mask = upper_mask
else:
lower_prefix_mask = prefix_mask.clone()
lower_prefix_mask[:, -cfg.num_workspace_slots:] = 0
lower_mask = torch.cat([lower_prefix_mask, answer_mask], dim=-1)
decoded_answer = self._decode_token_block(
ans_embeds,
ans_positions,
lower_mask,
upper_mask,
cache,
cache_position,
return_lower=need_answer_lower,
) # [B,N_a,d]
if need_answer_lower:
h_ans, answer_lower_hidden = decoded_answer
else:
h_ans = decoded_answer
hidden_for_logits = torch.cat([hidden_for_logits, h_ans[:, :-1]], dim=1)
logits = self.backbone.lm_head(final_norm(tm, hidden_for_logits)) # [B, N_a, V]
# Both training-only counterfactual objectives replay only L22-L27.
# Question/answer states below that boundary are shared with the
# factual pass, so neither branch performs another visual read.
q_mask = None
q_keep = None
q_positions = None
if compute_residual_swap or compute_step_retention:
q_mask = question_attention_mask
if q_mask is None:
q_mask = torch.ones_like(question_ids, dtype=torch.long)
q_keep = q_mask.gt(0)
q_pos_1d = q_keep.long().cumsum(dim=-1) - 1
q_pos_1d = q_pos_1d.masked_fill(~q_keep, 0)
q_positions = q_pos_1d.unsqueeze(0).expand(3, -1, -1)
step_answer_nll = None
step_answer_valid = None
if compute_step_retention:
if not getattr(cfg, "counterfactual_residual_recurrence", False):
raise RuntimeError(
"step retention is implemented only for CRR"
)
if len(z_states) < 3:
raise RuntimeError(
"step retention requires at least two CRR residual steps"
)
if q_mask is None or q_keep is None or q_positions is None:
raise RuntimeError("step retention is missing question layout")
# The factual path already scores Z_T. Score Z_1..Z_{T-1} in one
# stacked upper-decoder replay to keep the extra launches small.
base_anchor = z_states[0]
partial_prefixes = []
residual_prefix: list[torch.Tensor] = []
for residual in z_states[1:-1]:
residual_prefix.append(residual)
partial_prefixes.append(
base_anchor
+ self._aggregate_crr_residuals(residual_prefix)
)
num_partial = len(partial_prefixes)
partial_prefix = torch.cat(partial_prefixes, dim=0)
partial_q_positions = q_positions.repeat(1, num_partial, 1)
partial_q_mask = q_mask.repeat(num_partial, 1)
if answer_ids is not None and answer_ids.shape[1] > 1:
if (
answer_lower_hidden is None
or ans_positions is None
or upper_mask is None
):
raise RuntimeError(
"step retention is missing teacher-forced lower states"
)
partial_input = torch.cat(
[
partial_prefix.to(answer_lower_hidden.dtype),
answer_lower_hidden.repeat(num_partial, 1, 1),
],
dim=1,
)
partial_positions = torch.cat(
[
partial_q_positions,
ans_positions.repeat(1, num_partial, 1),
],
dim=-1,
)
partial_attention_mask = upper_mask.repeat(num_partial, 1)
else:
partial_input = partial_prefix
partial_positions = partial_q_positions
partial_attention_mask = partial_q_mask
partial_ctx = make_split_context(
tm,
partial_input,
partial_positions,
partial_attention_mask,
)
partial_upper = run_layer_range(
tm,
partial_ctx,
self._upper_decoder_start(),
None,
hidden_states=partial_input,
use_cache=False,
)
last_q = (q_keep.sum(dim=-1).long() - 1).clamp_min(0)
partial_last_q = last_q.repeat(num_partial)
partial_hidden_for_logits = partial_upper[
torch.arange(
partial_upper.shape[0], device=partial_upper.device
),
partial_last_q,
].unsqueeze(1)
if answer_ids is not None and answer_ids.shape[1] > 1:
n_a = answer_ids.shape[1]
partial_hidden_for_logits = torch.cat(
[
partial_hidden_for_logits,
partial_upper[:, n_q : n_q + n_a - 1],
],
dim=1,
)
partial_logits = self.backbone.lm_head(
final_norm(tm, partial_hidden_for_logits)
)
partial_nll, partial_valid = _answer_nll_per_example(
partial_logits, labels.repeat(num_partial, 1)
)
factual_nll, factual_valid = _answer_nll_per_example(logits, labels)
batch_size = labels.shape[0]
step_answer_nll = torch.cat(
[
partial_nll.view(num_partial, batch_size),
factual_nll.unsqueeze(0),
],
dim=0,
).transpose(0, 1)
step_answer_valid = torch.cat(
[
partial_valid.view(num_partial, batch_size),
factual_valid.unsqueeze(0),
],
dim=0,
).transpose(0, 1)
residual_swap_score_gap = None
residual_swap_valid = None
residual_swap_transition = None
corrective_current_nll = None
corrective_previous_nll = None
corrective_swap_nll = None
corrective_valid = None
if compute_residual_swap:
if not getattr(cfg, "counterfactual_residual_recurrence", False):
raise RuntimeError(
"counterfactual reliance is implemented only for CRR"
)
donors, residual_swap_valid = _different_answer_residual_donors(
labels, question_attention_mask
)
base_anchor = z_states[0]
if q_mask is None or q_keep is None or q_positions is None:
raise RuntimeError("CRR residual swap is missing question layout")
corrective_transition = bool(
getattr(cfg, "corrective_transition_reliance_loss", False)
)
if corrective_transition:
if reliance_transition_index is None:
raise ValueError(
"corrective transition reliance requires "
"reliance_transition_index"
)
if isinstance(reliance_transition_index, torch.Tensor):
if reliance_transition_index.numel() != 1:
raise ValueError(
"reliance_transition_index must be a scalar"
)
transition_index = int(reliance_transition_index.item())
else:
transition_index = int(reliance_transition_index)
horizon = int(cfg.num_workspace_steps)
if not 2 <= transition_index <= horizon:
raise ValueError(
"corrective reliance transition must be in "
f"[2,{horizon}], got {transition_index}"
)
# Score the selected *local* transition through the same
# decoder used at inference. Previous history is preserved;
# only C_k is donor-swapped in the negative prefix. Unlike
# the historical tail rollout, this directly asks whether
# transition k corrects Z_{k-1} before any later step can hide
# or amplify its effect.
current_prefix = base_anchor + self._aggregate_crr_residuals(
z_states[1 : transition_index + 1]
)
previous_prefix = (
base_anchor
+ self._aggregate_crr_residuals(
z_states[1:transition_index]
)
).detach()
aligned = q_mask.gt(0) & q_mask.index_select(0, donors).gt(0)
donor_residual = z_states[transition_index].index_select(
0, donors
).detach()
donor_residual = donor_residual.masked_fill(
~aligned.unsqueeze(-1), 0
)
swapped_steps = [
state.detach()
for state in z_states[1:transition_index]
]
swapped_steps.append(donor_residual)
swapped_prefix = (
base_anchor.detach()
+ self._aggregate_crr_residuals(swapped_steps)
).detach()
prefixes = torch.cat(
[current_prefix, previous_prefix, swapped_prefix], dim=0
)
num_branches = 3
if answer_ids is not None and answer_ids.shape[1] > 1:
if (
answer_lower_hidden is None
or ans_positions is None
or upper_mask is None
):
raise RuntimeError(
"corrective transition scoring is missing "
"teacher-forced lower states"
)
replay_input = torch.cat(
[
prefixes.to(answer_lower_hidden.dtype),
answer_lower_hidden.repeat(num_branches, 1, 1),
],
dim=1,
)
replay_positions = torch.cat(
[
q_positions.repeat(1, num_branches, 1),
ans_positions.repeat(1, num_branches, 1),
],
dim=-1,
)
replay_attention_mask = upper_mask.repeat(num_branches, 1)
else:
replay_input = prefixes
replay_positions = q_positions.repeat(1, num_branches, 1)
replay_attention_mask = q_mask.repeat(num_branches, 1)
replay_ctx = make_split_context(
tm,
replay_input,
replay_positions,
replay_attention_mask,
)
replay_upper = run_layer_range(
tm,
replay_ctx,
self._upper_decoder_start(),
None,
hidden_states=replay_input,
use_cache=False,
)
batch_size = labels.shape[0]
last_q = (q_keep.sum(dim=-1).long() - 1).clamp_min(0)
replay_last_q = last_q.repeat(num_branches)
replay_hidden_for_logits = replay_upper[
torch.arange(
replay_upper.shape[0], device=replay_upper.device
),
replay_last_q,
].unsqueeze(1)
if answer_ids is not None and answer_ids.shape[1] > 1:
n_a = answer_ids.shape[1]
replay_hidden_for_logits = torch.cat(
[
replay_hidden_for_logits,
replay_upper[:, n_q : n_q + n_a - 1],
],
dim=1,
)
replay_logits = self.backbone.lm_head(
final_norm(tm, replay_hidden_for_logits)
)
replay_nll, replay_valid = _answer_nll_per_example(
replay_logits, labels.repeat(num_branches, 1)
)
replay_nll = replay_nll.view(num_branches, batch_size)
replay_valid = replay_valid.view(num_branches, batch_size)
corrective_current_nll = replay_nll[0]
corrective_previous_nll = replay_nll[1]
corrective_swap_nll = replay_nll[2]
corrective_valid = (
residual_swap_valid
& replay_valid[0]
& replay_valid[1]
& replay_valid[2]
)
residual_swap_score_gap = (
corrective_swap_nll - corrective_current_nll
)
residual_swap_valid = corrective_valid
residual_swap_transition = torch.tensor(
transition_index,
dtype=torch.long,
device=labels.device,
)
else:
if getattr(cfg, "transition_aware_reliance_loss", False):
if reliance_transition_index is None:
raise ValueError(
"transition-aware reliance requires "
"reliance_transition_index"
)
if isinstance(reliance_transition_index, torch.Tensor):
if reliance_transition_index.numel() != 1:
raise ValueError(
"reliance_transition_index must be a scalar"
)
transition_index = int(reliance_transition_index.item())
else:
transition_index = int(reliance_transition_index)
swapped_prefix = self._counterfactual_transition_rollout_prefix(
z_states,
donors,
transition_index,
q_mask,
)
residual_swap_transition = torch.tensor(
transition_index,
dtype=torch.long,
device=labels.device,
)
else:
aggregated_residual = self._aggregate_crr_residuals(
z_states[1:]
)
donor_residual = aggregated_residual.index_select(0, donors)
# Never copy right-padding as semantic residual content.
aligned = (
q_mask.gt(0)
& q_mask.index_select(0, donors).gt(0)
)
donor_residual = donor_residual.masked_fill(
~aligned.unsqueeze(-1), 0
)
swapped_prefix = base_anchor + donor_residual
if answer_ids is not None and answer_ids.shape[1] > 1:
if (
answer_lower_hidden is None
or ans_positions is None
or upper_mask is None
):
raise RuntimeError(
"CRR residual swap is missing teacher-forced lower states"
)
swapped_input = torch.cat(
[
swapped_prefix.to(answer_lower_hidden.dtype),
answer_lower_hidden,
],
dim=1,
)
swapped_positions = torch.cat(
[q_positions, ans_positions], dim=-1
)
swapped_attention_mask = upper_mask
else:
swapped_input = swapped_prefix
swapped_positions = q_positions
swapped_attention_mask = q_mask
# Only L22-L27 are replayed. Vision, the lower backbone, L21,
# and teacher-forced answer states below L22 are shared with
# the factual pass; no second multimodal read is constructed.
swapped_ctx = make_split_context(
tm,
swapped_input,
swapped_positions,
swapped_attention_mask,
)
swapped_upper = run_layer_range(
tm,
swapped_ctx,
self._upper_decoder_start(),
None,
hidden_states=swapped_input,
use_cache=False,
)
last_q = (q_keep.sum(dim=-1).long() - 1).clamp_min(0)
swapped_hidden_for_logits = swapped_upper[
torch.arange(
swapped_upper.shape[0], device=swapped_upper.device
),
last_q,
].unsqueeze(1)
if answer_ids is not None and answer_ids.shape[1] > 1:
n_a = answer_ids.shape[1]
swapped_hidden_for_logits = torch.cat(
[
swapped_hidden_for_logits,
swapped_upper[:, n_q : n_q + n_a - 1],
],
dim=1,
)
swapped_logits = self.backbone.lm_head(
final_norm(tm, swapped_hidden_for_logits)
)
factual_nll, factual_valid = _answer_nll_per_example(
logits, labels
)
swapped_nll, swapped_valid = _answer_nll_per_example(
swapped_logits, labels
)
# score = -NLL, hence factual - swap = NLL_swap - NLL_pos.
residual_swap_score_gap = swapped_nll - factual_nll
residual_swap_valid &= factual_valid & swapped_valid
loss = None
if labels is not None:
loss = answer_cross_entropy(
logits,
labels,
per_example=bool(
getattr(cfg, "per_example_answer_loss", False)
),
)
iface = getattr(self, "_iface_loss", None)
self._iface_loss = None # consume-once
functional = getattr(self, "_functional_state_loss", None)
self._functional_state_loss = None # consume-once
return CloseCausalLMOutput(
loss=loss,
logits=logits,
past_key_values=cache if use_cache else None,
workspace_states=tuple(z_states),
v_star_pooled=v_pooled,
iface_loss=iface,
functional_state_loss=functional,
residual_swap_score_gap=residual_swap_score_gap,
residual_swap_valid=residual_swap_valid,
residual_swap_transition=residual_swap_transition,
corrective_current_nll=corrective_current_nll,
corrective_previous_nll=corrective_previous_nll,
corrective_swap_nll=corrective_swap_nll,
corrective_valid=corrective_valid,
step_answer_nll=step_answer_nll,
step_answer_valid=step_answer_valid,
)
__all__ = [
"CloseQwen2_5_VLForConditionalGeneration",
"WorkspaceRead",
"WorkspaceTransition",
"RawEvidenceQuestionTransition",
"answer_cross_entropy",
"build_adapter_config",
"CloseCausalLMOutput",
"LatentPolicyTrace",
"CRRLatentPolicyAdapter",
]