JAX-Privacy: A library for differentially private machine learning

Ryan McKenna; Galen Andrew; Borja Balle; Vadym Doroshenko; Arun Ganesh; Weiwei Kong; Alex Kurakin; Brendan McMahan

cs.LGFeb 19, 2026

JAX-Privacy: A library for differentially private machine learning

Authors: Ryan McKenna, Galen Andrew, Borja Balle, Vadym Doroshenko, Arun Ganesh, Weiwei Kong, Alex Kurakin, Brendan McMahan,

Abstract

JAX-Privacy is a library designed to simplify the deployment of robust and performant mechanisms for differentially private machine learning. Guided by design principles of usability, flexibility, and efficiency, JAX-Privacy serves both researchers requiring deep customization and practitioners who want a more out-of-the-box experience. The library provides verified, modular primitives for critical components for all aspects of the mechanism design including batch selection, gradient clipping, noise addition, accounting, and auditing, and brings together a large body of recent research on differentially private ML.

Figures & tables

		JAX-Privacy	JAX	Opacus	PyTorch
CNN	Small	31837.17 (0.50x)	63898.91 (1.00x)	8916.34 (0.14x)	58820.44 (0.92x)
	Medium	2466.76 (0.48x)	5180.58 (1.00x)	510.19 (0.10x)	5368.72 (1.04x)
	Large	437.03 (0.41x)	1057.17 (1.00x)	182.84 (0.17x)	1032.12 (0.98x)
State Space	Small	1326.49 (0.99x)	1346.26 (1.00x)	292.56 (0.22x)	1532.19 (1.14x)
	Medium	323.50 (1.00x)	322.28 (1.00x)	51.89 (0.16x)	363.28 (1.13x)
	Large	51.57 (0.80x)	64.25 (1.00x)	—	57.55 (0.90x)

Table 1: Throughput, measured as examples processed per second (higher is better), for training loops implemented with JAX-Privacy, JAX, Opacus, and PyTorch, on standard model architectures. Relative throughput compared to JAX is shown in parentheses.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Feature	TensorFlow Privacy	Opacus	JAX-Privacy
Base Framework	TensorFlow / Keras	PyTorch	JAX
Programming Model	Mixed (OO & Functional)	Object-Oriented	Purely Functional
Mechanism Logic	Optimizer Wrapper	nn.Module Hooks	vmap & grad transforms
Per-Example Gradients	Micro-batching / Fast Clip	Hooks + einsum	Automatic Vectorization
Memory Efficiency	XLA Optimization	Virtual Steps (Accumulation)	XLA Fusion (No Materialization)
Hardware Support	GPU & TPU	GPU (No Native TPU)	GPU & TPU

Appendix

Table 2: Comparison of major differentially private machine learning libraries. JAX-Privacy distinguishes itself through a purely functional design that leverages XLA for seamless compilation on TPUs and GPUs, avoiding the complexity of backward hooks used in PyTorch-based solutions.

Privacy-Preserving ML

Privacy-preserving machine learning develops methods to learn, share, and run models while limiting exposure of sensitive data.

See all 423 papers

Certification-Based Differentially Private Learning
Sep 30, 2026Mihnea Ghitu, Matthew WickerAbstract Interpretation Privacy-Preserving ML
Department of Computing Imperial College London London, United Kingdom
Enhancing Differentially Private Mechanisms via Empirical Bayes
Jun 19, 2026Minwoo Kim, Junyong Park, Sungkyu JungEmpirical Bayes Differential Privacy
Department of Statistics Seoul National University Seoul, 08826, South Korea · Department of Statistics and Institute for Data Innovation in Science Seoul National University Seoul, 08826, South Korea
Revisiting the Provable-Auditable Privacy Gap of DP-SGD
Aug 28, 2026Saloni Modi, Srivi Balaji, Yusong Zhu +2Differential Privacy Differentially Private Stochastic Gradient Descent
University of Texas at Austin · University of Waterloo and Vector Institute

Sep 30, 2026J/K move · Enter open · S save

Certification-Based Differentially Private Learning

Mihnea Ghitu, Matthew Wicker

Department of Computing Imperial College London London, United Kingdom

Abstract Interpretation Privacy-Preserving ML Differential Privacy Gaussian Differential Privacy

Open paper AskPDF arXiv

Differential privacy (DP) in machine learning is typically achieved by adding noise to model parameters (private learning) or to model outputs (private prediction). Recent work uses formal methods, namely abstract interpretation, to provide tighter privacy guarantees, but only for private prediction in classification settings. In this work, we investigate the use of formal methods as a general tool for tighter privacy analysis. First, we generalize the abstract gradient training (AGT) framework to private prediction in continuous, unbounded regression. Second, by reducing learning in parameterized models to a regression problem over the parameter space, we introduce Abstract Gradient Sampling (AGS), an algorithm that enables reachability-based analysis to provide guarantees for private learning. In both private prediction and private learning, we provide tightened privacy accounting for the AGT framework and a theoretical analysis demonstrating when our smooth sensitivity upper-bounds yield favourable privacy-utility trade-off. In practice, we validate that our regression bounds are tighter than global-sensitivity baselines on regression benchmarks, and, notably, yield the first finite privacy guarantees in settings where global prediction sensitivity is a priori unbounded. We also find that under matched conditions, our private learning algorithm can outperform standard private learners.

Abstract

Privacy-Preserving ML

Explore similar work

Certification-Based Differentially Private Learning

Enhancing Differentially Private Mechanisms via Empirical Bayes

Revisiting the Provable-Auditable Privacy Gap of DP-SGD