Buckets:

EmergentTwilight
/

MMCoT

Files
xet
EmergentTwilight/MMCoT / OURS /src /sft /objective.py
EmergentTwilight
20 days ago
download
raw
79.1 kB
"""Text/bit CE, Soft-BSQ features, and bidirectional Information Gain."""

from __future__ import annotations

import contextlib
import hashlib
from types import SimpleNamespace
from typing import NamedTuple

import torch
import torch.nn.functional as F
from torch import nn
from shared.distributions import (
    contrast_shaped_log_target,
    event_jsd_loss,
    gt_positive_gate,
)
from shared.modeling import unwrap_ar_model
from shared.supervised import select_noisy_sample_ids

from .config import ObjectiveConfig, RELATIVE_IG_EPSILON
from .data import InterleavedBatch, Span, TransitionEvent


def information_gain_loss(
    full_score: torch.Tensor,
    blind_score: torch.Tensor,
    temperature: float,
    *,
    formulation: str = "absolute",
    factor_count: int = 1,
) -> torch.Tensor:
    """Configured blind-stop-gradient Softplus objective from the design doc."""

    if temperature <= 0:
        raise ValueError("Information Gain temperature must be positive")
    if formulation == "absolute":
        return temperature * F.softplus(
            (blind_score.detach() - full_score) / temperature
        )
    if formulation != "relative_per_unit":
        raise ValueError(f"unknown Information Gain formulation: {formulation}")
    if type(factor_count) is not int or factor_count <= 0:
        raise ValueError("relative Information Gain factor_count must be positive")
    full_unit = full_score / factor_count
    blind_unit = blind_score.detach() / factor_count
    relative_scale = (-blind_unit).clamp_min(RELATIVE_IG_EPSILON)
    relative_margin = (full_unit - blind_unit) / relative_scale
    return relative_scale * temperature * F.softplus(
        -relative_margin / temperature
    )


def _span_positions(
    spans: tuple[Span, ...], *, offset: int = 0, device=None
) -> torch.Tensor:
    values = [
        torch.arange(span.start + offset, span.end + offset, device=device)
        for span in spans
    ]
    if not values:
        raise ValueError("score/mask spans cannot be empty")
    return torch.cat(values)


def _event_key(event: TransitionEvent) -> tuple[int, int]:
    return event.sample_index, event.index


def _text_event_logits(ar_model, hidden_states, input_ids, event: TransitionEvent):
    if min(span.start for span in event.target_sequence_spans) == 0:
        raise ValueError("position zero cannot be an autoregressive text target")
    targets = _span_positions(event.target_sequence_spans, device=input_ids.device)
    carriers = targets - 1
    return ar_model.project_text(hidden_states[event.sample_index, carriers])


def _project_text_event_logits(
    ar_model,
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    locations: list[tuple[tuple[int, int], TransitionEvent, int]],
) -> dict[tuple[int, int], torch.Tensor]:
    """Project every selected text carrier with one vocabulary-head call."""

    if not locations:
        return {}
    chunks = []
    lengths = []
    for _key, event, sample_index in locations:
        if min(span.start for span in event.target_sequence_spans) == 0:
            raise ValueError("position zero cannot be an autoregressive text target")
        targets = _span_positions(
            event.target_sequence_spans, device=input_ids.device
        )
        chunks.append(hidden_states[sample_index, targets - 1])
        lengths.append(targets.numel())
    projected = ar_model.project_text(torch.cat(chunks, dim=0))
    return {
        key: logits
        for (key, _event, _sample), logits in zip(
            locations, projected.split(lengths), strict=True
        )
    }


def _text_event_score(ar_model, hidden_states, input_ids, event: TransitionEvent):
    targets = _span_positions(event.target_sequence_spans, device=input_ids.device)
    logits = _text_event_logits(ar_model, hidden_states, input_ids, event).float()
    token_ids = input_ids[event.sample_index, targets]
    return -F.cross_entropy(logits, token_ids, reduction="mean")


def _visual_shape(ar_model) -> tuple[int, int, int]:
    vision = ar_model.model.visual
    return (
        1 + len(vision.config.deepstack_visual_indexes),
        vision.spatial_merge_unit,
        vision.config.bsq_dim,
    )


def _visual_event_score(
    ar_model,
    visual_logits,
    visual_codes,
    event: TransitionEvent,
):
    if min(span.start for span in event.target_sequence_spans) == 0:
        raise ValueError("position zero cannot be an autoregressive visual target")
    sequence = _span_positions(event.target_sequence_spans, device=visual_logits.device)
    code_positions = _span_positions(
        event.target_visual_spans, device=visual_codes.device
    )
    if sequence.numel() != code_positions.numel():
        raise ValueError("visual target sequence/code lengths differ")
    return _visual_selected_score(
        ar_model,
        visual_logits[event.sample_index, sequence - 1],
        visual_codes,
        event,
    )


def _visual_event_logits(
    ar_model,
    visual_logits,
    visual_codes,
    event: TransitionEvent,
) -> tuple[torch.Tensor, torch.Tensor]:
    sequence = _span_positions(event.target_sequence_spans, device=visual_logits.device)
    return _visual_selected_logits(
        ar_model,
        visual_logits[event.sample_index, sequence - 1],
        visual_codes,
        event,
    )


def _visual_selected_score(
    ar_model,
    selected_logits,
    visual_codes,
    event: TransitionEvent,
):
    """Score one visual event from logits already gathered at its carriers."""

    logits, targets = _visual_selected_logits(
        ar_model, selected_logits, visual_codes, event
    )
    logits = logits.float()
    bit_logp = (
        F.log_softmax(logits, dim=-1).gather(-1, targets.unsqueeze(-1)).squeeze(-1)
    )
    # One AR visual position is a 1,024-factor composite token.  Sum its
    # Bernoulli log probabilities and normalize only over AR steps J.
    return bit_logp.flatten(1).sum(-1).mean()


def _visual_selected_logits(
    ar_model,
    selected_logits,
    visual_codes,
    event: TransitionEvent,
) -> tuple[torch.Tensor, torch.Tensor]:
    code_positions = _span_positions(
        event.target_visual_spans, device=visual_codes.device
    )
    levels, spatial, bits = _visual_shape(ar_model)
    logits = selected_logits.reshape(-1, levels, spatial, bits, 2)
    targets = visual_codes[code_positions].to(torch.long)
    if logits.shape[:-1] != targets.shape:
        raise ValueError("selected visual logits and target codes differ")
    return logits, targets


def soft_bsq_feature_terms(
    ar_model,
    visual_logits: torch.Tensor,
    target_visual_features,
    generated_flat_mask: torch.Tensor,
    *,
    cosine_weight: float,
    layer_norm_epsilon: float,
    cosine_epsilon: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
    """Return summed LN-MSE, cosine, combined loss, and ``N_V * L``.
    The predicted path uses the probability of bit one as a continuous BSQ
    code and remains differentiable through UniAR's frozen level-specific
    output projections and mergers. Ground-truth features come from the one
    hard-BSQ replay already used to build the teacher-forced model input.
    """

    if visual_logits.ndim != 5 or visual_logits.shape[-1] != 2:
        raise ValueError("Soft-BSQ logits must have shape [N, L, U, D, 2]")
    bits = visual_logits.shape[-2]
    probabilities = visual_logits.softmax(-1)[..., 1]
    soft_values = (2 * probabilities - 1) / bits**0.5
    vision = ar_model.model.visual
    predicted_main, predicted_deepstack = vision.bsq_values_to_features(
        soft_values, with_merger=True
    )
    predicted = [predicted_main, *predicted_deepstack]
    target = [
        target_visual_features.image_embeds[generated_flat_mask],
        *[
            value[generated_flat_mask]
            for value in target_visual_features.deepstack_visual_embeds
        ],
    ]
    if len(predicted) != len(target) or len(predicted) != visual_logits.shape[1]:
        raise ValueError("Soft-BSQ feature levels differ from visual logits")

    mse_sum = visual_logits.new_zeros((), dtype=torch.float32)
    cosine_sum = visual_logits.new_zeros((), dtype=torch.float32)
    items = 0
    for predicted_level, target_level in zip(predicted, target):
        predicted_level = predicted_level.float()
        target_level = target_level.detach().to(predicted_level).float()
        if predicted_level.shape != target_level.shape:
            raise ValueError(
                "Soft-BSQ predicted/target feature shapes differ: "
                f"{predicted_level.shape} != {target_level.shape}"
            )
        dimension = predicted_level.shape[-1]
        predicted_normalized = F.layer_norm(
            predicted_level, (dimension,), eps=layer_norm_epsilon
        )
        target_normalized = F.layer_norm(
            target_level, (dimension,), eps=layer_norm_epsilon
        )
        mse_sum = mse_sum + (predicted_normalized - target_normalized).square().mean(
            -1
        ).sum()
        cosine_sum = cosine_sum + (
            1
            - F.cosine_similarity(
                predicted_level,
                target_level,
                dim=-1,
                eps=cosine_epsilon,
            )
        ).sum()
        items += predicted_level.shape[0]
    combined = mse_sum + cosine_weight * cosine_sum
    return mse_sum, cosine_sum, combined, items


class ObjectiveOutput(NamedTuple):
    """DDP-visible outputs, including its static-graph backward sink.
    An unregistered dataclass is opaque to PyTorch's output pytree traversal:
    with static_graph=True that silently bypasses the first backward reduction.
    NamedTuple preserves attribute access and exposes the loss tensor to DDP.
    """

    loss: torch.Tensor
    statistics: dict[str, torch.Tensor]
    events: list[dict]


def objective_counts(
    batch: InterleavedBatch,
    selected_events: list[list[TransitionEvent]],
    *,
    image_gen_id: int,
    marker_weight: float,
) -> dict[str, float]:
    """Cheap denominators used to normalize across ranks and accumulation."""

    targets = batch.labels[:, 1:]
    supervised = targets != -100
    marker = (targets == image_gen_id) & supervised
    text_weight = float(supervised.sum()) + (marker_weight - 1.0) * float(marker.sum())
    bits_per_position = math_prod(batch.visual_codes.shape[1:])
    result = {
        "text": text_weight,
        "bit": float(batch.generated_flat_mask.sum()) * bits_per_position,
        "soft_bsq_feature": float(batch.generated_flat_mask.sum())
        * int(batch.visual_codes.shape[1]),
        "ig_text_to_visual": 0.0,
        "ig_visual_to_text": 0.0,
        "ema_text_to_visual": 0.0,
        "ema_visual_to_text": 0.0,
    }
    for events in selected_events:
        # IG first averages events within one trajectory, so the global
        # denominator is the number of contributing trajectories rather than
        # the number of selected events.
        for direction in {event.direction for event in events}:
            result[f"ig_{direction}"] += 1.0
            result[f"ema_{direction}"] += 1.0
    return result


def math_prod(values) -> int:
    result = 1
    for value in values:
        result *= int(value)
    return result


def _sum_trajectory_means(
    values: list[list[torch.Tensor]], zero: torch.Tensor
) -> tuple[torch.Tensor, int]:
    """Sum per-trajectory event means and count contributing trajectories."""

    means = [torch.stack(trajectory).mean() for trajectory in values if trajectory]
    if not means:
        return zero, 0
    return torch.stack(means).sum(), len(means)


def _materialize_event_records(records: list[dict]) -> list[dict]:
    """Transfer every scalar event diagnostic to the host in one synchronization."""

    locations = []
    values = []
    for record in records:
        for key, value in record.items():
            if torch.is_tensor(value):
                if value.numel() != 1:
                    raise ValueError(f"event diagnostic {key} is not scalar")
                locations.append((record, key))
                values.append(value.detach().float().reshape(()))
    if values:
        host_values = torch.stack(values).cpu().tolist()
        for (record, key), value in zip(locations, host_values, strict=True):
            record[key] = float(value)
    return records


def bit_flip_visual_context(
    batch: InterleavedBatch,
    visual_codes: torch.Tensor,
    generated_flat_mask: torch.Tensor,
    *,
    noisy_sample_ids: frozenset[str],
    flip_probability: float,
    probability_mode: str = "fixed",
    generator: torch.Generator | None = None,
    corruption_seed: int | None = None,
    corruption_update: int = 0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Corrupt generated-state inputs while leaving observations and GT intact.
    Returns input codes, a per-sample noisy mask, noisy generated positions,
    and the exact per-bit flip mask.  ``visual_codes`` itself is never mutated.
    """

    if not 0.0 <= flip_probability <= 1.0:
        raise ValueError("flip_probability must be in [0, 1]")
    if probability_mode not in {"fixed", "uniform"}:
        raise ValueError("probability_mode must be fixed or uniform")
    if len(batch.visual_code_spans) != len(batch.sample_ids):
        raise ValueError("visual_code_spans must contain one span per sample")
    if generated_flat_mask.shape != visual_codes.shape[:1]:
        raise ValueError("generated_flat_mask and visual_codes differ")

    device = visual_codes.device
    noisy_flags = [sample_id in noisy_sample_ids for sample_id in batch.sample_ids]
    noisy_samples = torch.tensor(
        noisy_flags,
        device=device,
        dtype=torch.bool,
    )
    noisy_generated = torch.zeros_like(generated_flat_mask, dtype=torch.bool)
    for is_noisy, span in zip(noisy_flags, batch.visual_code_spans):
        if span.start < 0 or span.end > visual_codes.shape[0]:
            raise ValueError("visual code span falls outside the flattened codes")
        if is_noisy:
            noisy_generated[span.start : span.end] = generated_flat_mask[
                span.start : span.end
            ]

    flips = torch.zeros_like(visual_codes, dtype=torch.bool)
    if generator is not None and corruption_seed is not None:
        raise ValueError("provide generator or deterministic corruption seed, not both")
    if flip_probability and any(noisy_flags):
        if probability_mode == "fixed" and corruption_seed is None:
            flips = torch.rand(
                visual_codes.shape,
                device=device,
                generator=generator,
            ) < flip_probability
        elif probability_mode == "fixed":
            for sample_id, is_noisy, span in zip(
                batch.sample_ids, noisy_flags, batch.visual_code_spans, strict=True
            ):
                if not is_noisy:
                    continue
                seed = int.from_bytes(
                    hashlib.sha256(
                        f"{corruption_seed}\0{corruption_update}\0{sample_id}".encode()
                    ).digest()[:8],
                    "big",
                ) & ((1 << 63) - 1)
                sample_generator = torch.Generator(device=device).manual_seed(seed)
                flips[span.start : span.end] = torch.rand(
                    visual_codes[span.start : span.end].shape,
                    device=device,
                    generator=sample_generator,
                ) < flip_probability
        else:
            probabilities = sample_bit_flip_probabilities(
                batch.sample_ids,
                noisy_sample_ids=noisy_sample_ids,
                maximum_probability=flip_probability,
                device=device,
                generator=generator,
                corruption_seed=corruption_seed,
                corruption_update=corruption_update,
            )
            for sample_id, is_noisy, span, probability in zip(
                batch.sample_ids,
                noisy_flags,
                batch.visual_code_spans,
                probabilities,
                strict=True,
            ):
                if not is_noisy:
                    continue
                mask_generator = generator
                if corruption_seed is not None:
                    mask_seed = int.from_bytes(
                        hashlib.sha256(
                            (
                                f"{corruption_seed}\0{corruption_update}\0"
                                f"{sample_id}\0mask"
                            ).encode()
                        ).digest()[:8],
                        "big",
                    ) & ((1 << 63) - 1)
                    mask_generator = torch.Generator(device=device).manual_seed(
                        mask_seed
                    )
                flips[span.start : span.end] = torch.rand(
                    visual_codes[span.start : span.end].shape,
                    device=device,
                    generator=mask_generator,
                ) < probability
        flips &= noisy_generated[(slice(None),) + (None,) * (visual_codes.ndim - 1)]
        input_codes = visual_codes ^ flips.to(dtype=visual_codes.dtype)
    else:
        input_codes = visual_codes
    return input_codes, noisy_samples, noisy_generated, flips


def sample_bit_flip_probabilities(
    sample_ids: list[str],
    *,
    noisy_sample_ids: frozenset[str],
    maximum_probability: float,
    device: torch.device,
    generator: torch.Generator | None = None,
    corruption_seed: int | None = None,
    corruption_update: int = 0,
) -> torch.Tensor:
    """Sample trajectory/update BitFlip rates reproducibly from ``[0, p_max]``."""

    if not 0.0 <= maximum_probability <= 1.0:
        raise ValueError("maximum_probability must be in [0, 1]")
    if generator is not None and corruption_seed is not None:
        raise ValueError("provide generator or deterministic corruption seed, not both")
    probabilities = torch.zeros(len(sample_ids), device=device, dtype=torch.float32)
    if not maximum_probability:
        return probabilities
    for index, sample_id in enumerate(sample_ids):
        if sample_id not in noisy_sample_ids:
            continue
        sample_generator = generator
        if corruption_seed is not None:
            probability_seed = int.from_bytes(
                hashlib.sha256(
                    (
                        f"{corruption_seed}\0{corruption_update}\0"
                        f"{sample_id}\0probability"
                    ).encode()
                ).digest()[:8],
                "big",
            ) & ((1 << 63) - 1)
            sample_generator = torch.Generator(device=device).manual_seed(
                probability_seed
            )
        probabilities[index] = (
            torch.rand((), device=device, generator=sample_generator)
            * maximum_probability
        )
    return probabilities


@contextlib.contextmanager
def _detached_evaluation(model, *, use_reentrant: bool):
    """Run detached blind forwards without checkpoint recomputation frames.
    Non-reentrant activation checkpointing assumes that every checkpointed
    forward belongs to the graph that will be recomputed. IG deliberately runs
    additional no-grad forwards before the differentiable full branch. Disable
    checkpoint wrapping directly instead of switching the whole PEFT model to
    eval and back. IG configuration enforces zero LoRA/backbone dropout, so
    blind and full branches use the same training-mode distribution.
    """

    checkpointing = bool(getattr(model, "is_gradient_checkpointing", False))
    if checkpointing:
        model.gradient_checkpointing_disable()
    try:
        yield
    finally:
        if checkpointing:
            model.gradient_checkpointing_enable(
                gradient_checkpointing_kwargs={"use_reentrant": use_reentrant}
            )
        # The frozen BSQ encoder remains a feature extractor in both modes.
        visual = model.model.visual
        if hasattr(visual, "eval"):
            visual.eval()


@contextlib.contextmanager
def _uncached_autocast(device_type: str):
    """Keep detached casts out of the surrounding autocast weight cache.
    Blind output heads run first under ``no_grad``.  If their FP32 weights are
    cached by an outer BF16 autocast context, the differentiable full branch
    reuses detached casts and those trainable heads receive no gradients.
    """

    with torch.autocast(
        device_type=device_type,
        dtype=torch.get_autocast_dtype(device_type),
        enabled=torch.is_autocast_enabled(device_type),
        cache_enabled=False,
    ):
        yield


class InterleavedSFTObjective(nn.Module):
    """Loss wrapper suitable for DDP and arbitrary custom-loss extensions."""

    def __init__(
        self,
        model,
        config: ObjectiveConfig,
        image_gen_id: int,
        *,
        noisy_sample_ids: frozenset[str] | None = None,
        ema_teacher=None,
    ):
        super().__init__()
        self.model = model
        self.config = config
        self.image_gen_id = image_gen_id
        if config.bit.input_flip_probability and noisy_sample_ids is None:
            raise ValueError(
                "bit-flip training requires an explicit sample-level partition"
            )
        self.noisy_sample_ids = noisy_sample_ids or frozenset()
        self.ema_teacher = ema_teacher

    @staticmethod
    def _event_visual_output_mask(attention_mask, selected_events):
        mask = torch.zeros_like(attention_mask, dtype=torch.bool)
        found = False
        for events in selected_events:
            for event in events:
                if event.direction != "text_to_visual":
                    continue
                targets = _span_positions(
                    event.target_sequence_spans, device=attention_mask.device
                )
                mask[event.sample_index, targets - 1] = True
                found = True
        return mask if found else None

    @staticmethod
    def _concatenate_visual_features(*features):
        if not features or any(value is None for value in features):
            return None
        first = features[0]
        return type(first)(
            torch.cat([value.image_embeds for value in features], dim=0),
            tuple(
                torch.cat(values, dim=0)
                for values in zip(
                    *(value.deepstack_visual_embeds for value in features),
                    strict=True,
                )
            ),
        )

    @staticmethod
    def _event_logits_from_output(
        ar_model,
        output,
        *,
        input_ids,
        visual_codes,
        selected_events,
        visual_output_mask=None,
    ) -> dict[tuple[int, int], torch.Tensor]:
        events = [event for values in selected_events for event in values]
        text_locations = [
            (_event_key(event), event, event.sample_index)
            for event in events
            if event.direction == "visual_to_text"
        ]
        result = _project_text_event_logits(
            ar_model, output.hidden_states, input_ids, text_locations
        )
        visual_output_index = None
        if visual_output_mask is not None:
            visual_output_index = torch.full_like(
                visual_output_mask, -1, dtype=torch.long
            )
            visual_output_index[visual_output_mask] = torch.arange(
                sum(
                    event.target_steps
                    for events in selected_events
                    for event in events
                    if event.direction == "text_to_visual"
                ),
                device=visual_output_mask.device,
            )
        for event in events:
            if event.direction != "text_to_visual":
                continue
            if visual_output_index is not None:
                targets = _span_positions(
                    event.target_sequence_spans, device=input_ids.device
                )
                indices = visual_output_index[event.sample_index, targets - 1]
                logits, _ = _visual_selected_logits(
                    ar_model, output.visual_logits[indices], visual_codes, event
                )
            else:
                logits, _ = _visual_event_logits(
                    ar_model, output.visual_logits, visual_codes, event
                )
            result[_event_key(event)] = logits
        return result

    def _full_event_logits(
        self,
        ar_model,
        *,
        input_ids,
        attention_mask,
        image_grid_thw,
        visual_codes,
        position_ids,
        selected_events,
        visual_features=None,
        detach: bool = True,
    ) -> dict[tuple[int, int], torch.Tensor]:
        visual_events = any(
            event.direction == "text_to_visual"
            for events in selected_events
            for event in events
        )
        visual_output_mask = (
            self._event_visual_output_mask(attention_mask, selected_events)
            if visual_features is not None
            else None
        )
        kwargs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "image_grid_thw": image_grid_thw,
            "visual_codes": visual_codes,
            "position_ids": position_ids,
            "compute_visual_logits": visual_events,
        }
        if visual_features is not None:
            kwargs["visual_features"] = visual_features
        if visual_output_mask is not None:
            kwargs["visual_output_mask"] = visual_output_mask
        output = ar_model.forward_interleaved(**kwargs)
        result = self._event_logits_from_output(
            ar_model,
            output,
            input_ids=input_ids,
            visual_codes=visual_codes,
            selected_events=selected_events,
            visual_output_mask=visual_output_mask,
        )
        return (
            {key: value.detach() for key, value in result.items()}
            if detach
            else result
        )

    @staticmethod
    def _blind_view_slots(
        selected_events: list[list[TransitionEvent]],
    ) -> list[list[TransitionEvent]]:
        """Return the canonical direction/slot views in their historical order."""

        views = []
        for direction in ("text_to_visual", "visual_to_text"):
            per_sample = [
                [event for event in events if event.direction == direction]
                for events in selected_events
            ]
            slots = max((len(events) for events in per_sample), default=0)
            for slot in range(slots):
                active = [events[slot] for events in per_sample if slot < len(events)]
                if active:
                    views.append(active)
        return views

    @staticmethod
    def _validate_event_layout(
        batch: InterleavedBatch,
        selected_events: list[list[TransitionEvent]],
    ) -> None:
        if len(selected_events) != len(batch.sample_ids):
            raise ValueError("selected events must contain one list per trajectory")
        sequence_lengths = batch.attention_mask.sum(-1).tolist()
        for sample_index, events in enumerate(selected_events):
            visual_span = batch.visual_code_spans[sample_index]
            for event in events:
                if event.sample_index != sample_index:
                    raise ValueError("transition sample index is inconsistent")
                for span in (
                    *event.source_sequence_spans,
                    *event.target_sequence_spans,
                ):
                    if span.end > sequence_lengths[sample_index]:
                        raise ValueError("transition span overlaps sequence padding")
                for span in event.target_visual_spans:
                    if span.start < visual_span.start or span.end > visual_span.end:
                        raise ValueError("transition visual target belongs to another row")

    @staticmethod
    def _repeat_visual_features(visual_features, repeats: int):
        if visual_features is None:
            return None

        def repeated(value: torch.Tensor) -> torch.Tensor:
            return value.repeat((repeats,) + (1,) * (value.ndim - 1))

        return type(visual_features)(
            repeated(visual_features.image_embeds),
            tuple(repeated(value) for value in visual_features.deepstack_visual_embeds),
        )

    @torch.no_grad()
    def _counterfactual_logits(
        self,
        ar_model,
        *,
        input_ids,
        full_attention_mask,
        image_grid_thw,
        visual_codes,
        position_ids,
        selected_events,
        visual_features=None,
        apply_source_mask: bool = True,
        score_only: bool = False,
        transform=None,
    ) -> dict[tuple[int, int], torch.Tensor]:
        """Evaluate distinct blind views in bounded virtual-batch chunks."""

        result = {}
        batch_size, sequence_length = input_ids.shape
        views = self._blind_view_slots(selected_events)
        chunk_size = self.config.information_gain.blind_view_chunk_size
        for start in range(0, len(views), chunk_size):
            chunk = views[start : start + chunk_size]
            view_masks = []
            locations: list[tuple[tuple[int, int], TransitionEvent, int]] = []
            for view_index, events in enumerate(chunk):
                attention_mask = full_attention_mask.clone()
                for event in events:
                    if event.sample_index < 0 or event.sample_index >= batch_size:
                        raise ValueError("transition sample index is inconsistent")
                    if apply_source_mask:
                        for span in event.source_sequence_spans:
                            if span.end > sequence_length:
                                raise ValueError("source span falls outside the sequence")
                            attention_mask[
                                event.sample_index, span.start : span.end
                            ] = 0
                    locations.append(
                        (
                            _event_key(event),
                            event,
                            view_index * batch_size + event.sample_index,
                        )
                    )
                view_masks.append(attention_mask)

            repeats = len(chunk)
            expanded_attention_mask = torch.cat(view_masks, dim=0)
            expanded_input_ids = input_ids.repeat(repeats, 1)
            expanded_grid = image_grid_thw.repeat(repeats, 1)
            expanded_position_ids = position_ids.repeat(1, repeats, 1)
            expanded_features = self._repeat_visual_features(visual_features, repeats)
            expanded_codes = (
                visual_codes
                if expanded_features is not None
                else visual_codes.repeat(
                    (repeats,) + (1,) * (visual_codes.ndim - 1)
                )
            )
            visual_locations = [
                location
                for location in locations
                if location[1].direction == "text_to_visual"
            ]
            visual_output_mask = None
            visual_output_index = None
            if visual_locations:
                visual_output_mask = torch.zeros_like(
                    expanded_attention_mask, dtype=torch.bool
                )
                for _key, event, sample_index in visual_locations:
                    targets = _span_positions(
                        event.target_sequence_spans,
                        device=expanded_attention_mask.device,
                    )
                    visual_output_mask[sample_index, targets - 1] = True
                visual_output_index = torch.full_like(
                    expanded_attention_mask, -1, dtype=torch.long
                )
                visual_output_index[visual_output_mask] = torch.arange(
                    sum(event.target_steps for _key, event, _sample in visual_locations),
                    device=expanded_attention_mask.device,
                )

            kwargs = {
                "input_ids": expanded_input_ids,
                "attention_mask": expanded_attention_mask,
                "image_grid_thw": expanded_grid,
                "visual_codes": expanded_codes,
                "position_ids": expanded_position_ids,
                "compute_visual_logits": bool(visual_locations),
            }
            if expanded_features is not None:
                kwargs["visual_features"] = expanded_features
            if visual_output_mask is not None:
                kwargs["visual_output_mask"] = visual_output_mask
            output = ar_model.forward_interleaved(**kwargs)

            text_locations = [
                location
                for location in locations
                if location[1].direction == "visual_to_text"
            ]
            text_logits = _project_text_event_logits(
                ar_model, output.hidden_states, expanded_input_ids, text_locations
            )
            if score_only:
                for key, event, _sample_index in text_locations:
                    targets = _span_positions(
                        event.target_sequence_spans, device=input_ids.device
                    )
                    token_ids = input_ids[event.sample_index, targets]
                    result[key] = -F.cross_entropy(
                        text_logits[key].float(), token_ids, reduction="mean"
                    )
            else:
                for key, event, _sample_index in text_locations:
                    result[key] = (
                        transform(key, event, text_logits[key])
                        if transform is not None
                        else text_logits[key]
                    )
            for key, event, sample_index in visual_locations:
                targets = _span_positions(
                    event.target_sequence_spans, device=expanded_attention_mask.device
                )
                indices = visual_output_index[sample_index, targets - 1]
                logits, _ = _visual_selected_logits(
                    ar_model, output.visual_logits[indices], visual_codes, event
                )
                value = (
                    _visual_selected_score(
                        ar_model,
                        logits.reshape(logits.shape[0], -1),
                        visual_codes,
                        event,
                    )
                    if score_only
                    else logits
                )
                result[key] = (
                    transform(key, event, value)
                    if transform is not None and not score_only
                    else value
                )
        return {
            key: value.detach() if torch.is_tensor(value) else value
            for key, value in result.items()
        }

    def _teacher_blind_logits(
        self,
        ar_model,
        *,
        input_ids,
        full_attention_mask,
        image_grid_thw,
        visual_codes,
        position_ids,
        selected_events,
        visual_features=None,
        transform=None,
    ) -> dict[tuple[int, int], torch.Tensor]:
        return self._counterfactual_logits(
            ar_model,
            input_ids=input_ids,
            full_attention_mask=full_attention_mask,
            image_grid_thw=image_grid_thw,
            visual_codes=visual_codes,
            position_ids=position_ids,
            selected_events=selected_events,
            visual_features=visual_features,
            transform=transform,
        )

    def _blind_scores(
        self,
        ar_model,
        *,
        input_ids,
        full_attention_mask,
        image_grid_thw,
        visual_codes,
        position_ids,
        selected_events,
        visual_features=None,
        apply_source_mask: bool = True,
    ) -> dict[tuple[int, int], torch.Tensor]:
        return self._counterfactual_logits(
            ar_model,
            input_ids=input_ids,
            full_attention_mask=full_attention_mask,
            image_grid_thw=image_grid_thw,
            visual_codes=visual_codes,
            position_ids=position_ids,
            selected_events=selected_events,
            visual_features=visual_features,
            apply_source_mask=apply_source_mask,
            score_only=True,
        )

    def forward(
        self,
        batch: InterleavedBatch,
        selected_events: list[list[TransitionEvent]],
        *,
        global_denominators: dict[str, float] | None = None,
        ddp_world_size: int = 1,
        verify_no_mask: bool = False,
        record_events: bool = True,
        bit_flip_seed: int | None = None,
        bit_flip_update: int = 0,
    ) -> ObjectiveOutput:
        ar_model = unwrap_ar_model(self.model)
        device = next(ar_model.parameters()).device
        ig = self.config.information_gain
        ema = self.config.ema_distillation
        if ema.enabled and self.ema_teacher is None:
            raise ValueError("enabled EMA distillation requires an EMA teacher")
        if not (ig.enabled or ema.enabled):
            selected_events = [[] for _ in batch.sample_ids]
        self._validate_event_layout(batch, selected_events)
        input_ids = batch.input_ids.to(device)
        attention_mask = batch.attention_mask.to(device)
        labels = batch.labels.to(device)
        image_grid_thw = batch.image_grid_thw.to(device)
        generated_sequence_mask = batch.generated_sequence_mask.to(device)
        generated_flat_mask = batch.generated_flat_mask.to(device)
        visual_codes = batch.visual_codes.to(device=device, dtype=torch.long)
        if (ig.enabled or ema.enabled) and getattr(
            ar_model.config.text_config, "attention_dropout", 0.0
        ):
            raise ValueError("paired contrast objectives require zero attention dropout")
        blind_views = len(self._blind_view_slots(selected_events))
        blind_chunks = (
            (blind_views + ig.blind_view_chunk_size - 1)
            // ig.blind_view_chunk_size
            if blind_views
            else 0
        )
        model_forward_calls = 1
        counterfactual_forward_calls = 0
        if ema.enabled and blind_views:
            model_forward_calls += 1 + blind_chunks
            counterfactual_forward_calls += blind_chunks
        if ig.enabled and blind_views:
            ig_blind_calls = blind_chunks * (1 + int(verify_no_mask))
            model_forward_calls += ig_blind_calls
            counterfactual_forward_calls += ig_blind_calls

        input_codes, noisy_samples, noisy_generated, flips = bit_flip_visual_context(
            batch,
            visual_codes,
            generated_flat_mask,
            noisy_sample_ids=self.noisy_sample_ids,
            flip_probability=self.config.bit.input_flip_probability,
            probability_mode=self.config.bit.probability_mode,
            corruption_seed=bit_flip_seed,
            corruption_update=bit_flip_update,
        )
        has_noisy_context = bool(
            self.config.bit.input_flip_probability
            and any(sample_id in self.noisy_sample_ids for sample_id in batch.sample_ids)
        )
        with torch.no_grad():
            image_tokens = int(
                (batch.input_ids == ar_model.config.image_token_id).sum()
            )
            if image_tokens != batch.visual_codes.shape[0]:
                raise ValueError(
                    "visual placeholder/code mismatch: "
                    f"{image_tokens} placeholders != {batch.visual_codes.shape[0]} codes"
                )
            position_ids, _ = ar_model.model.get_rope_index(
                input_ids,
                image_grid_thw=image_grid_thw,
                attention_mask=attention_mask,
            )
            input_visual_features = (
                ar_model.prepare_visual_features(input_codes)
                if hasattr(ar_model, "prepare_visual_features")
                else None
            )
            clean_visual_features = input_visual_features
            if has_noisy_context and (
                ig.enabled or ema.enabled or self.config.soft_bsq_feature.enabled
            ) and hasattr(ar_model, "prepare_visual_features"):
                clean_visual_features = ar_model.prepare_visual_features(visual_codes)
            target_visual_features = clean_visual_features

        ema_teacher_data = {}
        if ema.enabled and any(selected_events):
            with self.ema_teacher.apply(self.model) as teacher_model:
                teacher_ar = unwrap_ar_model(teacher_model)
                with _uncached_autocast(device.type), _detached_evaluation(
                    teacher_ar,
                    use_reentrant=getattr(
                        teacher_ar.config,
                        "gradient_checkpointing_use_reentrant",
                        True,
                    ),
                ):
                    ema_full_logits = self._full_event_logits(
                        teacher_ar,
                        input_ids=input_ids,
                        attention_mask=attention_mask,
                        image_grid_thw=image_grid_thw,
                        visual_codes=visual_codes,
                        position_ids=position_ids,
                        selected_events=selected_events,
                        visual_features=clean_visual_features,
                    )

                    def shape_teacher_target(key, event, blind_logits):
                        direction_config = (
                            ema.text_to_visual
                            if event.direction == "text_to_visual"
                            else ema.visual_to_text
                        )
                        target_logp, ratio, full_logp, blind_logp = (
                            contrast_shaped_log_target(
                                ema_full_logits.pop(key),
                                blind_logits,
                                contrast_alpha=direction_config.contrast_alpha,
                                ratio_clip=ema.ratio_clip,
                            )
                        )
                        if event.direction == "text_to_visual":
                            positions = _span_positions(
                                event.target_visual_spans, device=device
                            )
                            targets = visual_codes[positions].to(torch.long)
                        else:
                            positions = _span_positions(
                                event.target_sequence_spans, device=device
                            )
                            targets = input_ids[event.sample_index, positions]
                        gt_ratio = ratio.gather(
                            -1, targets.unsqueeze(-1)
                        ).squeeze(-1)
                        position_gate = gt_positive_gate(
                            full_logp,
                            blind_logp,
                            targets,
                            delta=ema.gt_positive_delta,
                        )
                        full_gt = full_logp.gather(
                            -1, targets.unsqueeze(-1)
                        ).squeeze(-1)
                        blind_gt = blind_logp.gather(
                            -1, targets.unsqueeze(-1)
                        ).squeeze(-1)
                        if event.direction == "text_to_visual":
                            full_score = full_gt.flatten(1).sum(-1).mean()
                            blind_score = blind_gt.flatten(1).sum(-1).mean()
                        else:
                            full_score = full_gt.mean()
                            blind_score = blind_gt.mean()
                        return (
                            target_logp,
                            gt_ratio.float().mean().detach(),
                            full_score.detach(),
                            blind_score.detach(),
                            (ratio > 0).float().mean().detach(),
                            (ratio > ema.ratio_clip).float().mean().detach(),
                            position_gate.detach(),
                        )

                    ema_teacher_data = self._teacher_blind_logits(
                        teacher_ar,
                        input_ids=input_ids,
                        full_attention_mask=attention_mask,
                        image_grid_thw=image_grid_thw,
                        visual_codes=visual_codes,
                        position_ids=position_ids,
                        selected_events=selected_events,
                        visual_features=clean_visual_features,
                        transform=shape_teacher_target,
                    )
                    if ema_full_logits:
                        raise RuntimeError("EMA full/blind event sets differ")

        blind_scores: dict[tuple[int, int], torch.Tensor] = {}
        no_mask_scores: dict[tuple[int, int], torch.Tensor] = {}
        if ig.enabled and any(selected_events):
            # The trainer owns one outer autocast context for the whole
            # objective.  Detached blind passes must not populate its weight
            # cache with no-grad casts of trainable output projections.
            with _uncached_autocast(device.type), _detached_evaluation(
                ar_model,
                use_reentrant=getattr(
                    ar_model.config, "gradient_checkpointing_use_reentrant", True
                ),
            ):
                blind_scores = self._blind_scores(
                    ar_model,
                    input_ids=input_ids,
                    full_attention_mask=attention_mask,
                    image_grid_thw=image_grid_thw,
                    visual_codes=visual_codes,
                    position_ids=position_ids,
                    selected_events=selected_events,
                    visual_features=clean_visual_features,
                )
                if verify_no_mask:
                    no_mask_scores = self._blind_scores(
                        ar_model,
                        input_ids=input_ids,
                        full_attention_mask=attention_mask,
                        image_grid_thw=image_grid_thw,
                        visual_codes=visual_codes,
                        position_ids=position_ids,
                        selected_events=selected_events,
                        visual_features=clean_visual_features,
                        apply_source_mask=False,
                    )

        has_selected_events = any(selected_events)
        needs_clean_event_logits = bool(
            (ig.enabled or ema.enabled) and has_noisy_context and has_selected_events
        )
        needs_noisy_ig_scores = (
            ig.enabled and has_selected_events and not needs_clean_event_logits
        )
        needs_noisy_ema_logits = (
            ema.enabled and has_selected_events and not needs_clean_event_logits
        )
        fuse_clean_event_logits = bool(
            needs_clean_event_logits
            and ig.enabled
            and ig.student_full_view_implementation == "fused"
            and input_visual_features is not None
            and clean_visual_features is not None
        )
        if needs_clean_event_logits and not fuse_clean_event_logits:
            model_forward_calls += 1

        visual_predict_mask = torch.zeros_like(generated_sequence_mask)
        visual_predict_mask[:, :-1] = generated_sequence_mask[:, 1:]
        full_forward_kwargs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "image_grid_thw": image_grid_thw,
            "visual_codes": input_codes,
            "position_ids": position_ids,
            "compute_visual_logits": True,
        }
        if input_visual_features is not None:
            full_forward_kwargs["visual_features"] = input_visual_features
            full_forward_kwargs["visual_output_mask"] = visual_predict_mask
        clean_event_logits: dict[tuple[int, int], torch.Tensor] = {}
        if fuse_clean_event_logits:
            clean_visual_output_mask = self._event_visual_output_mask(
                attention_mask, selected_events
            )
            clean_visual_mask = (
                clean_visual_output_mask
                if clean_visual_output_mask is not None
                else torch.zeros_like(visual_predict_mask)
            )
            fused_output = ar_model.forward_interleaved(
                input_ids=input_ids.repeat(2, 1),
                attention_mask=attention_mask.repeat(2, 1),
                image_grid_thw=image_grid_thw.repeat(2, 1),
                visual_codes=torch.cat([input_codes, visual_codes], dim=0),
                position_ids=position_ids.repeat(1, 2, 1),
                compute_visual_logits=True,
                visual_features=self._concatenate_visual_features(
                    input_visual_features, clean_visual_features
                ),
                visual_output_mask=torch.cat(
                    [visual_predict_mask, clean_visual_mask], dim=0
                ),
            )
            clean_visual_count = sum(
                event.target_steps
                for events in selected_events
                for event in events
                if event.direction == "text_to_visual"
            )
            noisy_visual_count = fused_output.visual_logits.shape[0] - (
                clean_visual_count
            )
            output = SimpleNamespace(
                hidden_states=fused_output.hidden_states[: input_ids.shape[0]],
                visual_logits=fused_output.visual_logits[:noisy_visual_count],
            )
            clean_output = SimpleNamespace(
                hidden_states=fused_output.hidden_states[input_ids.shape[0] :],
                visual_logits=fused_output.visual_logits[noisy_visual_count:],
            )
            clean_event_logits = self._event_logits_from_output(
                ar_model,
                clean_output,
                input_ids=input_ids,
                visual_codes=visual_codes,
                selected_events=selected_events,
                visual_output_mask=clean_visual_output_mask,
            )
        else:
            output = ar_model.forward_interleaved(**full_forward_kwargs)

        shifted_labels = labels[:, 1:]
        text_positions = shifted_labels != -100
        if not text_positions.any():
            raise ValueError("batch has no supervised text tokens")
        text_logits = ar_model.project_text(
            output.hidden_states[:, :-1][text_positions]
        ).float()
        text_targets = shifted_labels[text_positions]
        text_nll = F.cross_entropy(text_logits, text_targets, reduction="none")
        marker_mask = text_targets == self.image_gen_id
        text_weights = torch.ones_like(text_nll)
        text_weights[marker_mask] = self.config.text.image_gen_token_weight
        text_sum = (text_nll * text_weights).sum()
        text_weight = text_weights.sum()
        text_predictions = text_logits.argmax(-1)

        full_text_logp = None
        text_logit_index = None
        if needs_noisy_ig_scores or needs_noisy_ema_logits:
            target_sequence_mask = torch.zeros_like(input_ids, dtype=torch.bool)
            target_sequence_mask[:, 1:] = text_positions
            if needs_noisy_ig_scores:
                full_text_logp = torch.zeros(
                    input_ids.shape, device=device, dtype=torch.float32
                )
                full_text_logp[target_sequence_mask] = -text_nll
            if needs_noisy_ema_logits:
                text_logit_index = torch.full_like(input_ids, -1, dtype=torch.long)
                text_logit_index[target_sequence_mask] = torch.arange(
                    text_logits.shape[0], device=device
                )

        levels, spatial, bits = _visual_shape(ar_model)
        visual_output = output.visual_logits
        if input_visual_features is None:
            visual_output = visual_output[visual_predict_mask]
        visual_logits = visual_output.reshape(-1, levels, spatial, bits, 2).float()
        visual_targets = visual_codes[generated_flat_mask]
        if visual_logits.shape[:-1] != visual_targets.shape:
            raise ValueError(
                f"visual target mismatch: {visual_logits.shape[:-1]} != {visual_targets.shape}"
            )
        visual_bit_nll = F.cross_entropy(
            visual_logits.reshape(-1, 2), visual_targets.reshape(-1), reduction="none"
        )
        bit_sum = visual_bit_nll.sum()
        bit_count = visual_bit_nll.numel()
        visual_predictions = visual_logits.argmax(-1)
        feature = self.config.soft_bsq_feature
        feature_mse_sum = bit_sum.new_zeros(())
        feature_cosine_sum = bit_sum.new_zeros(())
        feature_sum = bit_sum.new_zeros(())
        feature_count = 0
        if feature.enabled:
            if target_visual_features is None:
                raise ValueError(
                    "Soft-BSQ feature loss requires prepare_visual_features"
                )
            (
                feature_mse_sum,
                feature_cosine_sum,
                feature_sum,
                feature_count,
            ) = soft_bsq_feature_terms(
                ar_model,
                visual_logits,
                target_visual_features,
                generated_flat_mask,
                cosine_weight=feature.cosine_weight,
                layer_norm_epsilon=feature.layer_norm_epsilon,
                cosine_epsilon=feature.cosine_epsilon,
            )
        full_visual_logp = None
        visual_logit_index = None
        if needs_noisy_ig_scores:
            step_logp = (
                -visual_bit_nll.reshape(visual_targets.shape).flatten(1).sum(-1)
            )
            full_visual_logp = torch.zeros(
                visual_codes.shape[0], device=device, dtype=torch.float32
            )
            full_visual_logp[generated_flat_mask] = step_logp
        if needs_noisy_ema_logits:
            visual_logit_index = torch.full(
                (visual_codes.shape[0],), -1, device=device, dtype=torch.long
            )
            visual_logit_index[generated_flat_mask] = torch.arange(
                visual_logits.shape[0], device=device
            )
        # Bit flip is an SFT input augmentation, not an IG intervention.  The
        # noisy branch above predicts the original labels; paired IG scores
        # remain on the clean GT trajectory. Both IG and EMA need only selected
        # event carriers, never a second full-sequence CE/projection.
        if needs_clean_event_logits and not fuse_clean_event_logits:
            clean_event_logits = self._full_event_logits(
                ar_model,
                input_ids=input_ids,
                attention_mask=attention_mask,
                image_grid_thw=image_grid_thw,
                visual_codes=visual_codes,
                position_ids=position_ids,
                selected_events=selected_events,
                visual_features=clean_visual_features,
                detach=False,
            )

        directions = ("text_to_visual", "visual_to_text")
        ig_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        full_score_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        blind_score_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        full_unit_score_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        blind_unit_score_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        relative_margin_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ig_gate_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ig_target_step_sums = {direction: 0 for direction in directions}
        ig_event_counts = {direction: 0 for direction in directions}
        ema_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ema_gt_ratio_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ema_full_score_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ema_blind_score_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ema_positive_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ema_clipped_values = {
            direction: [[] for _ in selected_events] for direction in directions
        }
        ema_gate_passed_values = {direction: [] for direction in directions}
        ema_gate_position_values = {direction: [] for direction in directions}
        ema_target_step_sums = {direction: 0 for direction in directions}
        ema_event_counts = {direction: 0 for direction in directions}
        event_records: list[dict] = []
        for trajectory_index, events in enumerate(selected_events):
            for event in events:
                direction = event.direction
                record = (
                    {
                        "sample_id": event.sample_id,
                        "event_index": event.index,
                        "direction": direction,
                        "source_spans": [
                            [span.start, span.end]
                            for span in event.source_sequence_spans
                        ],
                        "target_steps": event.target_steps,
                    }
                    if record_events
                    else None
                )
                if ig.enabled:
                    if direction == "text_to_visual":
                        positions = _span_positions(
                            event.target_visual_spans, device=device
                        )
                        if _event_key(event) in clean_event_logits:
                            selected_logits = clean_event_logits[_event_key(event)]
                            full = _visual_selected_score(
                                ar_model,
                                selected_logits.reshape(selected_logits.shape[0], -1),
                                visual_codes,
                                event,
                            )
                        else:
                            if full_visual_logp is None:
                                raise RuntimeError("missing noisy visual IG score map")
                            full = full_visual_logp[positions].mean()
                        direction_config = ig.text_to_visual
                        factor_count = math_prod(visual_targets.shape[1:])
                        per_bit = full / factor_count
                    else:
                        positions = _span_positions(
                            event.target_sequence_spans, device=device
                        )
                        if _event_key(event) in clean_event_logits:
                            token_ids = input_ids[event.sample_index, positions]
                            full = -F.cross_entropy(
                                clean_event_logits[_event_key(event)].float(),
                                token_ids,
                                reduction="mean",
                            )
                        else:
                            if full_text_logp is None:
                                raise RuntimeError("missing noisy text IG score map")
                            full = full_text_logp[event.sample_index, positions].mean()
                        direction_config = ig.visual_to_text
                        factor_count = 1
                        per_bit = None
                    blind = blind_scores[_event_key(event)]
                    full_unit = full / factor_count
                    blind_unit = blind.detach() / factor_count
                    relative_scale = (-blind_unit).clamp_min(
                        RELATIVE_IG_EPSILON
                    )
                    relative_margin = (
                        full_unit.detach() - blind_unit
                    ) / relative_scale
                    gate_argument = (
                        -relative_margin / direction_config.temperature
                        if ig.formulation == "relative_per_unit"
                        else (blind.detach() - full.detach())
                        / direction_config.temperature
                    )
                    event_loss = information_gain_loss(
                        full,
                        blind,
                        direction_config.temperature,
                        formulation=ig.formulation,
                        factor_count=factor_count,
                    )
                    ig_values[direction][trajectory_index].append(event_loss)
                    full_score_values[direction][trajectory_index].append(full)
                    blind_score_values[direction][trajectory_index].append(blind)
                    full_unit_score_values[direction][trajectory_index].append(
                        full_unit
                    )
                    blind_unit_score_values[direction][trajectory_index].append(
                        blind_unit
                    )
                    relative_margin_values[direction][trajectory_index].append(
                        relative_margin
                    )
                    ig_gate_values[direction][trajectory_index].append(
                        torch.sigmoid(gate_argument)
                    )
                    ig_target_step_sums[direction] += event.target_steps
                    ig_event_counts[direction] += 1
                    if record is not None:
                        record.update(
                            {
                                "full_score": full.detach(),
                                "blind_score": blind.detach(),
                                "gain": (full.detach() - blind).detach(),
                                "ig_loss": event_loss.detach(),
                                "ig_formulation": ig.formulation,
                                "factor_count": factor_count,
                                "full_score_per_unit": full_unit.detach(),
                                "blind_score_per_unit": blind_unit,
                                "relative_scale": relative_scale,
                                "relative_margin": relative_margin,
                                "ig_gate": torch.sigmoid(gate_argument),
                            }
                        )
                    if record is not None and per_bit is not None:
                        record["full_score_per_bit"] = per_bit.detach()
                        record["blind_score_per_bit"] = (
                            blind / math_prod(visual_targets.shape[1:])
                        )
                    if record is not None and verify_no_mask:
                        oracle = no_mask_scores[_event_key(event)]
                        record["no_mask_score"] = oracle.detach()
                        record["no_mask_abs_error"] = (
                            oracle - full.detach()
                        ).abs()
                if ema.enabled:
                    key = _event_key(event)
                    if direction == "text_to_visual":
                        positions = _span_positions(
                            event.target_visual_spans, device=device
                        )
                        targets = visual_codes[positions].to(torch.long)
                        if key in clean_event_logits:
                            student_logits = clean_event_logits[key]
                        else:
                            if visual_logit_index is None:
                                raise RuntimeError("missing noisy visual EMA logit map")
                            indices = visual_logit_index[positions]
                            student_logits = visual_logits[indices]
                    else:
                        positions = _span_positions(
                            event.target_sequence_spans, device=device
                        )
                        targets = input_ids[event.sample_index, positions]
                        if key in clean_event_logits:
                            student_logits = clean_event_logits[key]
                        else:
                            if text_logit_index is None:
                                raise RuntimeError("missing noisy text EMA logit map")
                            indices = text_logit_index[event.sample_index, positions]
                            student_logits = text_logits[indices]
                    (
                        target_logp,
                        gt_ratio_mean,
                        full_score,
                        blind_score,
                        positive_fraction,
                        clipped_fraction,
                        position_gate,
                    ) = ema_teacher_data[key]
                    event_jsd, _position_jsd = event_jsd_loss(
                        target_logp,
                        student_logits,
                        position_gate=(
                            position_gate if ema.gt_positive_gate else None
                        ),
                    )
                    if position_gate.numel() != event.target_steps:
                        raise RuntimeError("EMA GT gate and event target lengths differ")
                    ema_values[direction][trajectory_index].append(event_jsd)
                    ema_gt_ratio_values[direction][trajectory_index].append(
                        gt_ratio_mean
                    )
                    ema_full_score_values[direction][trajectory_index].append(
                        full_score
                    )
                    ema_blind_score_values[direction][trajectory_index].append(
                        blind_score
                    )
                    ema_positive_values[direction][trajectory_index].append(
                        positive_fraction
                    )
                    ema_clipped_values[direction][trajectory_index].append(
                        clipped_fraction
                    )
                    ema_gate_passed_values[direction].append(position_gate.sum())
                    ema_gate_position_values[direction].append(
                        position_gate.new_tensor(
                            position_gate.numel(), dtype=torch.long
                        )
                    )
                    ema_target_step_sums[direction] += event.target_steps
                    ema_event_counts[direction] += 1
                    if record is not None:
                        record.update(
                            {
                                "ema_necessity_jsd": event_jsd.detach(),
                                "ema_gt_ratio": gt_ratio_mean.detach(),
                                "ema_gt_gate_coverage": position_gate.float().mean(),
                                "ema_teacher_full_score": full_score.detach(),
                                "ema_teacher_blind_score": blind_score.detach(),
                                "ema_ratio_positive_fraction": positive_fraction.detach(),
                                "ema_ratio_clipped_fraction": clipped_fraction.detach(),
                            }
                        )
                    if record is not None and direction == "text_to_visual":
                        record["ema_necessity_jsd_per_factor"] = (
                            event_jsd / math_prod(targets.shape[1:])
                        ).detach()
                if record is not None:
                    event_records.append(record)

        ig_sums = {}
        full_score_sums = {}
        blind_score_sums = {}
        full_unit_score_sums = {}
        blind_unit_score_sums = {}
        relative_margin_sums = {}
        ig_gate_sums = {}
        trajectory_counts = {}
        ema_sums = {}
        ema_gt_ratio_sums = {}
        ema_full_score_sums = {}
        ema_blind_score_sums = {}
        ema_positive_sums = {}
        ema_clipped_sums = {}
        ema_trajectory_counts = {}
        zero = text_sum.new_zeros(())
        for direction in directions:
            ig_sums[direction], trajectory_counts[direction] = (
                _sum_trajectory_means(ig_values[direction], zero)
            )
            full_score_sums[direction], full_count = _sum_trajectory_means(
                full_score_values[direction], zero
            )
            blind_score_sums[direction], blind_count = _sum_trajectory_means(
                blind_score_values[direction], zero
            )
            full_unit_score_sums[direction], full_unit_count = (
                _sum_trajectory_means(full_unit_score_values[direction], zero)
            )
            blind_unit_score_sums[direction], blind_unit_count = (
                _sum_trajectory_means(blind_unit_score_values[direction], zero)
            )
            relative_margin_sums[direction], relative_margin_count = (
                _sum_trajectory_means(relative_margin_values[direction], zero)
            )
            ig_gate_sums[direction], ig_gate_count = _sum_trajectory_means(
                ig_gate_values[direction], zero
            )
            if (
                full_count != trajectory_counts[direction]
                or blind_count != full_count
                or full_unit_count != full_count
                or blind_unit_count != full_count
                or relative_margin_count != full_count
                or ig_gate_count != full_count
            ):
                raise RuntimeError("IG trajectory statistics have inconsistent counts")
            ema_sums[direction], ema_trajectory_counts[direction] = (
                _sum_trajectory_means(ema_values[direction], zero)
            )
            ema_gt_ratio_sums[direction], ema_ratio_count = _sum_trajectory_means(
                ema_gt_ratio_values[direction], zero
            )
            ema_full_score_sums[direction], ema_full_count = _sum_trajectory_means(
                ema_full_score_values[direction], zero
            )
            ema_blind_score_sums[direction], ema_blind_count = (
                _sum_trajectory_means(ema_blind_score_values[direction], zero)
            )
            ema_positive_sums[direction], ema_positive_count = (
                _sum_trajectory_means(ema_positive_values[direction], zero)
            )
            ema_clipped_sums[direction], ema_clipped_count = (
                _sum_trajectory_means(ema_clipped_values[direction], zero)
            )
            ema_counts = {
                ema_trajectory_counts[direction],
                ema_ratio_count,
                ema_full_count,
                ema_blind_count,
                ema_positive_count,
                ema_clipped_count,
            }
            if len(ema_counts) != 1:
                raise RuntimeError("EMA trajectory statistics have inconsistent counts")

        if global_denominators is None:
            denominators = {
                "text": float(text_weight.detach()),
                "bit": float(bit_count),
                "soft_bsq_feature": float(feature_count),
                "ig_text_to_visual": float(trajectory_counts["text_to_visual"]),
                "ig_visual_to_text": float(trajectory_counts["visual_to_text"]),
                "ema_text_to_visual": float(
                    ema_trajectory_counts["text_to_visual"]
                ),
                "ema_visual_to_text": float(
                    ema_trajectory_counts["visual_to_text"]
                ),
            }
        else:
            denominators = global_denominators

        def normalized(value, key: str):
            denominator = denominators[key]
            if denominator == 0:
                return value.new_zeros(())
            scale = ddp_world_size if global_denominators is not None else 1.0
            return value * scale / denominator

        loss = (
            self.config.text.weight * normalized(text_sum, "text")
            + self.config.bit.weight * normalized(bit_sum, "bit")
            + feature.weight * normalized(feature_sum, "soft_bsq_feature")
            + ig.text_to_visual.weight
            * normalized(ig_sums["text_to_visual"], "ig_text_to_visual")
            + ig.visual_to_text.weight
            * normalized(ig_sums["visual_to_text"], "ig_visual_to_text")
            + ema.text_to_visual.weight
            * normalized(ema_sums["text_to_visual"], "ema_text_to_visual")
            + ema.visual_to_text.weight
            * normalized(ema_sums["visual_to_text"], "ema_visual_to_text")
        )

        statistics = {
            "text_weighted_nll": text_sum.detach(),
            "text_weight": text_weight.detach(),
            "text_nll": text_nll.sum().detach(),
            "text_tokens": text_nll.new_tensor(text_nll.numel()).detach(),
            "text_correct": (text_predictions == text_targets).sum().detach(),
            "marker_nll": text_nll[marker_mask].sum().detach(),
            "markers": marker_mask.sum().detach(),
            "marker_correct": (
                (text_predictions[marker_mask] == text_targets[marker_mask])
                .sum()
                .detach()
            ),
            "bit_nll": bit_sum.detach(),
            "visual_bits": bit_sum.new_tensor(bit_count).detach(),
            "visual_correct": (visual_predictions == visual_targets).sum().detach(),
            "visual_input_noisy_samples": noisy_samples.sum().detach(),
            "visual_input_noisy_visual_bits": (
                noisy_generated.sum() * (levels * spatial * bits)
            ).detach(),
            "visual_input_flipped_bits": flips.sum().detach(),
            "microbatches": bit_sum.new_tensor(1).detach(),
            "model_forward_calls": bit_sum.new_tensor(model_forward_calls).detach(),
            "counterfactual_views": bit_sum.new_tensor(blind_views).detach(),
            "counterfactual_forward_calls": bit_sum.new_tensor(
                counterfactual_forward_calls
            ).detach(),
            "samples": bit_sum.new_tensor(len(batch.sample_ids)).detach(),
            "sequence_tokens": attention_mask.sum().detach(),
            "generated_visual_positions": generated_flat_mask.sum().detach(),
            "visual_factors_per_step": bit_sum.new_tensor(
                levels * spatial * bits
            ).detach(),
        }
        if feature.enabled:
            statistics |= {
                "soft_bsq_feature_mse": feature_mse_sum.detach(),
                "soft_bsq_feature_cosine": feature_cosine_sum.detach(),
                "soft_bsq_feature_loss": feature_sum.detach(),
                "soft_bsq_feature_items": feature_sum.new_tensor(
                    feature_count
                ).detach(),
            }
        for direction in ("text_to_visual", "visual_to_text"):
            prefix = f"ig_{direction}"
            statistics[f"{prefix}_loss"] = ig_sums[direction].detach()
            statistics[f"{prefix}_full_score"] = full_score_sums[direction].detach()
            statistics[f"{prefix}_blind_score"] = blind_score_sums[direction].detach()
            statistics[f"{prefix}_full_score_per_unit"] = full_unit_score_sums[
                direction
            ].detach()
            statistics[f"{prefix}_blind_score_per_unit"] = blind_unit_score_sums[
                direction
            ].detach()
            statistics[f"{prefix}_relative_margin"] = relative_margin_sums[
                direction
            ].detach()
            statistics[f"{prefix}_gate"] = ig_gate_sums[direction].detach()
            statistics[f"{prefix}_events"] = text_sum.new_tensor(
                ig_event_counts[direction]
            ).detach()
            statistics[f"{prefix}_trajectories"] = text_sum.new_tensor(
                trajectory_counts[direction]
            ).detach()
            statistics[f"{prefix}_target_steps"] = text_sum.new_tensor(
                ig_target_step_sums[direction]
            ).detach()
            ema_prefix = f"ema_{direction}"
            statistics[f"{ema_prefix}_necessity_jsd"] = ema_sums[
                direction
            ].detach()
            statistics[f"{ema_prefix}_gt_ratio"] = ema_gt_ratio_sums[
                direction
            ].detach()
            statistics[f"{ema_prefix}_teacher_full_score"] = ema_full_score_sums[
                direction
            ].detach()
            statistics[f"{ema_prefix}_teacher_blind_score"] = ema_blind_score_sums[
                direction
            ].detach()
            statistics[f"{ema_prefix}_ratio_positive"] = ema_positive_sums[
                direction
            ].detach()
            statistics[f"{ema_prefix}_ratio_clipped"] = ema_clipped_sums[
                direction
            ].detach()
            statistics[f"{ema_prefix}_events"] = text_sum.new_tensor(
                ema_event_counts[direction]
            ).detach()
            statistics[f"{ema_prefix}_trajectories"] = text_sum.new_tensor(
                ema_trajectory_counts[direction]
            ).detach()
            statistics[f"{ema_prefix}_target_steps"] = text_sum.new_tensor(
                ema_target_step_sums[direction]
            ).detach()
            statistics[f"{ema_prefix}_gt_gate_passed_positions"] = (
                torch.stack(ema_gate_passed_values[direction]).sum().detach()
                if ema_gate_passed_values[direction]
                else text_sum.new_zeros(())
            )
            statistics[f"{ema_prefix}_gt_gate_positions"] = (
                torch.stack(ema_gate_position_values[direction]).sum().detach()
                if ema_gate_position_values[direction]
                else text_sum.new_zeros(())
            )
        return ObjectiveOutput(
            loss=loss,
            statistics=statistics,
            events=_materialize_event_records(event_records),
        )


__all__ = [
    "InterleavedSFTObjective",
    "ObjectiveOutput",
    "bit_flip_visual_context",
    "sample_bit_flip_probabilities",
    "information_gain_loss",
    "objective_counts",
    "select_noisy_sample_ids",
    "soft_bsq_feature_terms",
    "unwrap_ar_model",
]
Xet Storage Details

Size:: 79.1 kB
Xet hash:: c2f126ea9b78f3e77abcaf29722fce5a368e5ee1699920534c5595c777340f50
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.