--- title: "Distributed Cox Regression with Threshold Key Generation" author: "Balasubramanian Narasimhan" date: '`r Sys.Date()`' bibliography: homomorphing.bib output: html_document: fig_caption: yes theme: cerulean toc: yes toc_depth: 2 vignette: > %\VignetteIndexEntry{Distributed Cox Regression with Threshold Key Generation} %\VignetteEngine{knitr::rmarkdown} \usepackage[utf8]{inputenc} --- ```{r ktab, echo=FALSE} ## Tables through kableExtra. Math written as $...$ in headers, cells ## and captions becomes \( ... \) and `code` becomes , because ## pandoc does not process math inside a raw HTML table. ktab <- function(x, ..., col.names = names(x), caption = NULL) { tex <- function(s) gsub("`([^`]*)`", "\\1", gsub("\\$([^$]+)\\$", "\\\\(\\1\\\\)", s)) chr <- vapply(x, is.character, logical(1)) x[chr] <- lapply(x[chr], tex) tab <- knitr::kable(x, format = "html", escape = FALSE, col.names = tex(col.names), caption = if (!is.null(caption)) tex(caption), ...) kableExtra::kable_styling(tab, bootstrap_options = c("striped", "condensed"), full_width = TRUE) } ``` ```{r echo=F, purl=FALSE} knitr::opts_chunk$set( message = FALSE, warning = FALSE, error = FALSE, tidy = FALSE, cache = FALSE ) ## The encrypted chunks are gated `eval = RECOMPUTE` and are shown but ## not run when the vignette builds; their results are read from ## data(cox_threshold_results), produced by ## data-raw/cox_threshold_results.R from these same chunks. ## HOMOMORPHER_RECOMPUTE=true runs them for real. if (!exists("RECOMPUTE")) RECOMPUTE <- isTRUE(as.logical(Sys.getenv("HOMOMORPHER_RECOMPUTE", "FALSE"))) data(cox_threshold_results, package = "homomorpheR") ``` ## Introduction The `cox` vignette 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. In this vignette 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 in `vignette("cox")`. 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: 1. Encrypted local contributions $\mathit{ct}_i = E_{\mathit{pk}_{1..n}}(\ell_i)$. None decryptable alone. 2. The encrypted sum $\mathit{ct}_{\text{sum}} = \boxplus_i \mathit{ct}_i$. Not decryptable alone. 3. Partial decryptions $\rho_i$ contributed by each site. Not decryptable individually. 4. 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 in `vignette("cox")`. 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 `cox.Rmd`) ```{r} suppressPackageStartupMessages(library(survival)) library(homomorpheR) data(DLBCL) cox_data <- split( DLBCL[, c("time", "status", "GCB_sig", "LN_sig", "Prolif_sig", "BMP6", "MHC2_sig", "Subgroup")], DLBCL$Subgroup) ``` ## The protocol **Setup** (once): 1. Site 1 calls `key_gen(cc)` to produce its keypair $(\mathit{pk}_1, \mathit{sk}_1)$. 2. Site 2 calls `multiparty_key_gen(cc, pk_1)` to produce $(\mathit{pk}_{12}, \mathit{sk}_2)$. 3. Site 3 calls `multiparty_key_gen(cc, pk_{12})` to produce $(\mathit{pk}_{123}, \mathit{sk}_3)$. 4. The final $\mathit{pk}_{123}$ is the **joint public key**. Each site keeps its own $\mathit{sk}_i$. **Per query** (called inside the optimizer): 1. Each site $i$ computes its local Cox negative log-likelihood $\ell_i(\beta)$ and encrypts it under the joint public key. 2. The aggregator sums the encrypted contributions homomorphically. 3. Each site partial-decrypts the sum using its own $\mathit{sk}_i$. 4. 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 `master_aggregate()` runner from `cox.Rmd` works unchanged. ## Implementation ```{r} cph_control <- replace(coxph.control(), "iter.max", 0) local_cox_nll <- function(data, beta) { fit <- tryCatch( coxph(Surv(time, status) ~ GCB_sig + LN_sig + Prolif_sig + BMP6 + MHC2_sig, data = data, init = beta, control = cph_control), error = function(e) NULL) if (is.null(fit)) NA_real_ else -fit$loglik[1] } ``` The CKKS context needs the `MULTIPARTY` feature enabled so the chained `multiparty_key_gen()` calls work: ```{r eval=RECOMPUTE} cc <- openfhe.R::fhe_context("CKKS", multiplicative_depth = 1L, scaling_mod_size = 59L, first_mod_size = 60L, batch_size = 8L, features = c(openfhe.R::Feature$MULTIPARTY)) ``` The sites come first, because the joint public key is built from them. `make_threshold_master()` then takes the sites and returns the master. ```{r eval=RECOMPUTE} worker_gcb <- make_worker(name = "GCB", data = cox_data[["GCB"]], contribution_fn = local_cox_nll) worker_abc <- make_worker(name = "ABC", data = cox_data[["ABC"]], contribution_fn = local_cox_nll) worker_t3 <- make_worker(name = "Type III", data = cox_data[["Type III"]], contribution_fn = local_cox_nll) master <- make_threshold_master("Aggregator", crypto_context = cc, sites = list(worker_gcb, worker_abc, worker_t3)) ``` The check below confirms that the master has no property holding key shares and that the GCB site holds its own share: ```{r eval=RECOMPUTE} share_check <- c(master_holds_shares = "secret_keys" %in% names(S7::props(master)), gcb_holds_own_share = !is.null(worker_gcb@state$sk)) share_check ``` ```{r echo=FALSE, purl=FALSE, eval=!RECOMPUTE} cox_threshold_results$share_check ``` ## Iterative MLE through the threshold protocol The optimizer code is the same as in `cox.Rmd`. Only the master class differs. ```{r eval=RECOMPUTE} library(stats4) encrypted_nLL <- function(GCB_sig, LN_sig, Prolif_sig, BMP6, MHC2_sig) { master_aggregate(master, c(GCB_sig, LN_sig, Prolif_sig, BMP6, MHC2_sig)) } fit <- mle(encrypted_nLL, start = list(GCB_sig = 0, LN_sig = 0, Prolif_sig = 0, BMP6 = 0, MHC2_sig = 0), method = "BFGS", control = list(reltol = 1e-7)) summary(fit) logLik(fit) ``` ```{r eval=RECOMPUTE, echo=FALSE} cox_threshold_results <- list(coef = summary(fit)@coef, loglik = as.numeric(logLik(fit)), counts = fit@details$counts, share_check = share_check) ``` ```{r echo=FALSE, purl=FALSE, eval=!RECOMPUTE} cox_threshold_results$coef cat(sprintf("'log Lik.' %f (df=5)\n", cox_threshold_results$loglik)) ``` ## Comparison with the cleartext fit As in `vignette("cox")`, the check is the identical `mle()` objective with the encrypted aggregation replaced by an ordinary sum of the three sites' cleartext values. ```{r} library(stats4) plain_nLL <- function(GCB_sig, LN_sig, Prolif_sig, BMP6, MHC2_sig) { beta <- c(GCB_sig, LN_sig, Prolif_sig, BMP6, MHC2_sig) sum(vapply(cox_data, local_cox_nll, numeric(1), beta = beta)) } fit_plain <- mle(plain_nLL, start = list(GCB_sig = 0, LN_sig = 0, Prolif_sig = 0, BMP6 = 0, MHC2_sig = 0), method = "BFGS", control = list(reltol = 1e-7)) ``` ```{r echo = FALSE} thr_coefs <- cox_threshold_results$coef[, "Estimate"] plain_coefs <- coef(fit_plain)[names(thr_coefs)] thr_diff <- abs(thr_coefs - plain_coefs) tex_sci <- function(x) { e <- floor(log10(abs(x))) sprintf("$%.2f \\times 10^{%d}$", x / 10^e, as.integer(e)) } comparison <- data.frame( Coefficient = names(thr_coefs), threshold = unname(thr_coefs), cleartext = unname(plain_coefs), abs_diff = tex_sci(unname(thr_diff)) ) ktab(comparison, digits = 7, row.names = FALSE, align = "lrrr", col.names = c("Coefficient", "$\\hat\\beta$, `mle()` threshold", "$\\hat\\beta$, `mle()` cleartext", "$\\lvert \\text{difference} \\rvert$"), caption = "Threshold-CKKS DLBCL Cox against the same objective evaluated in the clear.") ``` The threshold fit agrees with the cleartext fit to within `r formatC(max(thr_diff), format = "e", digits = 2)` in every coefficient. ## Discussion 1. **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 property in which a share could sit. Encrypted intermediate values are undecryptable by any single party in the system, the aggregator included. 2. **The fit does not change.** The coefficients agree with the cleartext fit of the same objective to CKKS precision. 3. **Same optimizer, same callback shape, same code structure.** Compared to `cox.Rmd` 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 `mle()` 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 vignette 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.