feat: add projection_dim support to kronecker factored Hessians

bergson/hessians/apply_hessian.py

@@ -3,12 +3,14 @@
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Literal
 
 import numpy as np
 import torch
 import torch.distributed as dist
 from simple_parsing import ArgumentParser
 
+from bergson.collector.collector import create_projection_matrix
 from bergson.config import InversionConfig
 from bergson.data import column_offsets, create_index, load_gradients
 from bergson.distributed import init_dist
@@ -28,6 +30,10 @@ class EkfacConfig:
     `HessianConfig.ev_correction=True`."""
     apply_batch_size: int = 32
     """Number of query gradients moved on-device and preconditioned at a time."""
+    projection_dim: int = 0
+    """When set, compress each module's IVHP output to a ``[p, p]`` Kronecker
+    random projection (``P_S @ (H^-1 G) @ P_A^T``)."""
+    projection_type: Literal["normal", "rademacher"] = "rademacher"
     debug: bool = False
 
 
@@ -52,6 +58,13 @@ def __init__(
         if inversion_cfg is not None and apply_fn is not None:
             raise ValueError("Pass either inversion_cfg or apply_fn, not both.")
 
+        if cfg.projection_dim > 0 and cfg.ev_correction:
+            raise ValueError(
+                "projection_dim compression is not supported with EK-FAC "
+                "(ev_correction=True); set ev_correction=False or "
+                "projection_dim=0."
+            )
+
         self.cfg = cfg
         self.path = cfg.hessian_method_path
         self.gradient_path = cfg.gradient_path
@@ -77,9 +90,18 @@ def compute_ivhp_sharded(self):
             ev_correction=self.cfg.ev_correction,
         )
 
-        grad_sizes = {
+        o_dims = {
             name: preconditioner.eigen_g[name].shape[1]
-            * preconditioner.eigen_a[name].shape[1]
+            for name in preconditioner.eigen_a
+        }
+        i_dims = {
+            name: preconditioner.eigen_a[name].shape[1]
+            for name in preconditioner.eigen_a
+        }
+
+        p = self.cfg.projection_dim
+        grad_sizes = {
+            name: p * p if p > 0 else o_dims[name] * i_dims[name]
             for name in preconditioner.eigen_a
         }
 
@@ -126,6 +148,30 @@ def compute_ivhp_sharded(self):
 
             self.logger.debug("Finished H^{-1} G = Q_S @ (G' / lambda) @ Q_A^T batch")
 
+            if p > 0:
+                for name, flat in transformed.items():
+                    if name not in o_dims:
+                        continue
+                    g = flat.view(-1, o_dims[name], i_dims[name])
+                    P_l = create_projection_matrix(
+                        f"{name}/left",
+                        p,
+                        o_dims[name],
+                        g.dtype,
+                        g.device,
+                        self.cfg.projection_type,
+                    )
+                    P_r = create_projection_matrix(
+                        f"{name}/right",
+                        p,
+                        i_dims[name],
+                        g.dtype,
+                        g.device,
+                        self.cfg.projection_type,
+                    )
+                    transformed[name] = torch.einsum("ps,nsa,ra->npr", P_l, g, P_r)
+                self.logger.debug("Compressed IVHP output to [p, p] per module")
+
             # Stage the async D2H copies, synchronize once, then read them
             # into the numpy buffer. `.numpy()` is a host read so the
             # sync must sit between it and the copies.

bergson/hessians/pipeline.py

@@ -56,6 +56,14 @@ def hessian_pipeline(
     3. Apply the inverse Hessian to the mean query gradient.
     4. Score each training example against the transformed query gradient.
     """
+    if preprocess_cfg.unit_normalize:
+        raise ValueError(
+            "preprocess_cfg.unit_normalize (cosine similarity) is not "
+            "supported with Kronecker-factored Hessians; hessian_pipeline "
+            "only ever fits and applies a factored (kfac/tkfac/shampoo) "
+            "Hessian."
+        )
+
     run_path = index_cfg.run_path
     method = hessian_cfg.method
     query_path = f"{run_path}/query"
@@ -110,6 +118,8 @@ def _validate(cfg: IndexConfig):
             gradient_path=query_path,
             run_path=transformed_query_path,
             ev_correction=hessian_cfg.ev_correction,
+            projection_dim=index_cfg.projection_dim,
+            projection_type=index_cfg.projection_type,
         )
         launch_distributed_run(
             "apply_hessian",
@@ -123,7 +133,6 @@ def _validate(cfg: IndexConfig):
     if not _step_complete(scores_path, resume):
         score_index_cfg = deepcopy(index_cfg)
         score_index_cfg.run_path = scores_path
-        score_index_cfg.projection_dim = 0
         score_cfg.query_path = transformed_query_path
         score_cfg.higher_is_better = True
         _validate(score_index_cfg)

bergson/score/score.py

@@ -378,6 +378,16 @@ def score_dataset(
         and index_cfg.projection_dim != 0
         and is_factored_hessian(preprocess_cfg.hessian_path)
     ):
+        if preprocess_cfg.unit_normalize:
+            raise ValueError(
+                f"Scoring with a factored (EKFAC) hessian at "
+                f"{preprocess_cfg.hessian_path} and unit_normalize=True (cosine "
+                f"similarity) requires projection_dim=0: split (two-sided "
+                f"H^-1/2) scoring needs full, unprojected index gradients, but "
+                f"index_cfg.projection_dim={index_cfg.projection_dim}. Rebuild "
+                f"the index and query with projection_dim=0, or use a dense "
+                f"(autocorrelation) hessian."
+            )
         raise ValueError(
             f"Scoring with a factored (EKFAC) hessian at "
             f"{preprocess_cfg.hessian_path} requires projection_dim=0, but "

tests/ekfac_tests/test_factored_preconditioner.py

@@ -16,6 +16,7 @@
 import torch
 from safetensors.torch import load_file, save_file
 
+from bergson.collector.collector import create_projection_matrix
 from bergson.config import InversionConfig
 from bergson.data import create_index, load_module_gradients
 from bergson.hessians.apply_hessian import EkfacApplicator, EkfacConfig
@@ -239,3 +240,73 @@ def test_factored_apply_invariant_to_shard_count(tmp_path, inversion: str):
         assert torch.allclose(
             out1[name], out2[name], atol=1e-5
         ), f"{inversion}: apply on {name} depends on shard count"
+
+
+def test_apply_hessian_compresses_per_module(tmp_path):
+    """``EkfacConfig.projection_dim > 0`` compresses the IVHP output to a
+    ``[p, p]`` per-module Kronecker random projection, applied *after* H^-1
+    on the full gradient. It must use the exact same
+    ``create_projection_matrix`` convention (identifiers ``f"{name}/left"``
+    / ``f"{name}/right"``) that ``bergson build`` uses on training
+    gradients, so a compressed query is directly comparable to a compressed
+    training-side gradient store.
+    """
+    modules = {"a": (4, 6), "b": (5, 3)}  # (O, I)
+    hessian_path = tmp_path / "hessian"
+    _write_factored_hessian(hessian_path, modules, num_shards=1, seed=0)
+
+    grad_sizes = {name: o * i for name, (o, i) in modules.items()}
+    num_grads = 3
+    query_path = str(tmp_path / "query")
+    _make_query_gradients(query_path, grad_sizes, num_grads)
+
+    inversion_cfg = InversionConfig(inversion="damped_inverse", damping_factor=0.1)
+
+    # Reference: full, uncompressed IVHP output.
+    ref = _apply(
+        str(hessian_path),
+        query_path,
+        str(tmp_path / "out_full"),
+        "damped_inverse",
+        ev_correction=False,
+    )
+
+    p = 2
+    compressed_cfg = EkfacConfig(
+        hessian_method_path=str(hessian_path),
+        gradient_path=query_path,
+        run_path=str(tmp_path / "out_compressed"),
+        ev_correction=False,
+        projection_dim=p,
+    )
+    EkfacApplicator(compressed_cfg, inversion_cfg=inversion_cfg).compute_ivhp_sharded()
+    got = load_module_gradients(str(tmp_path / "out_compressed"))
+
+    for name, (o, i) in modules.items():
+        full = torch.from_numpy(np.asarray(ref[name][:])).view(num_grads, o, i)
+        P_l = create_projection_matrix(
+            f"{name}/left", p, o, full.dtype, full.device, "rademacher"
+        )
+        P_r = create_projection_matrix(
+            f"{name}/right", p, i, full.dtype, full.device, "rademacher"
+        )
+        expected = torch.einsum("ps,nsa,ra->npr", P_l, full, P_r)
+
+        compressed = torch.from_numpy(np.asarray(got[name][:])).view(num_grads, p, p)
+        assert compressed.shape == (num_grads, p, p)
+        assert torch.allclose(
+            compressed, expected, atol=1e-4, rtol=1e-4
+        ), f"{name}: compressed IVHP output doesn't match manual post-projection"
+
+
+def test_apply_hessian_rejects_compression_with_ev_correction():
+    """projection_dim compression is not supported with EK-FAC (ev_correction)."""
+    cfg = EkfacConfig(
+        hessian_method_path="unused",
+        gradient_path="unused",
+        run_path="unused",
+        ev_correction=True,
+        projection_dim=8,
+    )
+    with pytest.raises(ValueError, match="EK-FAC"):
+        EkfacApplicator(cfg, inversion_cfg=InversionConfig())

tests/ekfac_tests/test_hessian_pipeline.py

@@ -0,0 +1,23 @@
+import pytest
+
+from bergson.config import (
+    HessianConfig,
+    HessianPipelineConfig,
+    IndexConfig,
+    PreprocessConfig,
+    ScoreConfig,
+)
+from bergson.hessians.pipeline import hessian_pipeline
+
+
+def test_hessian_pipeline_rejects_unit_normalize(tmp_path):
+    """Cosine similarity (unit_normalize) is not supported with the
+    Kronecker-factored Hessians hessian_pipeline fits and applies."""
+    with pytest.raises(ValueError, match="unit_normalize"):
+        hessian_pipeline(
+            IndexConfig(run_path=str(tmp_path / "run")),
+            HessianConfig(method="kfac"),
+            ScoreConfig(),
+            PreprocessConfig(unit_normalize=True),
+            HessianPipelineConfig(),
+        )

tests/test_score.py

@@ -690,3 +690,21 @@ def test_score_factored_hessian_rejects_projection(tmp_path: Path):
             ScoreConfig(query_path=str(tmp_path / "q")),
             PreprocessConfig(hessian_path=hessian_path),
         )
+
+
+def test_score_factored_hessian_rejects_projection_with_unit_normalize(
+    tmp_path: Path,
+):
+    """Kronecker-factored hessian + projection + unit_normalize (cosine
+    similarity) together fail fast, with a message naming the actual cause."""
+    modules = {"mod_a": (4, 6)}
+    hessian_path = str(tmp_path / "hessian")
+    _write_factored_hessian(Path(hessian_path), modules)
+
+    index_cfg = IndexConfig(run_path=str(tmp_path / "out"), projection_dim=16)
+    with pytest.raises(ValueError, match="unit_normalize=True"):
+        score_dataset(
+            index_cfg,
+            ScoreConfig(query_path=str(tmp_path / "q")),
+            PreprocessConfig(hessian_path=hessian_path, unit_normalize=True),
+        )