repro-bundle-v0/src/hyper_cleaning.py

"""Distributed data hyper-cleaning on Fashion-MNIST (paper Sec 4.2 / Appendix
C.4), reduced scale: a small 1-hidden-layer MLP trained across n agents, each
holding a locally label-corrupted training shard plus a small clean
validation shard. x = per-training-example reweighting logits (local to the
owning agent's block, dimension = total training examples across all
agents); y = shared MLP weights.

g_i(x, y) = (1/m_i) sum_j sigmoid(x_j) * CE(y; example_j) + (mu/2)||y||^2   [weighted train loss + L2]
f_i(x, y) = (1/p_i) sum_j CE(y; clean_val_example_j)                        [clean validation loss]

Matches the paper's own finding (Sec 4.3, Fig. 7): each agent's x-gradient is
zero outside its own block, so FAB only needs to *track* (not locally
compute) the other agents' slices -- consensus on those blocks is trivial
since nobody else ever pushes a nonzero gradient into them.
"""
from __future__ import annotations

import numpy as np


def _softmax(z):
    z = z - z.max(axis=1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=1, keepdims=True)


class MLP:
    """One hidden layer, ReLU, softmax+cross-entropy. y is a flat parameter vector."""

    def __init__(self, in_dim, hidden, out_dim, seed=0):
        rng = np.random.default_rng(seed)
        self.in_dim, self.hidden, self.out_dim = in_dim, hidden, out_dim
        s1 = np.sqrt(2.0 / in_dim)
        s2 = np.sqrt(2.0 / hidden)
        w1 = rng.normal(scale=s1, size=(in_dim, hidden))
        b1 = np.zeros(hidden)
        w2 = rng.normal(scale=s2, size=(hidden, out_dim))
        b2 = np.zeros(out_dim)
        self.shapes = [w1.shape, b1.shape, w2.shape, b2.shape]
        self.sizes = [int(np.prod(s)) for s in self.shapes]
        self.dim = sum(self.sizes)
        self.init = self.pack(w1, b1, w2, b2)

    def pack(self, w1, b1, w2, b2):
        return np.concatenate([w1.ravel(), b1.ravel(), w2.ravel(), b2.ravel()])

    def unpack(self, y):
        i = 0
        parts = []
        for shape, size in zip(self.shapes, self.sizes):
            parts.append(y[i:i + size].reshape(shape))
            i += size
        return parts  # w1, b1, w2, b2

    def forward(self, y, X):
        w1, b1, w2, b2 = self.unpack(y)
        h_pre = X @ w1 + b1
        h = np.maximum(h_pre, 0.0)
        logits = h @ w2 + b2
        probs = _softmax(logits)
        cache = (X, h_pre, h, w1, w2)
        return probs, cache

    def per_example_ce(self, y, X, labels):
        probs, cache = self.forward(y, X)
        p_true = probs[np.arange(len(labels)), labels]
        return -np.log(np.clip(p_true, 1e-12, None)), probs, cache

    def weighted_loss_and_grad(self, y, X, labels, weights, l2=0.0):
        """Returns (scalar loss, grad wrt y, per-example d(loss)/d(weight_j))."""
        ce, probs, (X_, h_pre, h, w1, w2) = self.per_example_ce(y, X, labels)
        m = len(labels)
        loss = float(np.sum(weights * ce) / m + 0.5 * l2 * np.sum(y ** 2))

        onehot = np.zeros_like(probs)
        onehot[np.arange(m), labels] = 1.0
        dlogits = (probs - onehot) * weights[:, None] / m  # weighted softmax-CE grad
        dw2 = h.T @ dlogits
        db2 = dlogits.sum(axis=0)
        dh = dlogits @ w2.T
        dh_pre = dh * (h_pre > 0)
        dw1 = X_.T @ dh_pre
        db1 = dh_pre.sum(axis=0)
        grad_y = self.pack(dw1, db1, dw2, db2) + l2 * y
        grad_weights = ce / m  # d(loss)/d(weight_j) = ce_j / m
        return loss, grad_y, grad_weights

    def unweighted_loss_and_grad(self, y, X, labels, l2=0.0):
        return self.weighted_loss_and_grad(y, X, labels, np.ones(len(labels)), l2=l2)[:2]

    def accuracy(self, y, X, labels):
        probs, _ = self.forward(y, X)
        preds = probs.argmax(axis=1)
        return float(np.mean(preds == labels))


def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-np.clip(z, -30, 30)))


def sigmoid_grad(z):
    s = sigmoid(z)
    return s * (1 - s)


class DistributedHyperCleaning:
    """n agents; agent i owns training slice [offsets[i]:offsets[i+1]] of the
    global per-example weight vector x (dimension = total training examples).
    y is the shared MLP parameter vector (same dimension for every agent).
    """

    def __init__(self, mlp: MLP, x_train_shards, y_train_shards, x_val_shards, y_val_shards,
                 mu=1e-3):
        self.mlp = mlp
        self.n = len(x_train_shards)
        self.x_train = x_train_shards
        self.y_train = y_train_shards
        self.x_val = x_val_shards
        self.y_val = y_val_shards
        self.mu = mu
        sizes = [len(s) for s in y_train_shards]
        self.offsets = np.concatenate([[0], np.cumsum(sizes)])
        self.d_x = int(self.offsets[-1])
        self.d_y = mlp.dim

    def my_slice(self, i):
        return slice(self.offsets[i], self.offsets[i + 1])

    def grad_x_full_and_grad_y(self, i, x_full, y):
        """Local gradients of L_i-relevant pieces at agent i: returns the FULL
        d_x-dim gradient (nonzero only on agent i's own slice) and the d_y-dim
        gradient of g_i wrt y, plus g_i's value and weighted-train loss pieces
        needed by the caller to assemble L_i = f_i + lambda*(g_i(x,y)-g_i(x,z)).
        """
        sl = self.my_slice(i)
        w = sigmoid(x_full[sl])
        _, grad_y_g, grad_w = self.mlp.weighted_loss_and_grad(
            y, self.x_train[i], self.y_train[i], w, l2=self.mu)
        grad_x_full = np.zeros(self.d_x)
        grad_x_full[sl] = grad_w * sigmoid_grad(x_full[sl])
        return grad_x_full, grad_y_g

    def grad_y_f_i(self, i, y):
        _, grad_y_f = self.mlp.unweighted_loss_and_grad(y, self.x_val[i], self.y_val[i], l2=0.0)
        return grad_y_f

    def val_loss_i(self, i, y):
        ce, _, _ = self.mlp.per_example_ce(y, self.x_val[i], self.y_val[i])
        return float(np.mean(ce))