| |
| """Validate the active three-encoder MMD composition and gradients on CPU.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| sys.dont_write_bytecode = True |
|
|
|
|
| def assert_same(left: torch.Tensor, right: torch.Tensor, label: str) -> None: |
| torch.testing.assert_close(left, right, rtol=0, atol=0, msg=label) |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--release", choices=("1.3B", "14B"), required=True) |
| args = parser.parse_args() |
|
|
| root = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(root / args.release / "code")) |
|
|
| from losses.multi_encoder_mmd import ( |
| MultiEncoderMMD, |
| normalize_mmd_with_detached_scale, |
| ) |
| from losses.video_nystrom_mmd import ( |
| VideoNystromMMD, |
| build_video_nystrom_reference, |
| ) |
|
|
| shapes = { |
| "dinov3": (1, 2), |
| "vjepa2": (2, 3), |
| "videomae_v2": (1, 4), |
| } |
|
|
| def objective() -> MultiEncoderMMD: |
| branches = {} |
| for index, (name, shape) in enumerate(shapes.items()): |
| generator = torch.Generator().manual_seed(9000 + index) |
| reference = torch.randn(12, *shape, generator=generator) |
| pack = build_video_nystrom_reference( |
| reference, |
| num_landmarks=5, |
| metadata={ |
| "feature": {"name": name, "matrix_shape": list(shape)}, |
| "reference": {"dataset": "active-mmd-check"}, |
| }, |
| bandwidth=1.7, |
| kmeans_iterations=3, |
| ridge=1.0e-5, |
| chunk_size=16, |
| ) |
| branches[name] = VideoNystromMMD( |
| pack, |
| kernel_chunk_size=3, |
| attraction_weight=2.0, |
| current_teacher_attraction_enabled=True, |
| current_teacher_attraction_weight=1.0, |
| ) |
| return MultiEncoderMMD( |
| branches, |
| weights={"dinov3": 10.0, "vjepa2": 1.0, "videomae_v2": 1.0}, |
| normalization_eps=1.0e-4, |
| aggregation_divisor=3.0, |
| ) |
|
|
| generator = torch.Generator().manual_seed(1862643) |
| for mode in ("full", "partial"): |
| combined = objective() |
| generated = { |
| name: torch.randn(4, *shape, generator=generator) |
| for name, shape in shapes.items() |
| } |
| teacher = { |
| name: torch.randn( |
| 4 if mode == "full" else 2, *shape, generator=generator) |
| for name, shape in shapes.items() |
| } |
|
|
| if mode == "full": |
| leaves = { |
| name: value.clone().requires_grad_() |
| for name, value in generated.items() |
| } |
| loss, metadata = combined( |
| leaves, current_teacher=teacher, stage=False) |
| else: |
| indices = torch.tensor([0, 2]) |
| leaves = { |
| name: value.index_select(0, indices).clone().requires_grad_() |
| for name, value in generated.items() |
| } |
| loss, metadata = combined( |
| gradient_students=leaves, |
| student_reference=generated, |
| current_teacher=teacher, |
| gradient_student_indices=indices, |
| stage=False, |
| ) |
|
|
| expected_terms = [] |
| for name, branch in combined.branches.items(): |
| if mode == "full": |
| raw, _ = branch( |
| leaves[name], current_teacher=teacher[name], stage=False) |
| else: |
| raw, _ = branch( |
| gradient_students=leaves[name], |
| student_reference=generated[name], |
| current_teacher=teacher[name], |
| gradient_student_indices=indices, |
| stage=False, |
| ) |
| normalized, denominator = normalize_mmd_with_detached_scale( |
| raw, combined.normalization_eps) |
| expected_terms.append(combined.branch_weights[name] * normalized) |
| assert_same( |
| metadata["branches"][name]["raw_mmd"], |
| raw.detach(), |
| f"{mode} {name} raw MMD", |
| ) |
| assert_same( |
| metadata["branches"][name]["normalization_denominator"], |
| denominator.detach(), |
| f"{mode} {name} normalization denominator", |
| ) |
|
|
| expected = torch.stack(expected_terms).sum() / 3.0 |
| assert_same(loss, expected, f"{mode} combined MMD") |
| observed_gradients = torch.autograd.grad( |
| loss, tuple(leaves.values()), retain_graph=True) |
| expected_gradients = torch.autograd.grad( |
| expected, tuple(leaves.values())) |
| for name, observed, reference in zip( |
| leaves, observed_gradients, expected_gradients): |
| assert_same(observed, reference, f"{mode} {name} feature gradient") |
| assert torch.isfinite(observed).all() and observed.norm() > 0 |
| print( |
| f"PASS {args.release} {mode}: weighted three-encoder MMD, " |
| "metadata and feature gradients") |
|
|
| eps = 1.0e-4 |
| for value in (-2.0e-4, 0.0, 1.0e-8, 0.1, 7.0): |
| observed_input = torch.tensor( |
| value, dtype=torch.float64, requires_grad=True) |
| expected_input = observed_input.detach().clone().requires_grad_(True) |
| observed, observed_denominator = normalize_mmd_with_detached_scale( |
| observed_input, eps) |
| expected_denominator = expected_input.detach() + eps |
| expected = expected_input / expected_denominator |
| assert_same(observed.detach(), expected.detach(), "normalization value") |
| assert_same( |
| observed_denominator, expected_denominator, |
| "normalization denominator") |
| observed_gradient, = torch.autograd.grad(observed, observed_input) |
| expected_gradient, = torch.autograd.grad(expected, expected_input) |
| assert_same( |
| observed_gradient, expected_gradient, "normalization gradient") |
| print(f"PASS {args.release}: detached MMD normalization behavior") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
