Package: geom-median
Differentiable and Fast Geometric Median in NumPy and PyTorchThis package implements a fast numerical algorithm to compute the geometric median of high dimensional vectors.
As a generalization of the median (of scalars), thegeometric medianis a robust estimator of the mean in the presence of outliers and contaminations (adversarial or otherwise).The geometric median is also known as the Fermat point, Weber's L1 median, Fréchet median among others.
It has a breakdown point of 0.5, meaning that it yields a robust aggregate even under arbitrary corruptions to points accounting for under half the total weight. We use the smoothed Weiszfeld algorithm to compute the geometric median.Features:Implementation in both NumPy and PyTorch.PyTorch implementation is fully differentiable (compatible with gradient backpropagation a.k.a. automatic differentiation) and can run on GPUs with CUDA tensors.Blazing fast algorithm that converges linearly in almost all practical settings.InstallationThis package can be installed via pip aspip install geom_median. Alternatively, for an editable install,
rungitclonegit@github.com:krishnap25/geom_median.gitcdgeom_median
pipinstall-e.You must have a working installation of PyTorch, version 1.7 or over in case you wish to use the PyTorch API.
See detailshere.Usage GuideWe describe the PyTorch usage here. The NumPy API is entirely analogous.importtorchfromgeom_median.torchimportcompute_geometric_median# PyTorch API# from geom_median.numpy import compute_geometric_median  # NumPy APIFor the simplest use case, supply a list of tensors:n=10# Number of vectorsd=25# dimensionality of each vectorpoints=[torch.rand(d)for_inrange(n)]# list of n tensors of shape (d,)# The shape of each tensor is the same and can be arbitrary (not necessarily 1-dimensional)weights=torch.rand(n)# non-negative weights of shape (n,)out=compute_geometric_median(points,weights)# Access the median via `out.median`, which has the same shape as the points, i.e., (d,)The termination condition can be examined throughout.termination, which gives a message such as"function value converged within tolerance"or"maximum iterations reached".We also support a use case where each point is given by list of tensors.
For instance, each point is the list of parameters of atorch.nn.Modulefor instance aspoint = list(module.parameters()).
In this case, this is equivalent to flattening and concatenating all the tensors into a single vector viaflatted_point = torch.stack([v.view(-1) for v in point]).
This functionality can be invoked as follows:models=[torch.nn.Linear(20,10)for_inrange(n)]# a list of n modelspoints=[list(model.parameters())formodelinmodels]# list of points, where each point is a list of tensorsout=compute_geometric_median(points,weights=None)# equivalent to `weights = torch.ones(n)`.# Access the median via `out.median`, also given as a list of tensorsWe also support computing the geometric median for each component separately in the list-of-tensors format:models=[torch.nn.Linear(20,10)for_inrange(n)]# a list of n modelspoints=[list(model.parameters())formodelinmodels]# list of points, where each point is a list of tensorsout=compute_geometric_median(points,weights=None,per_component=True)# Access the median via `out.median`, also given as a list of tensorsThis per-component geometric median is equivalent in functionality toout.median[j]=compute_geometric_median([p[j]forpinpoints],weights)Backpropagation supportWhen using the PyTorch API, the resultout.median, as a function ofpoints, supports gradient backpropagation, also known as reverse-mode automatic differentiation. Here is a toy example illustrating this behavior.points=[torch.rand(d).requires_grad_(True)for_inrange(n)]# list of tensors with `requires_grad=True`out=compute_geometric_median(points,weights=None)torch.linalg.norm(out.median).backward()# call backward on any downstream function of `out.median`gradients=[p.gradforpinpoints]# gradients with respect of `points` and upstream nodes in the computation graphGPU supportSimply use as above wherepointsandweightsare CUDA tensors.Authors and ContactKrishna PillutlaSham KakadeZaid HarchaouiIn case of questions, please raise an issue on GitHub.CitationIf you found this package useful, please consider citing this paper.@article{pillutla:etal:rfa ,
  title={{Robust Aggregation for Federated Learning}},
  author={Pillutla, Krishna and  Kakade, Sham M. and Harchaoui, Zaid},
  journal={arXiv preprint},
  year={2019}
}