Google DeepMind launches JAX-Privacy 1.0 for Scalable Privacy-Preserving AI

Google DeepMind and Google Research have announced the release of JAX-Privacy 1.0, an open-source library designed to enable differentially private machine learning (DP-ML) at scale. The library, built on the high-performance computing framework JAX, is intended to help researchers and developers train large AI models while safeguarding individual privacy.

AI models rely on large, high-quality datasets to achieve accuracy and generalizability, but the use of sensitive data raises privacy concerns. To address this, JAX-Privacy provides tools that allow private training to be implemented efficiently without compromising data confidentiality.

JAX, introduced in 2020, is a numerical computing library optimized for large-scale machine learning. It offers automatic differentiation, just-in-time compilation, and seamless scaling across multiple accelerators. Its ecosystem includes domain-specific libraries such as Flax, which simplifies neural network construction, and Optax, which provides advanced optimizers. Building on JAX, JAX-Privacy enables the implementation of differentially private algorithms for modern distributed training workflows.

Differential privacy (DP) is considered the gold standard for quantifying and bounding privacy leakage, ensuring that the output of an algorithm remains nearly identical whether or not a single individual’s data is included. However, implementing DP at the scale required for modern AI models is complex. Traditional approaches like differentially private stochastic gradient descent (DP-SGD) require per-example gradient clipping, noise addition, and customized batching, which are computationally intensive.

JAX-Privacy 1.0 addresses these challenges by providing robust building blocks for gradient clipping, correlated noise generation, and batch construction, all optimized for distributed environments. Advanced DP algorithms, including DP matrix factorization, are also supported, allowing researchers to experiment with state-of-the-art privacy-preserving techniques.

The library is designed for scalability. Leveraging JAX’s native parallelism features such as vmap and shard_map, JAX-Privacy allows large models to be trained across multiple GPUs or TPUs without complex custom code. Tools like micro-batching and padding make it easier to handle variable-sized batches, enabling better privacy-utility trade-offs.

Auditing and correctness are emphasized. JAX-Privacy integrates Google’s DP accounting library to ensure mathematically correct noise calibration. Metrics for empirical privacy loss are provided, and auditing techniques, including the injection of "canaries" (known data points), allow privacy guarantees to be tested and validated.

The practical applications of JAX-Privacy have already been demonstrated in training large language models (LLMs), including VaultGemma, a state-of-the-art differentially private model. Developers can now fine-tune models in the Gemma family using Keras, with examples provided for tasks such as dialogue summarization and synthetic data generation. This makes it easier to build privacy-preserving AI applications for healthcare, finance, personalized recommendations, and other sensitive domains.

JAX-Privacy is primarily targeted at researchers, developers, and organizations handling sensitive data, but its open-source release allows the broader AI community to benefit. By simplifying the integration of differential privacy, it lowers barriers to responsible AI development and helps ensure that user data is protected.

Google DeepMind and Google Research have emphasized that continuous support and updates will be provided, incorporating new research advances and responding to community needs. The library is available on GitHub and via PIP, allowing developers to start building privacy-preserving ML models immediately.