JAX Privacy documentation

JAX Privacy is an open-source library for differentially private (DP) training of machine learning models. Originally developed to support research on DP image classification within DeepMind, it has subsequently been extended to other models, data modalities, use cases, and contributors. JAX Privacy is currently being developed and maintained to support the following goals:

• Provide a production-focused API for differentially-private training of ML models in JAX and Keras.

• Enable reproducibility of DP training research done at Google.

• Provide a platform enabling external researchers to easily experiment with settings relevant to Google’s DP training ecosystem and problem set.

The library is still in development and we are actively working on improving its usability and functionality. If you have any feedback or feature requests, please don’t hesitate to contact us.

DP training is an active research area with many recent developments. It does not come for free therefore expect:

• Some accuracy decrease compared to non-DP model verions, in most of the cases it will be negligible.

• Larger training time.

• More hyperparameters to tune.

If you are unfamiliar with the following topics, we recommend reading the provided literature:

• Differential Privacy: an approach to bound the leakage of information about individuals during data processing, including training of ML models. Start with this Medium blog post, dive deeper if you want 1, 2, 3.

• JAX: numpy-like library for high-performance numerical computation and gradient calculations, successor of TensorFlow.

Note that the documentation is written with the assumption that you are familiar with DP (especially DP for machine learning) and JAX, in particular with terminology used there.

Additionally, depending on APIs you will use, you might need to get familiar with one of the following libraries:

• Flax NNX: PyTorch-like high-level library for building neural networks in JAX.

• Flax Linen: predecessor of Flax NNX, more JAX-like than Pytorch-like.

• Keras: even more high-level library for building and training ML models using different backends including JAX and PyTorch.

You can navigate the documentation in the following way:

1. Start with short Overview of the library main components to get better high-level understanding.

2. Install the library following instructions in Installation section.

3. If you want to use the library for your own use-case, choose the framework you want to use (e.g. Keras or JAX & Flax) and study its API capabilities with the provided examples.

4. If you are interested in doing some specific task (e.g. LLM fine-tuning), navigate to the corresponding example in the “Examples” section.

5. If you want to reproduce paper results, navigate to Paper Results Reproduction.

6. If you run into issues, Troubleshooting might help.

7. If you want to understand the library’s design, navigate to Library Design. It can be useful for library contributors.

8. If you want to contribute to the library, navigate to Contribution Guide.

9. If you need help, navigate to Support.

Getting Started

• Overview
• Installation
  • Option 1: Static Installation
  • Option 2: Local Installation

API Documentation

• Core Library
  • Public API
  • Matrix Factorization
  • Experimental Modules
• Keras API
  • Example Usage
  • API Reference

Examples

• DP-SGD tutorial using Flax Linen on MNIST
  • Install libraries
  • Define hyper-parameters
  • Load the MNIST dataset
  • Define the CNN model with Flax Linen
  • Create model and Flax Linen train state
  • Create DP-SGD gradient computer
  • Define loss function
  • Define evaluation function
  • Define non-DP train step
  • Define DP-SGD train step
  • Train
• Tutorial of DP-SGD LoRA fine-tuning Gemma3 in Keras on SAMSum dataset
  • Install and import dependencies
  • Login to Kaggle
  • Hyper-parameters setup
  • Training data
  • Gemma3 model setup
  • Let’s DP fine-tune Gemma3!
  • Calculate performance metrics on the test dataset
  • Results
• Tutorial: Generating Differentially Private Synthetic Data
  • Install and import dependencies
  • Hyperparameter Setup
  • Prepare the data
  • Model setup and fine-tuning
  • Synthetic Data Generation
  • Evaluation

Paper Results Reproduction

• Paper Reproductions Guide
  • Changes to the JAX Ecosystem
  • Changes to JAX-Privacy maintainers

Technical Documentation

• Sharp Edges
  • Handling Variable Batch Sizes
  • jax.vmap + jax.sharding
  • Mixed precision training
• Troubleshooting
  • Loss does not go down
• Library Design
  • NumPy vs. JAX
• Contribution Guide
  • How to Contribute
• Support