Distributed Cox Regression with Threshold Key Generation

Introduction

The Cox page fits a stratified Cox model across three sites under CKKS. There the master holds the secret key and decrypts the encrypted sum at every iteration of the optimizer, so the master must be trusted not to decrypt anything else, such as a single site’s contribution.

On this page the three sites generate the CKKS key pair jointly, and each keeps its own share of the secret key. No single party holds the whole key. The master becomes an aggregator: it adds the encrypted contributions and collects the sites’ partial decryptions, but it cannot decrypt anything by itself.

The data and the model are the same as on the Cox page. Only who can decrypt changes.

Threat model

Three sites and one untrusted aggregator:

Sites \(1, 2, 3\) each hold private patient data and a secret key share \(\mathit{sk}_i\). They are honest-but-curious among themselves and toward the aggregator.
Aggregator holds no secret-key material. It receives the encrypted contributions, adds them, and sends the sum back to the sites for partial decryption. The encrypted values it handles tell it nothing on their own.

What the aggregator sees, by stage:

Encrypted local contributions \(\mathit{ct}_i = E_{\mathit{pk}_{1..n}}(\ell_i)\). None decryptable alone.
The encrypted sum \(\mathit{ct}_{\text{sum}} = \boxplus_i \mathit{ct}_i\). Not decryptable alone.
Partial decryptions \(\rho_i\) contributed by each site. Not decryptable individually.
After combining the partial decryptions, the sum \(\ell(\beta) = \sum_i \ell_i\) in the clear.

Step 4 reveals \(\ell(\beta)\) to the aggregator, which is what the master saw on the Cox page. The difference is that no single party can decrypt an individual contribution or any intermediate value. That takes a partial decryption from every site.

The Cox setup (same DLBCL data as the Cox page)

import numpy as np
from homomorphepy import load_dlbcl, site_order

df = load_dlbcl()
COVARIATES = ["GCB_sig", "LN_sig", "Prolif_sig", "BMP6", "MHC2_sig"]

def site_frame(name):
    m = (df["Subgroup"] == name).to_numpy()
    return {
        "time":   df["time"].to_numpy(dtype=float)[m],
        "status": df["status"].to_numpy(dtype=int)[m],
        "X":      np.column_stack([df[c].to_numpy(dtype=float)[m]
                                   for c in COVARIATES]),
    }

cox_data = {name: site_frame(name) for name in site_order()}

The protocol

Setup (once):

Site 1 calls KeyGen() to produce its keypair \((\mathit{pk}_1, \mathit{sk}_1)\).
Site 2 calls MultipartyKeyGen(pk_1) to produce \((\mathit{pk}_{12}, \mathit{sk}_2)\).
Site 3 calls MultipartyKeyGen(pk_12) to produce \((\mathit{pk}_{123}, \mathit{sk}_3)\).
The final \(\mathit{pk}_{123}\) is the joint public key. Each site keeps its own \(\mathit{sk}_i\).

Per query (called inside the optimizer):

Each site \(i\) computes its local Cox negative log-likelihood \(\ell_i(\beta)\) and encrypts it under the joint public key.
The aggregator sums the encrypted contributions homomorphically.
Each site partial-decrypts the sum using its own \(\mathit{sk}_i\).
The aggregator fuses the partials to recover \(\ell(\beta)\).

make_threshold_master() runs this chain across the sites in one call and returns a ThresholdMaster holding the joint public key. Each site keeps the share it generated. To decrypt, the master asks every site for a partial decryption and combines them. This happens inside the decrypt() method, so the aggregate() runner from the Cox page works unchanged.

Implementation

from statsmodels.duration.hazard_regression import PHReg

def local_cox_nll(data, beta):
    b = np.asarray(beta, dtype=float).ravel()
    model = PHReg(data["time"], data["X"],
                  status=data["status"], ties="efron")
    return float(-model.loglike(b))

The CKKS context is the one from the Cox page. The chained MultipartyKeyGen() calls need the MULTIPARTY feature, which make_threshold_master() enables on the context it is given:

from homomorphepy import fhe_context, make_worker, make_threshold_master

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

The sites come first, because the joint public key is built from them. make_threshold_master() then takes the sites and returns the master.

workers = [make_worker(name, data, local_cox_nll)
           for name, data in cox_data.items()]

master = make_threshold_master("Aggregator", cc, workers)

The check below confirms that the master has no attribute holding key shares and that the GCB site holds its own share:

share_check = dict(master_holds_shares = hasattr(master, "_share"),
                   gcb_holds_own_share = workers[0].has_share)
share_check

{'master_holds_shares': False, 'gcb_holds_own_share': True}

Iterative MLE through the threshold protocol

The optimizer code is the same as on the Cox page. Only the master class differs.

from scipy.optimize import minimize

def encrypted_nll(beta):
    value = master.aggregate(np.asarray(beta, dtype=float))
    return 1e12 if (value is None or np.isnan(value)) else float(value)

fit = minimize(encrypted_nll, x0=np.zeros(len(COVARIATES)),
               method="L-BFGS-B", options={"ftol": 1e-9})
beta_hat = np.asarray(fit.x)

Standard errors come from the Hessian of the encrypted objective at the fit, as on the Cox page:

import pandas as pd
from homomorphepy.examples.cox import fd_hessian

se = np.sqrt(np.diag(np.linalg.inv(fd_hessian(encrypted_nll, beta_hat))))
pd.DataFrame({"estimate": beta_hat, "std. error": se}, index=COVARIATES)

	estimate	std. error
GCB_sig	-0.263872	0.119404
LN_sig	-0.254359	0.085152
Prolif_sig	0.303126	0.149813
BMP6	0.303637	0.107278
MHC2_sig	-0.319147	0.094129

float(-fit.fun)

-495.22902176698926

Comparison with the cleartext fit

As on the Cox page, the check is the identical minimize() objective with the encrypted aggregation replaced by an ordinary sum of the three sites’ cleartext values.

def cleartext_nll(beta):
    return sum(local_cox_nll(d, beta) for d in cox_data.values())

fit_plain = minimize(cleartext_nll, x0=np.zeros(len(COVARIATES)),
                     method="L-BFGS-B", options={"ftol": 1e-9})
beta_plain = np.asarray(fit_plain.x)

	minimize(), threshold	minimize(), cleartext	absolute difference
coefficient
GCB_sig	-0.263872	-0.263872	2.43e-07
LN_sig	-0.254359	-0.254359	7.65e-08
Prolif_sig	0.303126	0.303126	6.16e-07
BMP6	0.303638	0.303638	1.04e-08
MHC2_sig	-0.319146	-0.319147	4.32e-08

The threshold fit agrees with the cleartext fit to within 6.16e-07 in every coefficient.

assert enc_diff.max() < 1e-4

The same run, in one call

from homomorphepy.examples import cox

result = cox.run(backend="threshold")
result.loglik_encrypted, result.loglik_centralized, result.max_abs_difference

(-495.2290217669887, -495.22902176698074, 3.224813755475786e-07)

Discussion

No single party holds the decryption key. Each site generated its own share and kept it; the master holds only the joint public key, and has no attribute in which a share could sit. Encrypted intermediate values are undecryptable by any single party in the system, the aggregator included.
The fit does not change. The coefficients agree with the cleartext fit of the same objective to CKKS precision.
Same optimizer, same callback shape, same code structure. Compared to the Cox page only the setup changed: the workers are built first and handed to make_threshold_master() instead of being wired to a make_ckks_master() afterwards, because the joint key cannot exist before the sites do. The optimizer sees nothing different.

Limitations

The aggregator sees \(\ell(\beta)\) at every iteration. That is the function value minimize() asks for. The individual site contributions stay hidden. Hiding \(\ell(\beta)\) as well would require running the optimizer on encrypted values, which is possible but considerably more complex.
Honest-but-curious is the trust model. Sites are assumed to follow the protocol. A malicious site could submit a corrupted partial decryption to break the fit; detecting this requires additional protocol machinery (commitments, zero-knowledge proofs) that this page does not implement.
Output privacy is unchanged. The released coefficients \(\hat\beta\) are the same as the cleartext fit. Output-level attacks (membership inference, model inversion) remain in scope and motivate the differential-privacy demonstrations elsewhere in the package.