An example using federated learning for causal inference
The case above was simply a vanilla optimal transport problem that
could easily be solved by the calc_weight function in the
main package. Let’s look at a more complicated use case.
Solving the problem directly
Ideally, we’d simply dump the data together and run our OT framework.
nrow <- 100
ncol <- 2
a <- Measure(x = matrix(rnorm(nrow*ncol,mean=c(0.1,0.1)) + 0.1,nrow,ncol,byrow = TRUE), adapt = "weights")
b <- Measure(x = matrix(rnorm(nrow*ncol,mean=c(-0.1,-0.1),sd=0.25),nrow,ncol,byrow = TRUE), adapt = "weights")
c <- Measure(x = matrix(rnorm(nrow*ncol,mean=c(0.1,-0.1)),nrow,ncol,byrow = TRUE), adapt = "weights")
d <- Measure(x = matrix(rnorm(nrow*ncol,mean= c(-0.1,0.1),sd=0.25),nrow,ncol,byrow = TRUE), adapt = "weights")
overall <- Measure(x = torch::torch_vstack(lapply(list(a,b,c,d), function(meas) meas$x)),
adapt = "none")
overall_ot <- OTProblem(a,overall) + OTProblem(b, overall) +
OTProblem(c, overall) + OTProblem(d, overall)One thing to note that’s kind of cool is that we can add our
OTProblem objects together to make a unified objective
function.
overall_ot
#> OT Problem:
#> OT(0x7fdd7a0af5d0, 0x7fdd7a2b50a0) +
#> OT(0x7fdd7a15d830, 0x7fdd7a2b50a0) +
#> OT(0x7fdd7a1e1798, 0x7fdd7a2b50a0) +
#> OT(0x7fdd79b17810, 0x7fdd7a2b50a0)Neat!
We can also run the calc_weight function in each
treatment group targeting the overall population
source_measures <- list(a,b,c,d)
meas <- x_temp <- NULL
z_temp <- c(rep(1, nrow*4), rep(0,nrow))
wt <- list()
for(i in seq_along(source_measures)) {
meas <- source_measures[[i]]
x_temp <- as.matrix(torch::torch_vstack(list(overall$x,meas$x)))
wt[[i]] <- calc_weight(x = x_temp,
z = z_temp,
estimand = "ATT",
method = "COT")
}Data sharing restrictions exist
Only moments are available
If only moments are available, then each site can run essentially independently. We just need to collect the moments from each site and combine
target.values <-
as.numeric(a$x$mean(1) + b$x$mean(1) +
c$x$mean(1) + d$x$mean(1))/4
a_t <- Measure(x = a$x, adapt = "weights",
target.values = target.values)
b_t <- Measure(x = a$x, adapt = "weights",
target.values = target.values)
c_t <- Measure(x = a$x, adapt = "weights",
target.values = target.values)
d_t <- Measure(x = a$x, adapt = "weights",
target.values = target.values)
all.target.measures <- list(a_t, b_t, c_t, d_t)Then we can optimize the weights targeting the moments in a bit of a hacky way.
ot_targ <- NULL
for(meas in all.target.measures) {
ot_targ <- OTProblem(meas, meas)
ot_targ$setup_arguments(lambda = 100)
ot_targ$solve(torch_optim = torch::optim_lbfgs,
torch_args = list(line_search_fn = "strong_wolfe"))
}And we can check the final balance
final.bal <- as.numeric(a_t$x$mT()$matmul(a_t$weights$detach()))
original <- as.numeric(a_t$x$mean(1))
rbind(original,
`final balance` = final.bal,
`target values` = target.values)
#> [,1] [,2]
#> original 0.08015652 0.31129337
#> final balance -0.01093957 0.09104138
#> target values -0.01100336 0.09096446This will target the moments without information about the underlying distributions. Obviously, we would prefer to use more of the available information, as we describe next.
We can pass other data that respects privacy
Say we can pass any amount of data but are limited by the fact that privacy or other restrictions prevent us from sharing the full data at each site. We can instead construct a pseudo-overall population using Wasserstein Barycenters. These construct average distributions. Let’s see how it might work.
Option 1: pass the gradients back
In this option, we pass gradients back to the main site. From this, we can construct a pseudo average population. Let’s see how it might work. We first construct our pseudo data.
Importantly, each data point must be initialized to a separate value otherwise all of the points will move together Then we pass this pseudo data and set up a problem at each site.
pseudo_a <- pseudo$detach()
pseudo_b <- pseudo$detach()
pseudo_c <- pseudo$detach()
pseudo_d <- pseudo$detach()
pseudo_a$requires_grad <- pseudo_b$requires_grad <-
pseudo_c$requires_grad <- pseudo_d$requires_grad <- "x"
ota <- OTProblem(a$detach(), # don't update a
pseudo_a)
otb <- OTProblem(b$detach(), # don't update b
pseudo_b)
otc <- OTProblem(c$detach(), # don't update c
pseudo_c)
otd <- OTProblem(d$detach(), # don't update c
pseudo_d)Then we setup the arguments. For simplicity, we will set
lambda = 0.1.
ota$setup_arguments(lambda = .1)
otb$setup_arguments(lambda = .1)
otc$setup_arguments(lambda = .1)
otd$setup_arguments(lambda = .1)Then we setup our optimizer at the main site
opt <- torch::optim_rmsprop(pseudo$x)
sched <- torch::lr_multiplicative(opt, lr_lambda = function(epoch) {0.99})Then we run our optimization loop like so:
#optimization loop
for (i in 1:100) {
# zero grad of main optimizer
opt$zero_grad()
# get gradients at each site
ota$loss$backward()
otb$loss$backward()
otc$loss$backward()
otd$loss$backward()
# pass grads back to main site
pseudo$grad <- pseudo_a$grad + pseudo_b$grad +
pseudo_c$grad + pseudo_d$grad
# update pseudo data at main site
opt$step()
# zero site gradients
torch::with_no_grad({
pseudo_a$grad$copy_(0.0)
pseudo_b$grad$copy_(0.0)
pseudo_c$grad$copy_(0.0)
pseudo_d$grad$copy_(0.0)
})
# update scheduler
sched$step()
}Then we pass the final pseudo data back to the sites and optimize the weights at each site:
pseudo_a$x <- pseudo_b$x <- pseudo_c$x <- pseudo_d$x <-
pseudo$x
ota_w <- OTProblem(a, pseudo_a$detach())
otb_w <- OTProblem(b, pseudo_b$detach())
otc_w <- OTProblem(c, pseudo_c$detach())
otd_w <- OTProblem(d, pseudo_d$detach())
ota_w$setup_arguments()
ota_w$solve(torch_args = list(line_search_fn = "strong_wolfe"))
ota_w$choose_hyperparameters()
otb_w$setup_arguments()
otb_w$solve(torch_args = list(line_search_fn = "strong_wolfe"))
otb_w$choose_hyperparameters()
otc_w$setup_arguments()
otc_w$solve(torch_args = list(line_search_fn = "strong_wolfe"))
otc_w$choose_hyperparameters()
otd_w$setup_arguments()
otd_w$solve(torch_args = list(line_search_fn = "strong_wolfe"))
otd_w$choose_hyperparameters()Note we haven’t checked for convergence when constructing the pseudo-data to save time. You should, however, do this in your own work.
Of course, maybe we can’t pass gradients. Instead, we can create pseudo data at each site.
Option 2: create pseudo data at each site and then pass this back to the main site
The second option is to create pseudo-data for each site and then use this to generate an overall average data set. This will allow us to create privacy respecting pseudo-data in each site, i.e., data that is close to the population at A but with different values. Then we take these pseudo-data and create an overall average data set like follows.
First, we need to reinitialize sites again since they were changed in the previous example.
a$weights <- a$init_weights
b$weights <- b$init_weights
c$weights <- c$init_weights
d$weights <- d$init_weightsThen again create pseudo data
Then we pass this pseudo data and set up a problem at each site.
pseudo_a <- pseudo$detach()
pseudo_b <- pseudo$detach()
pseudo_c <- pseudo$detach()
pseudo_d <- pseudo$detach()
pseudo_a$requires_grad <- pseudo_b$requires_grad <-
pseudo_c$requires_grad <- pseudo_d$requires_grad <- "x"
ota <- OTProblem(a$detach(), # don't update a
pseudo_a)
otb <- OTProblem(b$detach(), # don't update b
pseudo_b)
otc <- OTProblem(c$detach(), # don't update c
pseudo_c)
otd <- OTProblem(d$detach(), # don't update c
pseudo_d)and setup the arguments. For simplicity, we will set
lambda = 0.1.
ota$setup_arguments(lambda = .1)
otb$setup_arguments(lambda = .1)
otc$setup_arguments(lambda = .1)
otd$setup_arguments(lambda = .1)Then we solve for the barycenters.
# run separately at each site
ota$solve(torch_optim = torch::optim_rmsprop)
otb$solve(torch_optim = torch::optim_rmsprop)
otc$solve(torch_optim = torch::optim_rmsprop)
otd$solve(torch_optim = torch::optim_rmsprop)Looking at the pseudo data in group B we can see that the pseudo-data
is a much better approximation to B after optimization. 
Then we send the pseudo-data back to our main site to create the overall pseudo-data
# send back to the main site and create overall problem
ot_overall <-
OTProblem(pseudo_a$detach(),
pseudo) +
OTProblem(pseudo_b$detach(),
pseudo) +
OTProblem(pseudo_c$detach(),
pseudo) +
OTProblem(pseudo_d$detach(),
pseudo)
ot_overall$setup_arguments(lambda = 0.1)
ot_overall$solve(torch_optim = torch::optim_rmsprop)Now we have an average population to target at each site, which we can do like so:
# pass pseudo to each site then setup the problems again
ota2 <- OTProblem(a,
pseudo$detach())
otb2 <- OTProblem(b, # don't update b
pseudo$detach())
otc2 <- OTProblem(c,
pseudo$detach())
otd2 <- OTProblem(d,
pseudo$detach())
all.problems <- list(ota2,
otb2,
otc2,
otd2)
# then we optimize the weights at each site separately.
for (prob in all.problems) {
prob$setup_arguments()
prob$solve(
torch_optim = torch::optim_lbfgs,
torch_args = list(line_search_fn = "strong_wolfe")
)
prob$choose_hyperparameters()
}