dp-ferret 1.0.7
FERRET: Mutual-Information Differential Privacy via grouped sign updates.
License: Apache-2.0
Requires Python: >=3.10
Keywords: differential-privacy, machine-learning, privacy, ferret

# FERRET

Modular implementation of the FERRET mechanism for mutual-information differential privacy in large language model fine-tuning. The refactor converts the original research notebooks into a reusable Python package with configuration-driven experiments, privacy accounting, membership-inference tooling, and GPU-ready training loops.

## Features
- **Config-driven runs:** Dataclass-backed configuration models with YAML overlays (`ferret.config`).
- **MI-DP accounting:** Standalone FERRET privacy accountant with unit tests (`ferret.privacy`).
- **Sign update mechanism:** Modular implementation of grouped sign projections ready for CPU or GPU execution (`ferret.mechanisms`).
- **Trainer integration:** Hugging Face `SFTTrainer` wrapper that enforces FERRET updates and logs privacy diagnostics (`ferret.training`).
- **Experiment suite:** Scripts for multi-model sweeps, perplexity evaluation, and membership inference attacks (`ferret.experiments`, `scripts/`).
- **Optional telemetry:** Opt-in Weights & Biases logging for training, evaluation, and sweep runs.

## Getting Started
**Install from PyPI**

```bash
python -m pip install dp-ferret
```

You can now import the trainer, utilities, and bundled config presets directly:

```python
from ferret import FERRETSFTTrainer, resolve_variant, compute_perplexity, load_run_config

config = load_run_config("defaults")  # uses the packaged defaults.yaml preset
```

**Install for local development**

```bash
python -m venv .venv
source .venv/bin/activate
pip install -e .[dev]

# Optional: run unit tests
pytest
```

## CLI Entrypoints
- `scripts/ferret_train.py` – fine-tunes a model with FERRET and saves the resulting checkpoint.
- `scripts/ferret_evaluate.py` – loads a saved checkpoint, reports train/test perplexity, and (optionally) runs MIA.
- `scripts/ferret_sweep.py` – reproduces the multi-model, multi-ε study.
- `scripts/ferret_mia.py` – standalone membership-inference analysis for a checkpoint.

Every CLI accepts `--config` (default `configs/defaults.yaml`) so you can override datasets, epochs, and FERRET settings. Hugging Face tokens go through `--hf_token` or the `HF_TOKEN` environment variable.

### Common workflows

**Train & save a checkpoint**

```bash
python scripts/ferret_train.py \
  --config configs/defaults.yaml \
  --hf_token $HF_TOKEN \
  --save-dir results/gpt2_run
```

Defaults mirror the paper’s pipeline: TinyPixel/orca-mini, 3 epochs (~600 steps), and `dynamic_buckets="eighth"`. Use `--ferret_variant` to try `MAX`, `EIGHTH`, or `2`-grouping schemes.
The script prints train/eval perplexities and saves the model + tokenizer in the chosen directory (default `./results/checkpoint-last`).

**Run the FERRET sweep from the paper**

```bash
python scripts/ferret_sweep.py \
  --config configs/defaults.yaml \
  --output_dir results/paper \
  --models openai-community/gpt2 TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
          bigscience/bloom-560m cognitivecomputations/smollm-360m \
          deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B \
  --epsilons 0.1 0.5 1.0 2.0 \
  --epochs 1 3 5 \
  --variants MAX EIGHTH 2 \
  --hf_token $HF_TOKEN
```

The runner writes metrics and metadata to `results/paper/<model>__<variant>__eps_<value>/` and saves trained checkpoints under the configured `output_dir`.

**Evaluate an existing FERRET run**

```bash
python scripts/ferret_evaluate.py \
  --model-path results/gpt2_run \
  --config configs/defaults.yaml \
  --hf_token $HF_TOKEN \
  --reference-model openai-community/gpt2
```

The evaluator reports train/test perplexity, runs both the basic confidence MIA and the advanced reference-enhanced MIA (unless `--skip-mia`), and saves metrics to `evaluation_results/<checkpoint_name>/metrics.json`.
It also generates text samples (configurable via CLI) and writes them to `samples.jsonl`; if `vllm` is installed the evaluator automatically prefers it for faster batched generation and falls back to standard Hugging Face sampling otherwise.

### Dataset customization

The defaults fine-tune on `TinyPixel/orca-mini`, matching the paper’s setup (10k examples with a 50/50 train/test split). To substitute a different Hugging Face dataset, edit `configs/defaults.yaml` or provide overrides on the CLI:

```bash
python scripts/ferret_train.py --config configs/defaults.yaml \
  --dataset.dataset_name my/dataset --dataset.split train[:5000]
```

### Weights & Biases logging

All entrypoints can stream metrics to W&B when enabled in the config or via CLI:

```yaml
# configs/defaults.yaml
wandb:
  enabled: true
  project: ferret
  entity: your-team
  run_name: null        # optional; scripts will auto-name runs
  tags: ["gpt2", "ferret"]
  notes: null
  mode: online          # or "offline"
```

Override per-invocation with `--wandb-enabled` / `--wandb-enabled false`. The scripts log trainer metrics, FERRET privacy diagnostics, perplexities, and MIA scores; saved `metrics.json` files are also uploaded to the run. Set `FERRET_WANDB_DISABLED=1` to force local-only execution regardless of config.

## Project Layout
```
ferret/
  config/        # configuration dataclasses & loaders
  data/          # dataset utilities
  experiments/   # experiment runner, metrics, storage, MIA tools
  mechanisms/    # FERRET sign projection
  privacy/       # MI-DP accountant
  training/      # trainer wrapper, callbacks, factories
  utils/         # logging, seeding, torch helpers
scripts/         # CLI entrypoints
tests/           # pytest suite for accountant/mechanism utilities
docs/            # reproduction guide and supplemental docs
```

## Citation
If you build on this work, please cite:

```
@article{zagardo2025ferret,
  title   = {FERRET: Private Deep Learning Faster And Better Than DPSGD},
  author  = {David Zagardo},
  journal = {arXiv},
  year    = {2025},
  eprint  = {2506.05416},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CR},
  url     = {https://doi.org/10.48550/arXiv.2506.05416}
}
```

## License
Apache-2.0.