JAX-Privacy: A library for differentially private machine learning

2026-02-19

AI Summary

This work proposes a modular and verifiable framework for differentially private machine learning built on JAX, addressing the persistent challenge of balancing usability, flexibility, and efficiency in real-world deployments. The framework integrates core components—including batch sampling, gradient clipping, noise injection, privacy accounting, and auditing—into a unified system that supports both out-of-the-box usage and deep customization. By incorporating recent advances in the field, it significantly enhances the engineering maturity of differentially private ML, offering strong formal privacy guarantees without compromising computational efficiency or ease of use. This design enables broad applicability across diverse research explorations and practical deployment scenarios.

Technology Category

• Machine Learning: Privacy
• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Security and Privacy: Security and privacy of machine learning and AI applications
• User Modeling, Personalization and Recommendation: User privacy protection in personalized systems
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• differentially private machine learning
• privacy-preserving ML
• JAX-Privacy
• machine learning privacy

Innovation

• differentially private machine learning
• modular primitives
• gradient clipping
• privacy accounting
• JAX-Privacy