# JAX-Privacy: A library for differentially private machine learning

arXiv: https://arxiv.org/abs/2602.17861 · HTML: https://mlsanity.com/paper/2602.17861/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 25 references that are in the map.
- First posted on arXiv 19 Feb 2026.

## Builds on (most cited first, 19 of 19)

- [Deep Learning with Differential Privacy](https://mlsanity.com/paper/1607.00133/) (2016)
- [Opacus: User-Friendly Differential Privacy Library in PyTorch](https://mlsanity.com/paper/2109.12298/) (2021)
- [Privacy Amplification by Subsampling: Tight Analyses via Couplings and Divergences](https://mlsanity.com/paper/1807.01647/) (2018)
- [Auditing Differentially Private Machine Learning: How Private is Private SGD?](https://mlsanity.com/paper/2006.07709/) (2020)
- [Adversary Instantiation: Lower Bounds for Differentially Private Machine Learning](https://mlsanity.com/paper/2101.04535/) (2021)
- [Privacy Auditing with One (1) Training Run](https://mlsanity.com/paper/2305.08846/) (2023)
- [Differentially Private Optimization on Large Model at Small Cost](https://mlsanity.com/paper/2210.00038/) (2022)
- [(Amplified) Banded Matrix Factorization: A unified approach to private training](https://mlsanity.com/paper/2306.08153/) (2023)
- [Scaling up Differentially Private Deep Learning with Fast Per-Example Gradient Clipping](https://mlsanity.com/paper/2009.03106/) (2020)
- [Debugging Differential Privacy: A Case Study for Privacy Auditing](https://mlsanity.com/paper/2202.12219/) (2022)
- [Auditing f-Differential Privacy in One Run](https://mlsanity.com/paper/2410.22235/) (2024)
- [Fine-Tuning Large Language Models with User-Level Differential Privacy](https://mlsanity.com/paper/2407.07737/) (2024)
- [Near Exact Privacy Amplification for Matrix Mechanisms](https://mlsanity.com/paper/2410.06266/) (2024)
- [Scalable DP-SGD: Shuffling vs. Poisson Subsampling](https://mlsanity.com/paper/2411.04205/) (2024)
- [VaultGemma: A Differentially Private Gemma Model](https://mlsanity.com/paper/2510.15001/) (2025)
- [Balls-and-Bins Sampling for DP-SGD](https://mlsanity.com/paper/2412.16802/) (2024)
- [Correlated Noise Mechanisms for Differentially Private Learning](https://mlsanity.com/paper/2506.08201/) (2025)
- [Sparsity-Preserving Differentially Private Training of Large Embedding Models](https://mlsanity.com/paper/2311.08357/) (2023)
- [Tighter Privacy Analysis for Truncated Poisson Sampling](https://mlsanity.com/paper/2508.15089/) (2025)

## Built on by (most cited first, 3 of 3)

- [Beyond Square Roots: A Memory-Efficient Explicit Factorization for Multi-Epoch Private Learning](https://mlsanity.com/paper/2605.18379/) (2026)
- [ContinuousBench: Can Differentially Private Synthetic Text Improve Capabilities?](https://mlsanity.com/paper/2606.01849/) (2026)
- [Can R'enyi Accountants Beat Lossless Accountants for DP-SGD?](https://mlsanity.com/paper/2610.17120/) (2026)

## Abstract (arXiv)

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.