Distributed Maximum Likelihood Estimation

The statistical problem

Suppose we have count data \(y_1, y_2, \ldots, y_n\) that we model as independent draws from a Poisson distribution with unknown parameter \(\lambda\). The maximum likelihood estimate is \(\hat{\lambda} = \bar{y}\), obtained by minimizing the negative log-likelihood

\[ -\ell(\lambda \mid y) \;=\; -\sum_{i=1}^{n} \log p(y_i \mid \lambda). \]

In Python, this is a few lines using scipy.optimize.minimize():

import numpy as np
from scipy.optimize import minimize
from scipy.special import gammaln

rng = np.random.default_rng(17822)
y = rng.poisson(10.0, size=40)

def nLL(lam):
    lam = float(np.ravel(lam)[0])
    return -np.sum(y * np.log(lam) - lam - gammaln(y + 1.0))

fit0 = minimize(nLL, x0=[5.0], method="BFGS")
float(fit0.x[0]), float(np.mean(y))
(10.374998677343168, 10.375)

The privacy constraint

Now suppose the same data is distributed across three sites — say, three hospitals counting adverse events. None will share its raw counts with the others or with a central aggregator, but they are willing to jointly compute the same MLE provided no party learns anything about another party’s contribution.

To simulate this, partition y:

y1, y2, y3 = y[:20], y[20:27], y[27:]

The negative log-likelihood factorizes additively:

\[ -\ell(\lambda \mid y) \;=\; -\ell_1(\lambda \mid y_1) - \ell_2(\lambda \mid y_2) - \ell_3(\lambda \mid y_3) \]

so each site can compute its local term in the clear and only the sum of the three local likelihoods needs to travel between parties — and the sum must not reveal the individual addends.

The protocol

We use the master/worker topology that distcomp- and DataSHIELD-style federated analyses actually deploy: a star with the master at the center and one independent worker per site. There is no chain and no inter-site communication.

In words:

  1. The master generates a CKKS context and key pair, distributes the public key to the three workers, keeps the secret key.
  2. The master broadcasts the current \(\lambda\) to each worker.
  3. Each worker computes its local negative log-likelihood \(\ell_i(\lambda)\) on its private data, encrypts the result under the master’s public key, and returns \(E(\ell_i)\) to the master.
  4. The master sums the encrypted contributions homomorphically: \(E(\ell) = E(\ell_1) \boxplus E(\ell_2) \boxplus E(\ell_3) = E(\ell_1 + \ell_2 + \ell_3)\).
  5. The master decrypts \(E(\ell)\) to recover \(\ell\).
  6. The master hands \(\ell\) to the optimizer; the protocol repeats for each new guess of \(\lambda\) until convergence.

This is the realistic shape: each worker independently does its local computation and ships an encrypted summary; the master only sees the encrypted summaries (and, after homomorphic summation, the decrypted total). With a single-decrypter master, the master could decrypt individual \(E(\ell_i)\) in principle; the cryptographic story strengthens when paired with threshold key generation, where no single party holds the secret key. We will revisit that in the Cox threshold page.

Implementation

The computational topology — Site (worker), Master, the master/worker runner — is the same code as in any other distributed-stats page in this package; only the master’s backend changes. homomorphepy exports make_ckks_master() that takes a context and key pair. The master decrypts with the secret key; each site encrypts its own contribution under the public key it is given at setup. The worker class and the aggregate() runner are backend-agnostic.

The per-site negative log-likelihood is the same plain Python function it would be in the cleartext case:

def local_nll(data, lam):
    lam = float(np.ravel(lam)[0])
    if lam <= 0:
        return np.nan          # outside the parameter space
    data = np.asarray(data, dtype=float)
    return float(-np.sum(data * np.log(lam) - lam - gammaln(data + 1.0)))

1. Generate a CKKS key pair

from homomorphepy import fhe_context

cc = fhe_context("CKKS", multiplicative_depth=1,
                 scaling_mod_size=50, batch_size=8)
keys = cc.KeyGen()

2. Build workers and master

from homomorphepy import make_worker, make_ckks_master

workers = [make_worker("Site 1", y1, local_nll),
           make_worker("Site 2", y2, local_nll),
           make_worker("Site 3", y3, local_nll)]

master = make_ckks_master("Master", cc, keys).set_workers(workers)

3. Run minimize() through the encrypted channel

n_calls = 0

def objective(theta):
    global n_calls
    n_calls += 1
    value = master.aggregate(float(np.ravel(theta)[0]))
    # A parameter outside the support makes the objective
    # non-evaluable; the line search cannot act on NaN, so hand it a
    # large finite value it can step back from.
    return 1e12 if (value is None or np.isnan(value)) else float(value)

fit1 = minimize(objective, x0=[5.0], method="BFGS",
                options={"finite_diff_rel_step": 1e-3})

lambda_hat = float(fit1.x[0])
lambda_hat, n_calls
(10.375001912473042, 16)

The CKKS-based estimate differs from the cleartext estimate fit0 by 3.2e-06. No site ever revealed its individual counts to any other party.

assert abs(lambda_hat - float(np.mean(y))) < 1e-4

The same run, in one call

from homomorphepy.examples import mle

fit = mle.run()
fit.lambda_encrypted, fit.lambda_pooled, fit.abs_difference
(10.374997323397523, 10.375, 2.676602477436063e-06)

Beyond MLE

For Poisson MLE the protocol uses only additions, so the additive Paillier scheme would also suffice. CKKS saves the integer encoding of real numbers that Paillier needs. The same holds for stratified Cox regression, where the sites’ partial log-likelihoods are summed. Multiplying encrypted values is needed only when the computation itself runs on encrypted data, as in the sigmoid of encrypted regression.

CAVEAT

This is a teaching example. In production you would want a real communication transport, threshold key generation so no single party holds the full secret, and persistent serialization at site boundaries. The point of this page is not the deployment story but the structure of a privacy-preserving distributed computation built on homomorphic encryption.