JAX-Privacy: A library for differentially private machine learning

TL;DR

JAX-Privacy is a versatile library that simplifies implementing differentially private machine learning by offering modular, verified primitives, catering to both researchers and practitioners with a focus on usability and efficiency.

Contribution

It introduces a comprehensive, verified library with modular primitives for DP ML, integrating recent research and supporting customization and ease of use.

Findings

Provides verified, modular primitives for DP ML components

Supports both research and practical deployment needs

Integrates recent advances in differential privacy mechanisms

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.