Consensus ADMM with Gaussian Noise (Demonstration)

Introduction

This page is a demonstration. It fits an L2-regularized logistic regression across \(N\) sites by consensus ADMM under threshold FHE, the protocol the Cox-lasso page uses. Each site solves its local problem with cvxpy, and the aggregator decrypts only the average of the sites’ encrypted vectors.

The problem and its consensus split

With \(N\) sites, local data \((X_i, y_i)\) at site \(i\), and a shared coefficient \(x \in \mathbb{R}^p\), the global problem is

\[ \min_{x \in \mathbb{R}^p} \sum_{i=1}^{N} \ell_i(x; X_i, y_i) + \frac{\lambda}{2}\, \lVert x \rVert_2^2 \]

where \(\ell_i\) is the logistic loss on site \(i\). The standard consensus split (Boyd et al. 2011) introduces local copies \(x_i \in \mathbb{R}^p\) and a single global consensus \(z \in \mathbb{R}^p\):

\[ \min_{\{x_i\}, z} \sum_i \ell_i(x_i; X_i, y_i) + \frac{\lambda}{2}\, \lVert z \rVert_2^2 \quad \text{s.t.}\quad x_i = z, \ \forall i. \]

The augmented-Lagrangian iteration is:

\[ \begin{aligned} x_i^{k+1} &= \arg\min_{x_i}\, \ell_i(x_i) + \frac{\lambda}{2N}\lVert x_i\rVert_2^2 + \frac{\rho}{2}\lVert x_i - z^k + u_i^k\rVert_2^2, \\ z^{k+1} &= \frac{1}{N}\sum_i (x_i^{k+1} + u_i^k), \\ u_i^{k+1} &= u_i^k + (x_i^{k+1} - z^{k+1}). \end{aligned} \]

The \(x\)-update is local at each site; the \(z\)-update is the consensus average that has to traverse the encrypted channel; the \(u\)-update is local again. Only the \(z\)-update needs cryptography.

Adding noise

Here we ask what happens if each site also adds Gaussian noise of the kind an output differential privacy mechanism uses. The sensitivity \(\Delta = 1\) below is a placeholder, so the \(\varepsilon\) values are not a privacy guarantee. Adding the noise takes one extra rng.normal() call per site per iteration. With this optimizer, the fits are poor at any noise level that gives a small \(\varepsilon\). The page runs the fits and reports the numbers.

A brief output-DP primer

Given a query \(f: \mathcal{D} \to \mathbb{R}^p\), the Gaussian mechanism releases \(\tilde f(D) = f(D) + \mathcal{N}(0, \sigma^2 I)\). Sensitivity-driven \(\sigma\) gives single-query \((\varepsilon, \delta)\) DP. Multi-query composition uses zCDP (Bun and Steinke 2016): each release is \((\Delta/\sigma)^2/2\)-zCDP; \(T\) releases compose linearly to \(T \cdot \rho\); convert back to \((\varepsilon, \delta)\) via \(\varepsilon = \rho + 2\sqrt{\rho \log(1/\delta)}\).

The protocol modification

Without noise, site \(i\) encrypts \(x_i + u_i\) and the aggregator decrypts only their average \(z\). The change here is that at every ADMM iteration site \(i\) adds an independent draw \(\eta_i \sim \mathcal{N}(0, \sigma^2 N \cdot I)\) to \(x_i + u_i\) before encrypting. The encrypted noise terms sum under the joint key, the \(1/N\) scaling contracts the variance back to \(\sigma^2\) per coordinate, and the recovered \(z\) has noise \(\mathcal{N}(0, \sigma^2 I)\).

As on the Cox DP page, each site draws its own noise, so only site \(i\) sees its noiseless \(x_i + u_i\). The aggregator sees encrypted noisy contributions and, after decryption, their noisy average.

In the code below, site_contribution_dp() adds the noise and encrypts inside the site, so encrypted_consensus_dp() receives only encrypted values. If the aggregator added the noise instead, the numbers would be the same, but the aggregator would see each site’s noiseless \(x_i + u_i\).

The stopping rule changes

Under noise the residuals do not shrink below the per-iteration noise, so the number of iterations \(T\) is fixed in advance. Stopping on the residuals would also make \(T\) depend on the data, and \(T\) multiplies the privacy budget below. The next section picks \(\rho\) and \(T\) without using the cohort.

Setup

import warnings

import cvxpy as cp
import numpy as np

warnings.simplefilter("ignore", UserWarning)  # silence solver warnings

N, P, LAM = 3, 4, 1.0
SOLVER = cp.CLARABEL

from homomorphepy import Site


# Inherits the abstract `Site` (which supplies `name` and the key share),
# so it can take part in threshold key generation and keep its own
# share. z and u are cvxpy Parameters, so the problem canonicalizes
# once and every iteration reuses it.
class ConsensusSite(Site):
    _needs_local_fn = False  # ADMM does not use the contribute() round

    def __init__(self, name, X, y, rho):
        super().__init__(name, (X, y), None)
        self.x_var = cp.Variable(P)
        self.z_par = cp.Parameter(P, value=np.zeros(P))
        self.u_par = cp.Parameter(P, value=np.zeros(P))
        margins = -cp.multiply(2.0 * y - 1.0, X @ self.x_var)
        local_loss = (cp.sum(cp.logistic(margins))
                      + (LAM / (2 * N)) * cp.sum_squares(self.x_var))
        augmented = (rho / 2) * cp.sum_squares(self.x_var - self.z_par + self.u_par)
        self.problem = cp.Problem(cp.Minimize(local_loss + augmented))
        self.x_curr = np.zeros(P)
        self.u_curr = np.zeros(P)

    def local_update(self, z_curr):
        self.z_par.value = z_curr
        self.u_par.value = self.u_curr
        self.problem.solve(solver=SOLVER)
        if self.problem.status not in ("optimal", "optimal_inaccurate"):
            raise RuntimeError(f"local solve at {self.name} ended {self.problem.status}")
        self.x_curr = np.asarray(self.x_var.value).ravel()

    def dual_update(self, z_new):
        self.u_curr = self.u_curr + (self.x_curr - z_new)

# The loop takes the consensus step as a function, so the same loop
# runs in the clear and through the encrypted channel. It stops when
# both residuals fall below tol; tol = 0 runs exactly max_iter iterations.
def admm_loop(sites, rho, max_iter, tol, consensus_fn):
    z = np.zeros(P)
    for k in range(1, max_iter + 1):
        for s in sites:
            s.local_update(z)
        z_new = consensus_fn(sites)
        for s in sites:
            s.dual_update(z_new)
        primal = np.sqrt(np.mean([np.sum((s.x_curr - z_new) ** 2) for s in sites]))
        dual = rho * np.linalg.norm(z_new - z)
        z = z_new
        if primal < tol and dual < tol:
            return z, k, True
    return z, max_iter, False


def plain_consensus(sites):
    return sum(s.x_curr + s.u_curr for s in sites) / len(sites)

Simulated cohort

N_PER = (500, 1000, 1500)
BETA_TRUE = np.array([-0.5, 0.4, -0.3, 0.6])  # intercept, age, bmi, sex
COVARIATES = ["intercept", "age", "bmi", "sex"]


def draw_cohort(beta, rng):
    cohort = []
    for n in N_PER:
        X = np.column_stack([np.ones(n), rng.normal(size=n), rng.normal(size=n),
                             rng.binomial(1, 0.5, size=n)]).astype(float)
        prob = 1.0 / (1.0 + np.exp(-(X @ beta)))
        cohort.append((X, (rng.uniform(size=n) < prob).astype(float)))
    return cohort


site_data = draw_cohort(BETA_TRUE, np.random.default_rng(20260412))

Choosing \(\rho\) and \(T\) without touching the cohort

Tuning \(\rho\) on the real data would itself be a release: the chosen \(\rho\) and \(T\) depend on the records, and the budget below counts only the \(T\) Gaussian releases. So the sweep runs on a surrogate cohort built only from facts the protocol already treats as public: the number of sites, their sizes, and the covariate schema. The effect sizes are nominal values fixed in the analysis plan. No record from any site enters it, so the sweep needs no encryption.

import pandas as pd

TOL, MAX_ITER = 1e-3, 60
BETA_NOMINAL = np.array([0.0, 0.5, 0.5, 0.5])  # nominal, not the cohort's

surrogate_data = draw_cohort(BETA_NOMINAL, np.random.default_rng(20260413))


def build(cohort, rho):
    return [ConsensusSite(f"Site {i + 1}", X, y, rho) for i, (X, y) in enumerate(cohort)]


rows = []
for rho in (10.0, 20.0, 50.0, 100.0, 500.0):
    _, k, ok = admm_loop(build(surrogate_data, rho), rho, MAX_ITER, TOL, plain_consensus)
    rows.append({"rho": rho, "iters": k, "converged": ok})
rho_sweep = pd.DataFrame(rows)

converged = rho_sweep[rho_sweep["converged"]]
if converged.empty:
    raise RuntimeError("No rho in the grid converged within max_iter on the surrogate.")
best = converged["iters"].idxmin()  # the first minimum, in grid order
rho_chosen, T_fixed = float(converged.loc[best, "rho"]), int(converged.loc[best, "iters"])
rho_sweep

	rho	iters	converged
0	10.0	60	False
1	20.0	60	True
2	50.0	30	True
3	100.0	27	True
4	500.0	60	False

The value of \(\rho\) with the fewest iterations to convergence is \(\rho = 100\), with \(T = 27\). Both are now fixed. The DP-ADMM loop below runs exactly \(T = 27\) iterations, whatever the residuals do.

Threshold-FHE setup

from homomorphepy import fhe_context, make_threshold_master

cc = fhe_context("CKKS", multiplicative_depth=1, scaling_mod_size=59,
                 first_mod_size=60, batch_size=8)

The DP version of the consensus step. The only change from the Cox-lasso page’s encrypted_consensus() is the rng.normal(0.0, sigma * np.sqrt(N), size=P) term inside the per-site step:

def site_contribution_dp(site, sigma, rng):
    # Site-side: the site draws its own noise, adds it, and encrypts with
    # the public parameters it received at setup, all before anything
    # leaves the site. The noiseless x_i + u_i never leaves.
    noise = rng.normal(0.0, sigma * np.sqrt(N), size=P) if sigma > 0 else 0.0
    return site.encrypt(site.x_curr + site.u_curr + noise)


def encrypted_consensus_dp(master, sites, sigma, rngs):
    # Aggregator-side: sum the encrypted values, scale, threshold-decrypt.
    # The 1/N scaling contracts the summed noise variance back to sigma^2.
    cts = [site_contribution_dp(s, sigma, r) for s, r in zip(sites, rngs)]
    ct_avg = sum(cts) * (1.0 / len(sites))
    return np.asarray(master.decrypt(ct_avg, length=P), dtype=float)

The DP-ADMM loop

def run_dp_admm(sigma, seed):
    # The sites exist first: the joint public key is built from them,
    # each keeping the share it generates.
    sites = build(site_data, rho_chosen)
    master = make_threshold_master("Aggregator", cc, sites)
    rngs = [np.random.default_rng([seed, i]) for i in range(N)]
    z, _, _ = admm_loop(sites, rho_chosen, T_fixed, 0.0,
                        lambda s: encrypted_consensus_dp(master, s, sigma, rngs))
    return z

Centralized cvxpy fit

This fit pools the raw data, so it is not part of the protocol. We use it only to measure how far the noise moves the answer, in the last column of the table below. It is not released, so it is not charged to the budget.

X_pooled = np.vstack([X for X, _ in site_data])
y_pooled = np.concatenate([y for _, y in site_data])
beta_var = cp.Variable(P)
margins = -cp.multiply(2.0 * y_pooled - 1.0, X_pooled @ beta_var)
cp.Problem(cp.Minimize(cp.sum(cp.logistic(margins))
                       + (LAM / 2) * cp.sum_squares(beta_var))).solve(solver=SOLVER)
beta_central = np.asarray(beta_var.value).ravel()

The \(\sigma\) sweep

Six \(\sigma\) values from zero to one. The \(\sigma = 0\) row checks that the protocol without noise matches the centralized fit.

sigma_grid = [0.0, 1e-4, 1e-3, 1e-2, 1e-1, 1.0]
sweep_results = [run_dp_admm(s, seed=100 + j) for j, s in enumerate(sigma_grid, start=1)]

clean_dev = float(np.max(np.abs(sweep_results[0] - beta_central)))
agree_tol = 10 * TOL
if clean_dev > agree_tol:
    raise RuntimeError("DP-ADMM at sigma = 0 disagrees with the centralized fit.")

At \(\sigma = 0\) the largest coefficient deviation from the centralized fit is 3.19e-05, within \(10 \times\) the ADMM tolerance of 0.001.

summary = pd.DataFrame(
    [list(z) + [np.max(np.abs(z - beta_central))] for z in sweep_results]
    + [list(beta_central) + [0.0]],
    columns=COVARIATES + ["max_dev"],
    index=[f"sigma={s:g}" for s in sigma_grid] + ["centralized"],
)
summary.round(6)

	intercept	age	bmi	sex	max_dev
sigma=0	-0.488292	0.395497	-0.292848	0.602268	0.000032
sigma=0.0001	-0.488225	0.395567	-0.292813	0.602169	0.000131
sigma=0.001	-0.488050	0.393451	-0.291553	0.602096	0.002042
sigma=0.01	-0.482145	0.388394	-0.288560	0.609925	0.007624
sigma=0.1	-0.494716	0.422191	-0.351942	0.546630	0.059092
sigma=1	-1.183212	1.272715	-2.147010	2.059519	1.854161
centralized	-0.488296	0.395493	-0.292849	0.602300	0.000000

Fidelity decays monotonically as expected.

Privacy budget

Each iteration releases one noisy average \(z\), with noise \(\mathcal{N}(0, \sigma^2 I)\). Let \(\Delta\) be the largest change in \(z\) (L2 norm over all \(p\) coordinates) from adding or removing one record. Each record sits at one site, and the aggregator decrypts only the average, so one iteration is a single Gaussian release with \(\rho_{\text{iter}} = (\Delta/\sigma)^2/2\). Over \(T = 27\) iterations the total is \(\rho = T \cdot (\Delta/\sigma)^2/2\), whatever the number of sites. This is the same accounting as the Cox DP page. With sensitivity \(\Delta = 1\) (placeholder) and target \(\delta = 10^{-5}\), zCDP composition gives:

def zcdp_to_eps(rho, delta=1e-5):
    return rho + 2 * np.sqrt(rho * np.log(1 / delta))


budget = pd.DataFrame({"sigma": [s for s in sigma_grid if s > 0]})
budget["rho_total"] = T_fixed * (1 / budget["sigma"]) ** 2 / 2
budget["epsilon_at_delta_1e-5"] = zcdp_to_eps(budget["rho_total"])
budget.round(4)

	sigma	rho_total	epsilon_at_delta_1e-5
0	0.0001	1.350000e+09	1.350249e+09
1	0.0010	1.350000e+07	1.352493e+07
2	0.0100	1.350000e+05	1.374934e+05
3	0.1000	1.350000e+03	1.599339e+03
4	1.0000	1.350000e+01	3.843390e+01

The smallest \(\varepsilon\), at \(\sigma = 1\) where the fit is already poor, is 38. Whether that is acceptable depends on the application.

One could also explore a tighter sensitivity bound \(\Delta\), or pay for the choice of \(\rho\) and \(T\) with a DP selection mechanism, etc. We don’t do that here.

The same run, in one call

from homomorphepy.examples.dp import consensus_sweep

result = consensus_sweep()
result.choice.rho, result.choice.T, f"{result.clean_dev:.2e}"

(100.0, 27, '3.18e-05')

References

Boyd, Stephen, Neal Parikh, Eric Chu, Borja Peleato, and Jonathan Eckstein. 2011. “Distributed Optimization and Statistical Learning via the Alternating Direction Method of Multipliers.” Foundations and Trends in Machine Learning 3 (1): 1–122. https://doi.org/10.1561/2200000016.

Bun, Mark, and Thomas Steinke. 2016. “Concentrated Differential Privacy: Simplifications, Extensions, and Lower Bounds.” Theory of Cryptography Conference (TCC), 635–58. https://doi.org/10.1007/978-3-662-53641-4_24.