Buckets:
| """On-policy four-view distribution matching over sampled event prefixes.""" | |
| from __future__ import annotations | |
| from collections import defaultdict | |
| from dataclasses import dataclass | |
| from typing import Literal | |
| import torch | |
| from shared.distributions import ( | |
| contrast_shaped_log_target, | |
| event_jsd_loss, | |
| symmetric_jsd_from_logits, | |
| ) | |
| from .config import OPDObjectiveConfig | |
| from .text_support import TextSupport, clipped_text_jsd | |
| Direction = Literal["text_to_visual", "visual_to_text"] | |
| Modality = Literal["text", "visual"] | |
| class OPDEventLogits: | |
| trajectory_id: str | |
| sample_index: int | |
| event_index: int | |
| direction: Direction | |
| modality: Modality | |
| has_prior_generated_history: bool | |
| verifier_correct: bool | |
| privilege_tag: str | None | |
| teacher_11_logits: torch.Tensor | None | |
| teacher_10_logits: torch.Tensor | None | |
| teacher_01_logits: torch.Tensor | None | |
| teacher_00_logits: torch.Tensor | None | |
| student_11_logits: torch.Tensor | None | |
| student_01_logits: torch.Tensor | None | |
| unprivileged_teacher_11_logits: torch.Tensor | None = None | |
| unprivileged_teacher_10_logits: torch.Tensor | None = None | |
| unprivileged_teacher_01_logits: torch.Tensor | None = None | |
| unprivileged_teacher_00_logits: torch.Tensor | None = None | |
| record_event_metadata: bool = True | |
| text_support: TextSupport | None = None | |
| class OPDObjectiveOutput: | |
| loss: torch.Tensor | |
| statistics: dict[str, torch.Tensor] | |
| def _composite_jsd(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: | |
| factor_jsd = symmetric_jsd_from_logits(left, right) | |
| per_position = ( | |
| factor_jsd | |
| if factor_jsd.ndim == 1 | |
| else factor_jsd.flatten(1).mean(-1) | |
| ) | |
| return per_position.mean() | |
| def _factor_entropy(logits: torch.Tensor) -> torch.Tensor: | |
| logp = torch.log_softmax(logits.detach().float(), dim=-1) | |
| return -(logp.exp() * logp).sum(-1).mean() | |
| def _argmax_overlap(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: | |
| return (left.detach().argmax(-1) == right.detach().argmax(-1)).float().mean() | |
| def compute_opd_objective( | |
| events: list[OPDEventLogits], config: OPDObjectiveConfig | |
| ) -> OPDObjectiveOutput: | |
| """Compute contrast-shaped necessity and sufficiency JSD.""" | |
| if config.sampled_target_gate: | |
| raise ValueError("on-policy OPD does not permit a sampled target gate") | |
| live_tensor = next( | |
| ( | |
| tensor | |
| for event in events | |
| for tensor in (event.student_11_logits, event.student_01_logits) | |
| if tensor is not None | |
| ), | |
| None, | |
| ) | |
| zero = ( | |
| live_tensor.new_zeros((), dtype=torch.float32) | |
| if live_tensor is not None | |
| else torch.zeros((), dtype=torch.float32) | |
| ) | |
| values = { | |
| direction: { | |
| objective: defaultdict(list) | |
| for objective in ("necessity", "sufficiency") | |
| } | |
| for direction in ("text_to_visual", "visual_to_text") | |
| } | |
| outcome_values = { | |
| direction: { | |
| objective: { | |
| outcome: defaultdict(list) | |
| for outcome in ("correct", "incorrect") | |
| } | |
| for objective in ("necessity", "sufficiency") | |
| } | |
| for direction in ("text_to_visual", "visual_to_text") | |
| } | |
| outcome_positions = { | |
| direction: { | |
| objective: {"correct": 0, "incorrect": 0} | |
| for objective in ("necessity", "sufficiency") | |
| } | |
| for direction in ("text_to_visual", "visual_to_text") | |
| } | |
| diagnostics: dict[str, list[torch.Tensor]] = defaultdict(list) | |
| trajectory_samples: dict[str, int] = {} | |
| trajectory_verdicts: dict[str, bool] = {} | |
| trajectory_privilege: dict[str, str | None] = {} | |
| seen_events: set[tuple[str, int]] = set() | |
| for event in events: | |
| event_key = event.trajectory_id, event.event_index | |
| if event_key in seen_events: | |
| raise ValueError(f"duplicate event identity: {event_key!r}") | |
| seen_events.add(event_key) | |
| previous_sample = trajectory_samples.setdefault( | |
| event.trajectory_id, event.sample_index | |
| ) | |
| if previous_sample != event.sample_index: | |
| raise ValueError("one trajectory_id cannot alias multiple sample indices") | |
| previous_verdict = trajectory_verdicts.setdefault( | |
| event.trajectory_id, event.verifier_correct | |
| ) | |
| if previous_verdict is not event.verifier_correct: | |
| raise ValueError("events in one trajectory must reuse one verifier result") | |
| if event.trajectory_id in trajectory_privilege: | |
| if trajectory_privilege[event.trajectory_id] != event.privilege_tag: | |
| raise ValueError("events in one trajectory must share one privilege tag") | |
| else: | |
| trajectory_privilege[event.trajectory_id] = event.privilege_tag | |
| expected_modality = { | |
| "text_to_visual": "visual", | |
| "visual_to_text": "text", | |
| }.get(event.direction) | |
| if expected_modality != event.modality: | |
| raise ValueError("event direction/modality target mismatch") | |
| if event.student_11_logits is not None and ( | |
| event.teacher_11_logits is None or event.teacher_10_logits is None | |
| ): | |
| raise ValueError("necessity requires teacher 11/10 and student 11 logits") | |
| if event.student_01_logits is not None and ( | |
| event.teacher_01_logits is None or event.teacher_00_logits is None | |
| ): | |
| raise ValueError("sufficiency requires teacher 01/00 and student 01 logits") | |
| event_logits = tuple( | |
| logits | |
| for logits in ( | |
| event.teacher_11_logits, | |
| event.teacher_10_logits, | |
| event.teacher_01_logits, | |
| event.teacher_00_logits, | |
| event.student_11_logits, | |
| event.student_01_logits, | |
| ) | |
| if logits is not None | |
| ) | |
| if not event_logits: | |
| raise ValueError("an OPD event must contain one enabled objective") | |
| if len({tuple(logits.shape) for logits in event_logits}) != 1: | |
| raise ValueError("all enabled event logits must have identical shapes") | |
| shape = event_logits[0].shape | |
| if event.modality == "text" and (len(shape) != 2 or shape[-1] < 2): | |
| raise ValueError("text event logits must have shape [J, vocabulary]") | |
| if event.modality == "visual" and (len(shape) < 3 or shape[-1] != 2): | |
| raise ValueError( | |
| "visual event logits must have shape [J, factors..., 2]" | |
| ) | |
| if shape[0] == 0: | |
| raise ValueError("event must contain at least one target position") | |
| if event.verifier_correct and event.privilege_tag is not None: | |
| raise ValueError("teacher privilege is allowed only on incorrect events") | |
| unprivileged_logits = ( | |
| event.unprivileged_teacher_11_logits, | |
| event.unprivileged_teacher_10_logits, | |
| event.unprivileged_teacher_01_logits, | |
| event.unprivileged_teacher_00_logits, | |
| ) | |
| if any(value is not None for value in unprivileged_logits): | |
| if event.privilege_tag is None: | |
| raise ValueError( | |
| "unprivileged teacher diagnostics require a privileged event" | |
| ) | |
| required = ( | |
| (0, 1) if event.student_11_logits is not None else () | |
| ) + ((2, 3) if event.student_01_logits is not None else ()) | |
| if any(unprivileged_logits[index] is None for index in required): | |
| raise ValueError( | |
| "unprivileged diagnostics must cover every enabled view pair" | |
| ) | |
| if any( | |
| value.shape != shape | |
| for value in unprivileged_logits | |
| if value is not None | |
| ): | |
| raise ValueError( | |
| "unprivileged teacher diagnostics must match event shapes" | |
| ) | |
| direction_config = ( | |
| config.text_to_visual | |
| if event.direction == "text_to_visual" | |
| else config.visual_to_text | |
| ) | |
| outcome = "correct" if event.verifier_correct else "incorrect" | |
| if event.text_support is not None: | |
| support = event.text_support | |
| diagnostics["text_support_events"].append(zero.new_tensor(1)) | |
| diagnostics["text_support_retained_mass"].append(support.retained_mass.mean()) | |
| diagnostics["text_support_minimum_retained_mass"].append(support.retained_mass.min()) | |
| diagnostics["text_support_size"].append(support.valid.sum(-1).float().mean()) | |
| diagnostics["text_support_effective_top_k"].append(zero.new_tensor(support.top_k)) | |
| if event.student_11_logits is not None: | |
| necessity_target, necessity_ratio, *_ = contrast_shaped_log_target( | |
| event.teacher_11_logits, | |
| event.teacher_10_logits, | |
| contrast_alpha=direction_config.necessity_contrast_alpha, | |
| ratio_clip=config.ratio_clip, | |
| ) | |
| necessity, _ = event_jsd_loss( | |
| necessity_target, | |
| event.student_11_logits, | |
| factor_reduction=( | |
| config.visual_factor_reduction | |
| if event.modality == "visual" | |
| else "sum" | |
| ), | |
| ) | |
| if event.modality == "text" and config.text_contribution_clip is not None: | |
| necessity, clipped = clipped_text_jsd(necessity_target, event.student_11_logits, contribution_clip=config.text_contribution_clip) | |
| valid = event.text_support.valid if event.text_support is not None else torch.ones_like(clipped, dtype=torch.bool) | |
| diagnostics["text_contribution_clipped_fraction"].append(clipped[valid].mean()) | |
| diagnostics["text_contribution_events"].append(zero.new_tensor(1)) | |
| diagnostics["teacher_student_full_overlap"].append( | |
| _argmax_overlap(event.teacher_11_logits, event.student_11_logits) | |
| ) | |
| diagnostics["teacher_full_entropy_per_factor"].append( | |
| _factor_entropy(event.teacher_11_logits) | |
| ) | |
| diagnostics["student_full_entropy_per_factor"].append( | |
| _factor_entropy(event.student_11_logits) | |
| ) | |
| diagnostics["teacher_student_full_jsd"].append( | |
| _composite_jsd( | |
| event.teacher_11_logits, event.student_11_logits.detach() | |
| ) | |
| ) | |
| diagnostics["teacher_necessity_jsd"].append( | |
| _composite_jsd(event.teacher_11_logits, event.teacher_10_logits) | |
| ) | |
| if event.text_support is not None: | |
| necessity_ratio = necessity_ratio[event.text_support.valid] | |
| diagnostics["necessity_ratio_mean"].append(necessity_ratio.mean()) | |
| diagnostics["necessity_ratio_positive_fraction"].append( | |
| (necessity_ratio > 0).float().mean() | |
| ) | |
| diagnostics["necessity_ratio_clipped_fraction"].append( | |
| (necessity_ratio > config.ratio_clip).float().mean() | |
| ) | |
| values[event.direction]["necessity"][event.trajectory_id].append( | |
| necessity | |
| ) | |
| outcome_values[event.direction]["necessity"][outcome][ | |
| event.trajectory_id | |
| ].append(necessity) | |
| outcome_positions[event.direction]["necessity"][outcome] += int( | |
| event.student_11_logits.shape[0] | |
| ) | |
| if event.has_prior_generated_history and event.student_01_logits is not None: | |
| sufficiency_target, sufficiency_ratio, *_ = contrast_shaped_log_target( | |
| event.teacher_01_logits, | |
| event.teacher_00_logits, | |
| contrast_alpha=direction_config.sufficiency_contrast_alpha, | |
| ratio_clip=config.ratio_clip, | |
| ) | |
| sufficiency, _ = event_jsd_loss( | |
| sufficiency_target, | |
| event.student_01_logits, | |
| factor_reduction=( | |
| config.visual_factor_reduction | |
| if event.modality == "visual" | |
| else "sum" | |
| ), | |
| ) | |
| if event.modality == "text" and config.text_contribution_clip is not None: | |
| sufficiency, clipped = clipped_text_jsd(sufficiency_target, event.student_01_logits, contribution_clip=config.text_contribution_clip) | |
| valid = event.text_support.valid if event.text_support is not None else torch.ones_like(clipped, dtype=torch.bool) | |
| diagnostics["text_contribution_clipped_fraction"].append(clipped[valid].mean()) | |
| diagnostics["text_contribution_events"].append(zero.new_tensor(1)) | |
| diagnostics["teacher_student_local_overlap"].append( | |
| _argmax_overlap(event.teacher_01_logits, event.student_01_logits) | |
| ) | |
| diagnostics["teacher_local_entropy_per_factor"].append( | |
| _factor_entropy(event.teacher_01_logits) | |
| ) | |
| diagnostics["student_local_entropy_per_factor"].append( | |
| _factor_entropy(event.student_01_logits) | |
| ) | |
| diagnostics["teacher_student_local_jsd"].append( | |
| _composite_jsd( | |
| event.teacher_01_logits, event.student_01_logits.detach() | |
| ) | |
| ) | |
| diagnostics["teacher_sufficiency_jsd"].append( | |
| _composite_jsd(event.teacher_01_logits, event.teacher_00_logits) | |
| ) | |
| if event.text_support is not None: | |
| sufficiency_ratio = sufficiency_ratio[event.text_support.valid] | |
| diagnostics["sufficiency_ratio_mean"].append(sufficiency_ratio.mean()) | |
| diagnostics["sufficiency_ratio_positive_fraction"].append( | |
| (sufficiency_ratio > 0).float().mean() | |
| ) | |
| diagnostics["sufficiency_ratio_clipped_fraction"].append( | |
| (sufficiency_ratio > config.ratio_clip).float().mean() | |
| ) | |
| values[event.direction]["sufficiency"][event.trajectory_id].append( | |
| sufficiency | |
| ) | |
| outcome_values[event.direction]["sufficiency"][outcome][ | |
| event.trajectory_id | |
| ].append(sufficiency) | |
| outcome_positions[event.direction]["sufficiency"][outcome] += int( | |
| event.student_01_logits.shape[0] | |
| ) | |
| if ( | |
| event.record_event_metadata | |
| and unprivileged_logits[0] is not None | |
| and event.teacher_11_logits is not None | |
| and event.teacher_10_logits is not None | |
| ): | |
| privileged_necessity = _composite_jsd( | |
| event.teacher_11_logits, event.teacher_10_logits | |
| ) | |
| unprivileged_necessity = _composite_jsd( | |
| unprivileged_logits[0], unprivileged_logits[1] | |
| ) | |
| diagnostics["privileged_teacher_necessity_jsd"].append( | |
| privileged_necessity | |
| ) | |
| diagnostics["unprivileged_teacher_necessity_jsd"].append( | |
| unprivileged_necessity | |
| ) | |
| diagnostics["privilege_necessity_retention"].append( | |
| privileged_necessity | |
| / (unprivileged_necessity + config.privilege_retention_epsilon) | |
| ) | |
| if ( | |
| event.has_prior_generated_history | |
| and unprivileged_logits[2] is not None | |
| and event.teacher_01_logits is not None | |
| and event.teacher_00_logits is not None | |
| ): | |
| privileged_sufficiency = _composite_jsd( | |
| event.teacher_01_logits, event.teacher_00_logits | |
| ) | |
| unprivileged_sufficiency = _composite_jsd( | |
| unprivileged_logits[2], unprivileged_logits[3] | |
| ) | |
| diagnostics["privileged_teacher_sufficiency_jsd"].append( | |
| privileged_sufficiency | |
| ) | |
| diagnostics["unprivileged_teacher_sufficiency_jsd"].append( | |
| unprivileged_sufficiency | |
| ) | |
| diagnostics["privilege_sufficiency_retention"].append( | |
| privileged_sufficiency | |
| / ( | |
| unprivileged_sufficiency | |
| + config.privilege_retention_epsilon | |
| ) | |
| ) | |
| statistics: dict[str, torch.Tensor] = {} | |
| reduced = {} | |
| for direction in ("text_to_visual", "visual_to_text"): | |
| for objective in ("necessity", "sufficiency"): | |
| trajectories = values[direction][objective] | |
| trajectory_means = [ | |
| torch.stack(event_values).mean() | |
| for event_values in trajectories.values() | |
| ] | |
| reduced[direction, objective] = ( | |
| torch.stack(trajectory_means).mean() if trajectory_means else zero | |
| ) | |
| prefix = f"{direction}_{objective}" | |
| statistics[f"{prefix}_loss"] = reduced[direction, objective].detach() | |
| statistics[f"{prefix}_events"] = zero.new_tensor( | |
| sum(len(event_values) for event_values in trajectories.values()) | |
| ) | |
| statistics[f"{prefix}_trajectories"] = zero.new_tensor(len(trajectories)) | |
| for outcome in ("correct", "incorrect"): | |
| outcome_trajectories = outcome_values[direction][objective][outcome] | |
| outcome_means = [ | |
| torch.stack(event_values).mean() | |
| for event_values in outcome_trajectories.values() | |
| ] | |
| outcome_loss = ( | |
| torch.stack(outcome_means).mean() if outcome_means else zero | |
| ) | |
| outcome_prefix = f"{prefix}_{outcome}" | |
| statistics[f"{outcome_prefix}_loss"] = outcome_loss.detach() | |
| statistics[f"{outcome_prefix}_events"] = zero.new_tensor( | |
| sum( | |
| len(event_values) | |
| for event_values in outcome_trajectories.values() | |
| ) | |
| ) | |
| statistics[f"{outcome_prefix}_trajectories"] = zero.new_tensor( | |
| len(outcome_trajectories) | |
| ) | |
| statistics[f"{outcome_prefix}_target_positions"] = zero.new_tensor( | |
| outcome_positions[direction][objective][outcome] | |
| ) | |
| statistics["teacher_privileged_events"] = zero.new_tensor( | |
| sum( | |
| event.privilege_tag is not None and event.record_event_metadata | |
| for event in events | |
| ) | |
| ) | |
| statistics["teacher_privileged_target_positions"] = zero.new_tensor( | |
| sum( | |
| int(event.student_11_logits.shape[0]) | |
| for event in events | |
| if event.privilege_tag is not None | |
| and event.record_event_metadata | |
| and event.student_11_logits is not None | |
| ) | |
| ) | |
| for name, values in diagnostics.items(): | |
| statistics[name] = torch.stack(values).mean().detach() | |
| loss = zero | |
| for direction, direction_config in ( | |
| ("text_to_visual", config.text_to_visual), | |
| ("visual_to_text", config.visual_to_text), | |
| ): | |
| loss = ( | |
| loss | |
| + direction_config.necessity_weight * reduced[direction, "necessity"] | |
| + direction_config.sufficiency_weight * reduced[direction, "sufficiency"] | |
| ) | |
| return OPDObjectiveOutput( | |
| loss=loss, | |
| statistics=statistics, | |
| ) | |
| __all__ = ["OPDEventLogits", "OPDObjectiveOutput", "compute_opd_objective"] | |
Xet Storage Details
- Size:
- 19.9 kB
- Xet hash:
- 6aaa37030ed22b78b3957902cec4b5bcc95111e41a8f3148305aeda508281861
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.