File size: 3,229 Bytes
b88c26d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | #!/usr/bin/env python3
"""
MMEP Gradient vs Backprop Gradient Comparison
Computes gradients via both MMEP (free/nudged phase difference) and
standard backpropagation, then reports cosine similarity and L2 distance.
Usage:
python scripts/gradient_comparison.py --hidden-dim 256 --num-layers 4
Contact: jessica@collectivekitty.com
"""
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
def compute_mmep_gradient(model, x, target, T_free=20, T_nudge=4,
alpha=0.5, beta=0.1):
"""Compute EP gradient via free/nudged phase difference."""
# Free phase
h = x.clone()
for _ in range(T_free):
h = (1 - alpha) * h + alpha * torch.tanh(model(h))
h_free = h.detach().clone()
# Nudged phase
h = h_free.clone()
for _ in range(T_nudge):
h = (1 - alpha) * h + alpha * torch.tanh(model(h))
h = h + beta * (target - h)
h_nudged = h.detach().clone()
# EP gradient: (1/beta) * (correlation_nudged - correlation_free)
# For weight gradient: dW ~ (1/beta) * (h_nudged @ x^T - h_free @ x^T)
grad_ep = (1.0 / beta) * (h_nudged - h_free).mean(dim=0)
return grad_ep
def compute_backprop_gradient(model, x, target):
"""Compute standard backprop gradient."""
model.zero_grad()
output = model(x)
loss = F.mse_loss(output, target)
loss.backward()
# Return gradient of first parameter as representative
for p in model.parameters():
if p.grad is not None:
return p.grad.flatten()[:target.shape[-1]]
return torch.zeros(target.shape[-1])
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--hidden-dim", type=int, default=64)
parser.add_argument("--num-layers", type=int, default=2)
parser.add_argument("--T-free", type=int, default=20)
parser.add_argument("--T-nudge", type=int, default=4)
parser.add_argument("--beta", type=float, default=0.1)
parser.add_argument("--num-samples", type=int, default=100)
args = parser.parse_args()
d = args.hidden_dim
model = nn.Linear(d, d, bias=False)
torch.manual_seed(42)
cos_sims = []
for _ in range(args.num_samples):
x = torch.randn(8, d)
target = torch.randn(8, d)
grad_ep = compute_mmep_gradient(model, x, target,
T_free=args.T_free, T_nudge=args.T_nudge,
beta=args.beta)
grad_bp = compute_backprop_gradient(model, x, target)
# Align dimensions
min_dim = min(grad_ep.shape[0], grad_bp.shape[0])
cos_sim = F.cosine_similarity(
grad_ep[:min_dim].unsqueeze(0),
grad_bp[:min_dim].unsqueeze(0)
).item()
cos_sims.append(cos_sim)
avg_cos = sum(cos_sims) / len(cos_sims)
print(f"Gradient Comparison (d={d}, L={args.num_layers})")
print(f" Mean cosine similarity: {avg_cos:.4f}")
print(f" Min: {min(cos_sims):.4f}, Max: {max(cos_sims):.4f}")
print(f" (1.0 = perfect alignment, expected to improve with more T_free)")
if __name__ == "__main__":
main()
|