Initial submission code: FA evaluation protocol + reproduction scripts

Reference implementation of the three-diagnostic FA evaluation protocol
(scale stability, reference validity, depth utility) from the NeurIPS 2026
E&D track paper. Includes models, metrics, and full reproduction pipeline.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

3 files changed, 321 insertions, 0 deletions

diff --git a/protocol/__init__.py b/protocol/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/protocol/__init__.py
diff --git a/protocol/example_usage.py b/protocol/example_usage.py
new file mode 100644
index 0000000..2b2c65e
--- /dev/null
+++ b/protocol/example_usage.py
@@ -0,0 +1,106 @@
+"""
+Minimal example: apply the FA evaluation protocol to a DFA-trained ResMLP.
+
+This script trains a model with DFA, then runs the three-diagnostic protocol.
+Expected output: FAIL(D1+D2+D3) — DFA on terminal-LN ResMLP triggers all diagnostics.
+"""
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import torchvision
+import torchvision.transforms as transforms
+import numpy as np
+
+from models.residual_mlp import ResidualMLP
+from protocol.fa_protocol import FAProtocol
+
+
+def get_cifar10(batch_size=128):
+    tv = transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
+    ])
+    tv_train = transforms.Compose([
+        transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616)),
+    ])
+    tr = torchvision.datasets.CIFAR10('./data', True, download=True, transform=tv_train)
+    te = torchvision.datasets.CIFAR10('./data', False, download=True, transform=tv)
+    return (torch.utils.data.DataLoader(tr, batch_size=batch_size, shuffle=True, num_workers=2),
+            torch.utils.data.DataLoader(te, batch_size=batch_size, shuffle=False, num_workers=2))
+
+
+def train_dfa(model, train_loader, device, epochs=30):
+    """Minimal DFA training (canonical: no clipping, mean reduction)."""
+    d = model.d_hidden
+    L = model.num_blocks
+    C = 10
+    Bs = [torch.randn(d, C, device=device) / np.sqrt(C) for _ in range(L)]
+    block_opts = [optim.AdamW(b.parameters(), lr=1e-3, weight_decay=0.01) for b in model.blocks]
+    embed_opt = optim.AdamW(model.embed.parameters(), lr=1e-3, weight_decay=0.01)
+    head_opt = optim.AdamW(list(model.out_head.parameters()) + list(model.out_ln.parameters()),
+                           lr=1e-3, weight_decay=0.01)
+    for ep in range(1, epochs + 1):
+        model.train()
+        for x, y in train_loader:
+            x = x.view(x.size(0), -1).to(device); y = y.to(device)
+            batch = x.size(0)
+            with torch.no_grad():
+                logits, hiddens = model(x, return_hidden=True)
+                e_T = logits.softmax(-1); e_T[torch.arange(batch), y] -= 1
+            hL = hiddens[-1].detach()
+            head_opt.zero_grad()
+            F.cross_entropy(model.out_head(model.out_ln(hL)), y).backward()
+            head_opt.step()
+            for l in range(L):
+                a = (e_T @ Bs[l].T).detach()
+                rms = (a ** 2).mean(-1, keepdim=True).sqrt() + 1e-6
+                f_l = model.blocks[l](hiddens[l].detach())
+                loss = (f_l * (a / rms)).sum(-1).mean()
+                block_opts[l].zero_grad(); loss.backward(); block_opts[l].step()
+            a0 = (e_T @ Bs[0].T).detach()
+            rms0 = (a0 ** 2).mean(-1, keepdim=True).sqrt() + 1e-6
+            h0 = model.embed(x)
+            embed_opt.zero_grad(); (h0 * (a0 / rms0)).sum(-1).mean().backward(); embed_opt.step()
+        if ep % 10 == 0:
+            print(f"  DFA ep {ep}/{epochs}", flush=True)
+
+
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    seed = 42
+    torch.manual_seed(seed); np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+    print("Loading CIFAR-10...")
+    train_loader, test_loader = get_cifar10()
+
+    # Prepare eval buffer
+    xs, ys = [], []
+    for x, y in test_loader:
+        xs.append(x.view(x.size(0), -1)); ys.append(y)
+        if sum(xb.size(0) for xb in xs) >= 128:
+            break
+    x_eval = torch.cat(xs)[:128].to(device)
+    y_eval = torch.cat(ys)[:128].to(device)
+
+    # Train with DFA
+    print("Training DFA (30 epochs)...")
+    model = ResidualMLP(3072, 256, 10, 4).to(device)
+    train_dfa(model, train_loader, device, epochs=30)
+
+    # Run protocol
+    print("\nRunning protocol...")
+    protocol = FAProtocol(model, x_eval, y_eval)
+    report = protocol.run(frozen_baseline_acc=0.349)
+    print(protocol.summary(report))
+
+
+if __name__ == '__main__':
+    main()
diff --git a/protocol/fa_protocol.py b/protocol/fa_protocol.py
new file mode 100644
index 0000000..8d12939
--- /dev/null
+++ b/protocol/fa_protocol.py
@@ -0,0 +1,215 @@
+"""
+Reference implementation of the three-diagnostic FA evaluation protocol.
+
+Usage:
+    from protocol.fa_protocol import FAProtocol
+
+    protocol = FAProtocol(model, x_eval, y_eval)
+    report = protocol.run(frozen_baseline_acc=0.349)
+    print(report['verdict'])
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from typing import Optional
+
+
+class FAProtocol:
+    """Three-diagnostic evaluation protocol for feedback alignment methods.
+
+    Diagnostics:
+        D1 (Scale stability): max per-block residual growth rho = max_l ||h_{l+1}|| / ||h_l||.
+            Flags if rho > threshold (default 50).
+        D2 (Reference validity): BP gradient norm at the deepest hidden state.
+            Flags if ||g_L|| < 10 * eps, where eps is the cosine clamp floor.
+        D3 (Depth utility): test accuracy vs frozen-blocks baseline.
+            Flags if trained acc < frozen_acc + margin (default 2 pp).
+
+    The protocol requires:
+        - A trained model with .blocks (nn.ModuleList) and forward(x, return_hidden=True)
+        - A test batch (x_eval, y_eval)
+        - A frozen-blocks baseline accuracy (must be computed separately)
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        x_eval: torch.Tensor,
+        y_eval: torch.Tensor,
+        d1_threshold: float = 50.0,
+        d2_eps: float = 1e-8,
+        d2_factor: float = 10.0,
+        d3_margin_pp: float = 2.0,
+    ):
+        self.model = model
+        self.x_eval = x_eval
+        self.y_eval = y_eval
+        self.d1_threshold = d1_threshold
+        self.d2_floor = d2_factor * d2_eps
+        self.d3_margin = d3_margin_pp / 100.0
+
+    def _compute_hidden_norms(self, hiddens):
+        """Compute median per-sample L2 norm at each hidden layer."""
+        norms = []
+        for h in hiddens:
+            if h.dim() == 4:  # conv: (B, C, H, W) -> pool to (B, C)
+                h_flat = F.adaptive_avg_pool2d(h, 1).flatten(1)
+            elif h.dim() == 3:  # transformer: (B, T, D) -> cls token
+                h_flat = h[:, 0]
+            else:
+                h_flat = h
+            norms.append(float(h_flat.norm(dim=-1).median().item()))
+        return norms
+
+    def _compute_bp_grad_norms(self, hiddens):
+        """Compute BP gradient norms at each hidden layer via manual forward."""
+        model = self.model
+        L = len(hiddens) - 1  # number of blocks
+
+        # Rebuild forward from hidden states with grad tracking
+        hs = [hiddens[0].detach().clone().requires_grad_(True)]
+        for i, block in enumerate(model.blocks):
+            if hasattr(block, 'forward'):
+                h_next = block(hs[-1])
+                # Check if block includes residual (output same shape, skip connection)
+                if h_next.shape == hs[-1].shape and not self._block_has_internal_skip(block):
+                    h_next = hs[-1] + h_next
+            hs.append(h_next)
+
+        # Forward through head
+        h_final = hs[-1]
+        if h_final.dim() == 4:  # conv
+            h_final = F.adaptive_avg_pool2d(h_final, 1).flatten(1)
+        elif h_final.dim() == 3:  # transformer cls token
+            h_final = h_final[:, 0]
+        if hasattr(model, 'out_ln'):
+            h_final = model.out_ln(h_final)
+        logits = model.out_head(h_final)
+        loss = F.cross_entropy(logits, self.y_eval)
+        grads = torch.autograd.grad(loss, hs, allow_unused=True)
+
+        norms = []
+        for g in grads:
+            if g is None:
+                norms.append(0.0)
+                continue
+            if g.dim() == 4:
+                g_flat = F.adaptive_avg_pool2d(g, 1).flatten(1)
+            elif g.dim() == 3:
+                g_flat = g[:, 0]
+            else:
+                g_flat = g
+            norms.append(float(g_flat.norm(dim=-1).median().item()))
+        return norms
+
+    @staticmethod
+    def _block_has_internal_skip(block):
+        """Heuristic: check if the block's forward already includes a residual skip."""
+        src = type(block).forward.__qualname__
+        # Blocks that compute x + f(x) internally (e.g., transformer blocks)
+        return False  # conservative default; override if needed
+
+    def run(self, frozen_baseline_acc: Optional[float] = None, test_acc: Optional[float] = None):
+        """Run all three diagnostics.
+
+        Args:
+            frozen_baseline_acc: accuracy of the frozen-blocks baseline (required for D3).
+            test_acc: test accuracy of the trained model. If None, computed from x_eval/y_eval.
+
+        Returns:
+            dict with 'diagnostics', 'verdict', and raw values.
+        """
+        self.model.eval()
+
+        # Forward pass to get hidden states
+        with torch.no_grad():
+            logits, hiddens = self.model(self.x_eval, return_hidden=True)
+
+        if test_acc is None:
+            test_acc = float((logits.argmax(-1) == self.y_eval).float().mean().item())
+
+        # D1: Scale stability
+        h_norms = self._compute_hidden_norms(hiddens)
+        growth_ratios = [h_norms[i+1] / max(h_norms[i], 1e-12)
+                         for i in range(len(h_norms) - 1)]
+        max_growth = max(growth_ratios) if growth_ratios else 1.0
+        d1_fires = max_growth > self.d1_threshold
+
+        # D2: Reference validity
+        bp_grad_norms = self._compute_bp_grad_norms(hiddens)
+        g_L = bp_grad_norms[-1] if bp_grad_norms else 0.0
+        d2_fires = g_L < self.d2_floor
+
+        # D3: Depth utility
+        if frozen_baseline_acc is not None:
+            margin = test_acc - frozen_baseline_acc
+            d3_fires = margin < self.d3_margin
+        else:
+            margin = None
+            d3_fires = None
+
+        # Verdict
+        mode1 = d1_fires and d2_fires
+        flags = []
+        if d1_fires:
+            flags.append('D1')
+        if d2_fires:
+            flags.append('D2')
+        if d3_fires:
+            flags.append('D3')
+
+        if not flags:
+            verdict = 'PASS'
+        else:
+            verdict = 'FAIL(' + '+'.join(flags) + ')'
+
+        return {
+            'verdict': verdict,
+            'test_acc': test_acc,
+            'diagnostics': {
+                'D1_scale_growth': {
+                    'max_growth': max_growth,
+                    'per_block_growth': growth_ratios,
+                    'hidden_norms': h_norms,
+                    'threshold': self.d1_threshold,
+                    'fires': d1_fires,
+                },
+                'D2_ref_validity': {
+                    'g_L': g_L,
+                    'bp_grad_norms': bp_grad_norms,
+                    'floor': self.d2_floor,
+                    'fires': d2_fires,
+                },
+                'D3_depth_utility': {
+                    'test_acc': test_acc,
+                    'frozen_baseline_acc': frozen_baseline_acc,
+                    'margin': margin,
+                    'margin_threshold': self.d3_margin,
+                    'fires': d3_fires,
+                },
+            },
+        }
+
+    def summary(self, report: dict) -> str:
+        """Human-readable summary of a protocol report."""
+        d = report['diagnostics']
+        lines = [
+            f"Verdict: {report['verdict']}",
+            f"Test accuracy: {report['test_acc']:.4f}",
+            f"D1 Scale stability: max growth = {d['D1_scale_growth']['max_growth']:.1f}x "
+            f"(threshold {d['D1_scale_growth']['threshold']}x) -> "
+            f"{'FIRE' if d['D1_scale_growth']['fires'] else 'pass'}",
+            f"D2 Reference validity: ||g_L|| = {d['D2_ref_validity']['g_L']:.2e} "
+            f"(floor {d['D2_ref_validity']['floor']:.0e}) -> "
+            f"{'FIRE' if d['D2_ref_validity']['fires'] else 'pass'}",
+        ]
+        if d['D3_depth_utility']['fires'] is not None:
+            lines.append(
+                f"D3 Depth utility: margin = {d['D3_depth_utility']['margin']*100:+.1f} pp "
+                f"(threshold {d['D3_depth_utility']['margin_threshold']*100:.0f} pp) -> "
+                f"{'FIRE' if d['D3_depth_utility']['fires'] else 'pass'}"
+            )
+        else:
+            lines.append("D3 Depth utility: not evaluated (no frozen baseline provided)")
+        return '\n'.join(lines)