| """LoopQ sharing-gap scoring and progressive SLT selection (LQ4). |
| |
| Equation (6) scores a transform-weight group by summing cross-loop gradient |
| variance after normalization by diagonal Fisher curvature. This module |
| consumes saved calibration statistics; collection and trajectory calibration |
| belong to LQ6. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| from collections.abc import Callable, Mapping |
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import torch |
|
|
|
|
| DEFAULT_FISHER_EPSILON = 1e-8 |
| OURO_1_4B_GROUP_BUDGET = 4 |
| HUGINN_PARAMETER_FRACTION = 0.05 |
|
|
|
|
| def loss_value_gradients( |
| loss: torch.Tensor, |
| values: tuple[torch.Tensor, ...], |
| ) -> tuple[torch.Tensor | None, ...]: |
| """Evaluate loss-to-intermediate VJPs in one full-graph traversal.""" |
|
|
| return torch.autograd.grad( |
| loss, values, retain_graph=True, allow_unused=True |
| ) |
|
|
|
|
| def batched_chain_rule_gradients( |
| values: tuple[torch.Tensor, ...], |
| value_gradients: tuple[torch.Tensor | None, ...], |
| parameters: tuple[torch.nn.Parameter, ...], |
| ) -> tuple[torch.Tensor | None, ...]: |
| """Sum exact intermediate-to-parameter VJPs for supplied loss gradients.""" |
|
|
| active = tuple( |
| (value, gradient) |
| for value, gradient in zip(values, value_gradients) |
| if gradient is not None |
| ) |
| if not active: |
| return tuple(None for _ in parameters) |
| return torch.autograd.grad( |
| tuple(item[0] for item in active), parameters, |
| grad_outputs=tuple(item[1] for item in active), |
| retain_graph=True, allow_unused=True, |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class SharingGapStatistics: |
| """Saved sufficient statistics for one selectable transform-weight group.""" |
|
|
| gradients: torch.Tensor |
| fisher_diagonal: torch.Tensor |
| parameter_count: int |
|
|
| def to_dict(self) -> dict[str, Any]: |
| self.validate() |
| return { |
| "gradients": self.gradients.detach().cpu(), |
| "fisher_diagonal": self.fisher_diagonal.detach().cpu(), |
| "parameter_count": self.parameter_count, |
| } |
|
|
| @classmethod |
| def from_dict(cls, state: Mapping[str, Any]) -> "SharingGapStatistics": |
| required = {"gradients", "fisher_diagonal", "parameter_count"} |
| missing = required.difference(state) |
| if missing: |
| raise ValueError(f"sharing-gap statistics are missing fields: {sorted(missing)}") |
| item = cls( |
| gradients=torch.as_tensor(state["gradients"]), |
| fisher_diagonal=torch.as_tensor(state["fisher_diagonal"]), |
| parameter_count=int(state["parameter_count"]), |
| ) |
| item.validate() |
| return item |
|
|
| def validate(self, *, expected_loop_count: int | None = None) -> None: |
| if self.gradients.ndim < 2 or self.gradients.shape[0] < 2: |
| raise ValueError("gradients must have shape (at least 2 loops, coordinates...)") |
| if expected_loop_count is not None and self.gradients.shape[0] != expected_loop_count: |
| raise ValueError( |
| f"expected {expected_loop_count} loops, got {self.gradients.shape[0]}" |
| ) |
| if tuple(self.fisher_diagonal.shape) != tuple(self.gradients.shape[1:]): |
| raise ValueError( |
| "fisher_diagonal shape must equal the gradient coordinate shape, " |
| f"got {tuple(self.fisher_diagonal.shape)} and {tuple(self.gradients.shape[1:])}" |
| ) |
| if not self.gradients.is_floating_point() or not self.fisher_diagonal.is_floating_point(): |
| raise TypeError("gradients and fisher_diagonal must be floating point") |
| if not torch.isfinite(self.gradients).all() or not torch.isfinite(self.fisher_diagonal).all(): |
| raise ValueError("sharing-gap statistics must be finite") |
| if (self.fisher_diagonal < 0).any(): |
| raise ValueError("fisher_diagonal must be non-negative") |
| if self.parameter_count <= 0: |
| raise ValueError("parameter_count must be positive") |
|
|
|
|
| def sharing_gap_score( |
| statistics: SharingGapStatistics, |
| *, |
| epsilon: float = DEFAULT_FISHER_EPSILON, |
| expected_loop_count: int | None = None, |
| ) -> float: |
| """Compute Eq. (6) with population variance (division by T).""" |
|
|
| statistics.validate(expected_loop_count=expected_loop_count) |
| if not math.isfinite(epsilon) or epsilon <= 0: |
| raise ValueError("epsilon must be finite and positive") |
| gradients = statistics.gradients.detach().to(device="cpu", dtype=torch.float64) |
| fisher = statistics.fisher_diagonal.detach().to(device="cpu", dtype=torch.float64) |
| variance = gradients.var(dim=0, unbiased=False) |
| return float((variance / (fisher + epsilon)).sum().item()) |
|
|
|
|
| @dataclass(frozen=True) |
| class SelectionRound: |
| round_index: int |
| scores: dict[str, float] |
| selected_group: str |
| selected_parameter_count: int |
| cumulative_parameter_count: int |
|
|
|
|
| @dataclass(frozen=True) |
| class ProgressiveSelectionArtifact: |
| """Serializable audit trail proving selection and score recomputation.""" |
|
|
| format_version: int |
| loop_count: int |
| epsilon: float |
| stopping_rule: str |
| requested_budget: float | int |
| selected_groups: tuple[str, ...] |
| rounds: tuple[SelectionRound, ...] |
|
|
| def to_dict(self) -> dict[str, Any]: |
| return { |
| "format": "loopq_progressive_slt", |
| "format_version": self.format_version, |
| "loop_count": self.loop_count, |
| "epsilon": self.epsilon, |
| "stopping_rule": self.stopping_rule, |
| "requested_budget": self.requested_budget, |
| "selected_groups": list(self.selected_groups), |
| "rounds": [ |
| { |
| "round_index": item.round_index, |
| "scores": dict(sorted(item.scores.items())), |
| "selected_group": item.selected_group, |
| "selected_parameter_count": item.selected_parameter_count, |
| "cumulative_parameter_count": item.cumulative_parameter_count, |
| } |
| for item in self.rounds |
| ], |
| } |
|
|
|
|
| StatisticsCallback = Callable[ |
| [tuple[str, ...], int], Mapping[str, SharingGapStatistics] |
| ] |
|
|
|
|
| def _validate_candidate_names(statistics: Mapping[str, SharingGapStatistics]) -> None: |
| if not statistics: |
| raise ValueError("at least one SLT candidate is required") |
| if any(not isinstance(name, str) or not name for name in statistics): |
| raise ValueError("SLT candidate names must be non-empty strings") |
|
|
|
|
| def progressive_select_by_group_count( |
| recompute_statistics: StatisticsCallback, |
| *, |
| budget: int, |
| loop_count: int, |
| epsilon: float = DEFAULT_FISHER_EPSILON, |
| ) -> ProgressiveSelectionArtifact: |
| """Progressively select exactly ``budget`` groups, recomputing every round.""" |
|
|
| if budget <= 0: |
| raise ValueError("group budget must be positive") |
| return _progressive_select( |
| recompute_statistics, |
| loop_count=loop_count, |
| epsilon=epsilon, |
| group_budget=budget, |
| parameter_target=None, |
| requested_budget=budget, |
| stopping_rule="group_count", |
| ) |
|
|
|
|
| def huginn_parameter_target(parameter_counts: Mapping[str, int]) -> int: |
| """Return ceil(5% of all candidate transform parameters).""" |
|
|
| if not parameter_counts or any(count <= 0 for count in parameter_counts.values()): |
| raise ValueError("Huginn candidate parameter counts must be positive") |
| return math.ceil(sum(parameter_counts.values()) * HUGINN_PARAMETER_FRACTION) |
|
|
|
|
| def progressive_select_huginn( |
| recompute_statistics: StatisticsCallback, |
| *, |
| loop_count: int = 32, |
| epsilon: float = DEFAULT_FISHER_EPSILON, |
| ) -> ProgressiveSelectionArtifact: |
| """Select whole groups until their parameters first meet the 5% target.""" |
|
|
| initial = recompute_statistics(tuple(), 0) |
| _validate_candidate_names(initial) |
| target = huginn_parameter_target( |
| {name: item.parameter_count for name, item in initial.items()} |
| ) |
|
|
| def cached_first_callback(selected: tuple[str, ...], round_index: int): |
| if round_index == 0 and not selected: |
| return initial |
| return recompute_statistics(selected, round_index) |
|
|
| return _progressive_select( |
| cached_first_callback, |
| loop_count=loop_count, |
| epsilon=epsilon, |
| group_budget=None, |
| parameter_target=target, |
| requested_budget=HUGINN_PARAMETER_FRACTION, |
| stopping_rule="parameter_fraction_ceil_then_whole_group_overshoot", |
| ) |
|
|
|
|
| def _progressive_select( |
| recompute_statistics: StatisticsCallback, |
| *, |
| loop_count: int, |
| epsilon: float, |
| group_budget: int | None, |
| parameter_target: int | None, |
| requested_budget: float | int, |
| stopping_rule: str, |
| ) -> ProgressiveSelectionArtifact: |
| if loop_count < 2: |
| raise ValueError("loop_count must be at least 2") |
| selected: list[str] = [] |
| rounds: list[SelectionRound] = [] |
| cumulative_parameters = 0 |
| expected_names: set[str] | None = None |
| expected_parameter_counts: dict[str, int] | None = None |
| while True: |
| if group_budget is not None and len(selected) >= group_budget: |
| break |
| if parameter_target is not None and cumulative_parameters >= parameter_target: |
| break |
| statistics = recompute_statistics(tuple(selected), len(rounds)) |
| _validate_candidate_names(statistics) |
| current_names = set(statistics) |
| current_counts = {name: item.parameter_count for name, item in statistics.items()} |
| if expected_names is None: |
| expected_names = current_names |
| expected_parameter_counts = current_counts |
| elif current_names != expected_names or current_counts != expected_parameter_counts: |
| raise ValueError( |
| "progressive recomputation must preserve the candidate set and parameter counts" |
| ) |
| remaining = sorted(set(statistics).difference(selected)) |
| if not remaining: |
| raise ValueError("SLT budget cannot be satisfied by remaining candidates") |
| scored = { |
| name: sharing_gap_score( |
| statistics[name], epsilon=epsilon, expected_loop_count=loop_count |
| ) |
| for name in remaining |
| } |
| |
| chosen = min(remaining, key=lambda name: (-scored[name], name)) |
| count = statistics[chosen].parameter_count |
| cumulative_parameters += count |
| selected.append(chosen) |
| rounds.append( |
| SelectionRound( |
| round_index=len(rounds), |
| scores=dict(sorted(scored.items())), |
| selected_group=chosen, |
| selected_parameter_count=count, |
| cumulative_parameter_count=cumulative_parameters, |
| ) |
| ) |
| return ProgressiveSelectionArtifact( |
| format_version=1, |
| loop_count=loop_count, |
| epsilon=epsilon, |
| stopping_rule=stopping_rule, |
| requested_budget=requested_budget, |
| selected_groups=tuple(selected), |
| rounds=tuple(rounds), |
| ) |
|
|
