path: root/experiments/ep_baseline.py

"""
Equilibrium Propagation (Scellier & Bengio 2017) for ResidualMLP on CIFAR-10.
Feedforward EP with energy-based state optimization.

Usage: python ep_baseline.py --method ep --seed 42 --gpu 0
"""
import os, sys, json, argparse, numpy as np, torch, torch.nn as nn, torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from models.residual_mlp import ResidualMLP
from metrics.credit_metrics import cosine_similarity_batch, perturbation_correlation
import torchvision, torchvision.transforms as transforms


def get_cifar10(bs=128):
    tt = transforms.Compose([transforms.RandomCrop(32, 4), transforms.RandomHorizontalFlip(),
                              transforms.ToTensor(), transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))])
    tv = transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))])
    return (DataLoader(torchvision.datasets.CIFAR10('./data', True, download=True, transform=tt), bs, True, num_workers=4, pin_memory=True),
            DataLoader(torchvision.datasets.CIFAR10('./data', False, download=True, transform=tv), bs, False, num_workers=4, pin_memory=True))


def evaluate(m, tl, dev):
    m.eval(); c, t = 0, 0
    with torch.no_grad():
        for x, y in tl:
            x = x.view(x.size(0), -1).to(dev); y = y.to(dev)
            c += (m(x).argmax(1) == y).sum().item(); t += x.size(0)
    return c / t


def ep_energy(model, hiddens, lam=1.0):
    """
    Compute the EP energy E = 0.5 * sum_l ||h_{l+1} - h_l - F_l(h_l)||^2
    hiddens: list of L+1 tensors, [h_0, h_1, ..., h_L]
    F_l(h_l) is the residual branch output (block forward without the skip).
    lam: weight for the state consistency term (kept at 1.0).
    """
    L = model.num_blocks
    E = 0.0
    for l in range(L):
        f_l = model.blocks[l](hiddens[l])  # residual branch
        residual = hiddens[l + 1] - hiddens[l] - f_l
        E = E + 0.5 * (residual ** 2).sum(-1)  # (batch,)
    return E  # (batch,)


def ep_free_phase(model, x):
    """
    Free phase: standard forward pass. Returns hidden states h_0..h_L.
    """
    with torch.no_grad():
        _, hiddens = model(x, return_hidden=True)
    return hiddens  # list of L+1 tensors


def ep_nudged_phase(model, x, y, h_free, beta, T_nudge, alpha_nudge):
    """
    Nudged phase: minimize E(h) + beta * C(h_L, y) w.r.t. hidden states h_1..h_L.
    h_0 is fixed (output of embed layer).
    Returns list of nudged hidden states [h_0, h_1^*, ..., h_L^*].
    """
    L = model.num_blocks
    # Initialize nudged states from free phase (detached)
    h_nudged = [h.clone().detach() for h in h_free]
    # h_0 is fixed (embed output)
    h_nudged[0] = h_free[0].clone().detach()
    # Optimize h_1 .. h_L
    for i in range(1, L + 1):
        h_nudged[i].requires_grad_(True)

    params_to_opt = h_nudged[1:]
    inner_opt = optim.SGD(params_to_opt, lr=alpha_nudge)

    for _ in range(T_nudge):
        # Energy over all layers
        E = ep_energy(model, h_nudged)  # (batch,)
        # Cost at output layer: cross-entropy
        logits = model.out_head(model.out_ln(h_nudged[L]))
        C = F.cross_entropy(logits, y, reduction='none')  # (batch,)
        total = (E + beta * C).mean()
        inner_opt.zero_grad()
        total.backward()
        inner_opt.step()

    return [h.detach() for h in h_nudged]



def train_ep(model, trl, tel, dev, epochs=100, lr=1e-3, wd=0.01,
             beta=0.5, T_nudge=20, alpha_nudge=0.1, random_targets: bool = False):
    L = model.num_blocks

    # Separate optimizers for different parts
    block_opts = [optim.AdamW(b.parameters(), lr=lr, weight_decay=wd) for b in model.blocks]
    embed_opt = optim.AdamW(model.embed.parameters(), lr=lr, weight_decay=wd)
    head_opt = optim.AdamW(list(model.out_head.parameters()) + list(model.out_ln.parameters()), lr=lr, weight_decay=wd)
    all_opts = block_opts + [embed_opt, head_opt]
    schedulers = [optim.lr_scheduler.CosineAnnealingLR(o, T_max=epochs) for o in all_opts]

    for ep in range(1, epochs + 1):
        model.train()
        for x, y in trl:
            x = x.view(x.size(0), -1).to(dev); y = y.to(dev)
            if random_targets:
                y = torch.randint(0, 10, y.shape, device=dev)

            # ---- FREE PHASE ----
            # Standard forward pass to get free fixed point
            with torch.no_grad():
                _, h_free = model(x, return_hidden=True)

            # ---- NUDGED PHASE ----
            # Minimize E(h) + beta * C(h_L, y) w.r.t. hidden states
            h_nudged = ep_nudged_phase(model, x, y, h_free, beta, T_nudge, alpha_nudge)

            # ---- EP WEIGHT UPDATE ----
            # Δθ ∝ (∂E_nudged/∂θ - ∂E_free/∂θ) / beta
            # For blocks: dE/dθ_l comes from F_l(h_l) term in E
            # For embed: dE/dθ_embed comes from h_0 = embed(x) being the base state

            for o in all_opts:
                o.zero_grad()

            # Compute EP grads for residual blocks
            # E = sum_l 0.5 ||h_{l+1} - h_l - F_l(h_l)||^2
            # dE/dθ_l = - (h_{l+1} - h_l - F_l(h_l))^T * dF_l/dθ_l
            #          = - residual_l^T * dF_l/dθ_l

            for l in range(L):
                h_l_free = h_free[l].detach()
                h_lp1_free = h_free[l + 1].detach()
                h_l_nudge = h_nudged[l].detach()
                h_lp1_nudge = h_nudged[l + 1].detach()

                # Free phase: -residual_l_free dot dF_l/dtheta
                # = -(h_{l+1}^free - h_l^free - F_l(h_l^free)) dot dF_l/dtheta
                # We compute this by forward pass with grad for params
                f_l_free = model.blocks[l](h_l_free)
                res_free = h_lp1_free - h_l_free - f_l_free.detach()
                # Gradient: d/dtheta [ -0.5 * res_free^2 ] = res_free * dF/dtheta ... actually we want
                # to minimize E, so grad = dE/dtheta = -res * dF/dtheta
                # To use autograd: compute -res_free.detach() * f_l_free, sum, backward
                loss_free_l = -(res_free.detach() * f_l_free).sum()

                f_l_nudge = model.blocks[l](h_l_nudge)
                res_nudge = h_lp1_nudge - h_l_nudge - f_l_nudge.detach()
                loss_nudge_l = -(res_nudge.detach() * f_l_nudge).sum()

                # EP grad = (nudged - free) / beta  [we want d(E_nudge - E_free)/dtheta / beta]
                # Since loss_free_l = -res_free * F_l contributes dE/dtheta_free (the negative),
                # and loss_nudge_l similarly, we need:
                # grad_l = (dE_nudge/dtheta - dE_free/dtheta) / beta
                # dE/dtheta = -res * dF/dtheta => computed via backward of (res * F).sum()
                # So: ep_loss = (loss_nudge_l - loss_free_l) / beta
                ep_loss_l = (loss_nudge_l - loss_free_l) / beta
                ep_loss_l.backward()

            # Grad for embed layer:
            # h_0 = embed(x), so dE/dtheta_embed = dE/dh_0 * dh_0/dtheta_embed
            # dE/dh_0: E depends on h_0 via (h_1 - h_0 - F_0(h_0)) term
            # = -(h_1 - h_0 - F_0(h_0)) * (I + dF_0/dh_0)^T ... complex
            # Simpler: treat h_0 as part of the system and use chain rule via autograd
            h0_free = model.embed(x)  # differentiable
            # compute E contribution from layer 0 with h_0 = embed(x) fixed, block params fixed
            with torch.no_grad():
                f0_free = model.blocks[0](h0_free.detach())
                res0_free = h_free[1].detach() - h0_free.detach() - f0_free
            embed_loss_free = -(res0_free.detach() * h0_free).sum() / beta  # approximation: -res * dh0/dtheta

            h0_nudge_rg = model.embed(x)
            with torch.no_grad():
                f0_nudge = model.blocks[0](h_nudged[0].detach())
                res0_nudge = h_nudged[1].detach() - h_nudged[0].detach() - f0_nudge
            embed_loss_nudge = -(res0_nudge.detach() * h0_nudge_rg).sum() / beta

            embed_ep = (embed_loss_nudge - embed_loss_free)
            embed_ep.backward()

            # Grad for out_head + out_ln: standard BP at nudged output
            # EP: Δθ_out = ∂C_nudged/∂θ_out (since ∂E/∂θ_out = 0)
            # This is equivalent to standard BP loss at nudged hidden state
            logits_nudged = model.out_head(model.out_ln(h_nudged[L].detach()))
            head_loss = F.cross_entropy(logits_nudged, y)
            head_loss.backward()

            # Clip and step
            all_params = list(model.parameters())
            torch.nn.utils.clip_grad_norm_(all_params, 1.0)
            for o in all_opts:
                o.step()

        for s in schedulers:
            s.step()
        if ep % 20 == 0:
            print(f"  Ep {ep}: acc={evaluate(model, tel, dev):.4f}", flush=True)

    return model


def ep_credit_signals(model, x, y, beta, T_nudge, alpha_nudge):
    """
    Compute EP credit signals a_l^EP = (h_l^nudged - h_l^free) / beta for diagnostics.
    """
    with torch.no_grad():
        _, h_free = model(x, return_hidden=True)
    h_nudged = ep_nudged_phase(model, x, y, h_free, beta, T_nudge, alpha_nudge)
    L = model.num_blocks
    # Negate: EP nudge moves h toward lower loss, opposite to BP grad direction
    credits = [-(h_nudged[l] - h_free[l]) / beta for l in range(L)]
    return credits, h_free, h_nudged


def compute_diagnostics(model, tel, dev, beta, T_nudge, alpha_nudge):
    model.eval()
    L = model.num_blocks

    for x, y in tel:
        x = x.view(x.size(0), -1).to(dev); y = y.to(dev)
        break

    # EP credit signals
    ep_credits, h_free, h_nudged = ep_credit_signals(model, x, y, beta, T_nudge, alpha_nudge)

    # BP gradients for comparison
    h0 = model.embed(x.detach())
    hs = [h0.clone().requires_grad_(True)]
    for bl in model.blocks:
        hs.append(hs[-1] + bl(hs[-1]))
    lo = model.out_head(model.out_ln(hs[-1]))
    loss = F.cross_entropy(lo, y)
    gs = torch.autograd.grad(loss, hs)
    bp_grads = {l: gs[l].detach() for l in range(L)}

    # Gamma: cosine sim between EP credit and BP grad
    gammas = []
    for l in range(L):
        g = cosine_similarity_batch(ep_credits[l], bp_grads[l])
        gammas.append(g)

    # rho: perturbation correlation using EP credit
    rhos = []
    with torch.no_grad():
        _, hi = model(x, return_hidden=True)

    for l in range(L):
        h_l = hi[l].detach()
        a_l = ep_credits[l].detach()

        def mk(sl):
            def f(h):
                with torch.no_grad():
                    c = h
                    for i in range(sl, L):
                        c = c + model.blocks[i](c)
                    return F.cross_entropy(model.out_head(model.out_ln(c)), y, reduction='none')
            return f

        rhos.append(perturbation_correlation(h_l, a_l, mk(l), epsilon=1e-3, M=16))

    # naive state error
    with torch.no_grad():
        _, hi2 = model(x, return_hidden=True)
        nse = ((hi2[L // 2] - hi2[-1]).norm(-1) / hi2[-1].norm(-1).clamp(min=1e-8)).mean().item()

    return {'Gamma': float(np.mean(gammas)), 'rho': float(np.mean(rhos)),
            'naive_StateErr': nse, 'gammas_per_layer': [float(g) for g in gammas],
            'rhos_per_layer': [float(r) for r in rhos]}


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--method', type=str, default='ep')
    p.add_argument('--seed', type=int, required=True)
    p.add_argument('--gpu', type=int, default=0)
    p.add_argument('--output_dir', type=str, default='results/ep_baseline')
    p.add_argument('--epochs', type=int, default=100)
    p.add_argument('--beta', type=float, default=0.5, help='EP nudge strength')
    p.add_argument('--T_nudge', type=int, default=20, help='Inner optimization steps for nudged phase')
    p.add_argument('--alpha_nudge', type=float, default=0.1, help='Inner step size for nudged phase')
    p.add_argument('--lr', type=float, default=1e-3)
    p.add_argument('--wd', type=float, default=0.01)
    p.add_argument('--d_hidden', type=int, default=256)
    p.add_argument('--random_targets', action='store_true',
                   help='Replace each minibatch label with i.i.d. random class targets (codex round 36 OPTION EP).')
    args = p.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    dev = torch.device(f'cuda:{args.gpu}')
    torch.manual_seed(args.seed); np.random.seed(args.seed); torch.cuda.manual_seed_all(args.seed)

    trl, tel = get_cifar10()
    L, d = 4, args.d_hidden
    model = ResidualMLP(3072, d, 10, L).to(dev)

    print(f"[{args.method} s={args.seed}] Training EP  beta={args.beta} T={args.T_nudge} alpha={args.alpha_nudge}", flush=True)

    model = train_ep(model, trl, tel, dev, epochs=args.epochs, lr=args.lr, wd=args.wd,
                     beta=args.beta, T_nudge=args.T_nudge, alpha_nudge=args.alpha_nudge,
                     random_targets=args.random_targets)

    acc = evaluate(model, tel, dev)
    diag = compute_diagnostics(model, tel, dev, args.beta, args.T_nudge, args.alpha_nudge)

    torch.save(model.state_dict(), os.path.join(args.output_dir, f'{args.method}_s{args.seed}.pt'))

    result = {'method': args.method, 'seed': args.seed, 'acc': acc,
              'Gamma': diag['Gamma'], 'rho': diag['rho'],
              'naive_StateErr': diag['naive_StateErr'],
              'gammas_per_layer': diag['gammas_per_layer'],
              'rhos_per_layer': diag['rhos_per_layer'],
              'beta': args.beta, 'T_nudge': args.T_nudge, 'alpha_nudge': args.alpha_nudge}

    with open(os.path.join(args.output_dir, f'{args.method}_s{args.seed}.json'), 'w') as f:
        json.dump(result, f, indent=2, default=float)

    print(f"[{args.method} s={args.seed}] acc={acc:.4f} Γ={diag['Gamma']:.4f} ρ={diag['rho']:.4f} nse={diag['naive_StateErr']:.4f}", flush=True)


if __name__ == '__main__':
    main()