Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,613 @@
|
||||
//! Differential privacy mechanisms for federated learning.
|
||||
//!
|
||||
//! This module implements privacy-preserving mechanisms including gradient
|
||||
//! clipping, noise addition, and privacy budget accounting.
|
||||
|
||||
use rand::Rng;
|
||||
|
||||
use fedmed_shared::{PrivacyAccountantType, PrivacyConfig};
|
||||
|
||||
/// Differential privacy mechanism for adding noise to gradients.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DifferentialPrivacy {
|
||||
/// Privacy configuration.
|
||||
config: PrivacyConfig,
|
||||
/// Random number generator seed (reserved for reproducibility).
|
||||
#[allow(dead_code)]
|
||||
seed: u64,
|
||||
}
|
||||
|
||||
impl DifferentialPrivacy {
|
||||
/// Create a new differential privacy mechanism.
|
||||
#[must_use]
|
||||
pub fn new(config: PrivacyConfig) -> Self {
|
||||
Self { config, seed: 42 }
|
||||
}
|
||||
|
||||
/// Add Gaussian noise to gradients for differential privacy.
|
||||
#[must_use]
|
||||
pub fn add_noise(&self, gradients: &[f32]) -> Vec<f32> {
|
||||
let mut rng = rand::thread_rng();
|
||||
let sigma = self.config.noise_multiplier as f32 * self.config.clip_norm as f32;
|
||||
|
||||
gradients
|
||||
.iter()
|
||||
.map(|g| {
|
||||
// Box-Muller transform for Gaussian noise
|
||||
let u1: f32 = rng.r#gen::<f32>();
|
||||
let u2: f32 = rng.r#gen::<f32>();
|
||||
let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
|
||||
g + noise * sigma
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Add noise using local differential privacy (for client-side).
|
||||
#[must_use]
|
||||
pub fn add_local_noise(&self, gradients: &[f32]) -> Vec<f32> {
|
||||
if !self.config.local_dp {
|
||||
return gradients.to_vec();
|
||||
}
|
||||
|
||||
let mut rng = rand::thread_rng();
|
||||
// Local DP typically uses higher noise
|
||||
let sigma = self.config.noise_multiplier as f32 * self.config.clip_norm as f32 * 2.0;
|
||||
|
||||
gradients
|
||||
.iter()
|
||||
.map(|g| {
|
||||
let u1: f32 = rng.r#gen::<f32>();
|
||||
let u2: f32 = rng.r#gen::<f32>();
|
||||
let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
|
||||
g + noise * sigma
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Get the noise standard deviation.
|
||||
#[must_use]
|
||||
pub fn noise_std(&self) -> f64 {
|
||||
self.config.noise_multiplier * self.config.clip_norm
|
||||
}
|
||||
|
||||
/// Get the privacy configuration.
|
||||
#[must_use]
|
||||
pub fn config(&self) -> &PrivacyConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Update the noise multiplier.
|
||||
pub fn set_noise_multiplier(&mut self, multiplier: f64) {
|
||||
self.config.noise_multiplier = multiplier;
|
||||
}
|
||||
}
|
||||
|
||||
/// Gradient clipper for bounding sensitivity.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GradientClipper {
|
||||
/// Maximum L2 norm for gradients.
|
||||
clip_norm: f64,
|
||||
/// Whether to use per-sample clipping (reserved for advanced use).
|
||||
#[allow(dead_code)]
|
||||
per_sample: bool,
|
||||
}
|
||||
|
||||
impl Default for GradientClipper {
|
||||
fn default() -> Self {
|
||||
Self::new(1.0)
|
||||
}
|
||||
}
|
||||
|
||||
impl GradientClipper {
|
||||
/// Create a new gradient clipper.
|
||||
#[must_use]
|
||||
pub fn new(clip_norm: f64) -> Self {
|
||||
Self {
|
||||
clip_norm,
|
||||
per_sample: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create a per-sample gradient clipper.
|
||||
#[must_use]
|
||||
pub fn per_sample(clip_norm: f64) -> Self {
|
||||
Self {
|
||||
clip_norm,
|
||||
per_sample: true,
|
||||
}
|
||||
}
|
||||
|
||||
/// Clip gradients to bound their L2 norm.
|
||||
#[must_use]
|
||||
pub fn clip(&self, gradients: &[f32]) -> Vec<f32> {
|
||||
let norm = self.l2_norm(gradients);
|
||||
|
||||
if norm <= self.clip_norm as f32 {
|
||||
return gradients.to_vec();
|
||||
}
|
||||
|
||||
let scale = self.clip_norm as f32 / norm;
|
||||
gradients.iter().map(|g| g * scale).collect()
|
||||
}
|
||||
|
||||
/// Clip gradients with adaptive clipping.
|
||||
#[must_use]
|
||||
pub fn adaptive_clip(&self, gradients: &[f32], quantile: f32) -> Vec<f32> {
|
||||
// Compute per-coordinate absolute values
|
||||
let mut abs_grads: Vec<f32> = gradients.iter().map(|g| g.abs()).collect();
|
||||
abs_grads.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
|
||||
// Find the quantile threshold
|
||||
let idx = ((abs_grads.len() as f32 * quantile) as usize).min(abs_grads.len() - 1);
|
||||
let threshold = abs_grads[idx];
|
||||
|
||||
// Clip to threshold
|
||||
gradients
|
||||
.iter()
|
||||
.map(|g| g.clamp(-threshold, threshold))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Compute L2 norm of gradients.
|
||||
#[must_use]
|
||||
pub fn l2_norm(&self, gradients: &[f32]) -> f32 {
|
||||
gradients.iter().map(|g| g * g).sum::<f32>().sqrt()
|
||||
}
|
||||
|
||||
/// Compute L1 norm of gradients.
|
||||
#[must_use]
|
||||
pub fn l1_norm(&self, gradients: &[f32]) -> f32 {
|
||||
gradients.iter().map(|g| g.abs()).sum()
|
||||
}
|
||||
|
||||
/// Compute infinity norm of gradients.
|
||||
#[must_use]
|
||||
pub fn linf_norm(&self, gradients: &[f32]) -> f32 {
|
||||
gradients
|
||||
.iter()
|
||||
.map(|g| g.abs())
|
||||
.fold(0.0_f32, f32::max)
|
||||
}
|
||||
|
||||
/// Get the clip norm.
|
||||
#[must_use]
|
||||
pub fn clip_norm(&self) -> f64 {
|
||||
self.clip_norm
|
||||
}
|
||||
|
||||
/// Set the clip norm.
|
||||
pub fn set_clip_norm(&mut self, clip_norm: f64) {
|
||||
self.clip_norm = clip_norm;
|
||||
}
|
||||
}
|
||||
|
||||
/// Privacy accountant for tracking privacy budget.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PrivacyAccountant {
|
||||
/// Privacy configuration.
|
||||
config: PrivacyConfig,
|
||||
/// Total privacy spent (epsilon).
|
||||
epsilon_spent: f64,
|
||||
/// Number of composition steps.
|
||||
num_steps: usize,
|
||||
/// Sample rate for subsampling amplification.
|
||||
sample_rate: f64,
|
||||
/// Accountant type.
|
||||
accountant_type: PrivacyAccountantType,
|
||||
}
|
||||
|
||||
impl PrivacyAccountant {
|
||||
/// Create a new privacy accountant.
|
||||
#[must_use]
|
||||
pub fn new(config: PrivacyConfig) -> Self {
|
||||
Self {
|
||||
accountant_type: config.accountant,
|
||||
config,
|
||||
epsilon_spent: 0.0,
|
||||
num_steps: 0,
|
||||
sample_rate: 1.0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Create an accountant with a specific sample rate.
|
||||
#[must_use]
|
||||
pub fn with_sample_rate(config: PrivacyConfig, sample_rate: f64) -> Self {
|
||||
Self {
|
||||
accountant_type: config.accountant,
|
||||
config,
|
||||
epsilon_spent: 0.0,
|
||||
num_steps: 0,
|
||||
sample_rate: sample_rate.clamp(0.0, 1.0),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record a privacy-consuming step and return the epsilon spent.
|
||||
pub fn step(&mut self, num_clients: usize) -> f64 {
|
||||
self.num_steps += 1;
|
||||
|
||||
// Compute epsilon for this step based on accountant type
|
||||
let step_epsilon = match self.accountant_type {
|
||||
PrivacyAccountantType::Moments => self.moments_epsilon(num_clients),
|
||||
PrivacyAccountantType::RDP => self.rdp_epsilon(num_clients),
|
||||
PrivacyAccountantType::GDP => self.gdp_epsilon(num_clients),
|
||||
PrivacyAccountantType::PLD => self.pld_epsilon(num_clients),
|
||||
};
|
||||
|
||||
self.epsilon_spent += step_epsilon;
|
||||
step_epsilon
|
||||
}
|
||||
|
||||
/// Compute epsilon using moments accountant.
|
||||
fn moments_epsilon(&self, _num_clients: usize) -> f64 {
|
||||
// Simplified moments accountant
|
||||
let sigma = self.config.noise_multiplier;
|
||||
let q = self.sample_rate;
|
||||
|
||||
// Basic composition: epsilon per step
|
||||
if sigma > 0.0 {
|
||||
q * (1.0 / sigma.powi(2))
|
||||
} else {
|
||||
f64::INFINITY
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute epsilon using Renyi differential privacy.
|
||||
fn rdp_epsilon(&self, _num_clients: usize) -> f64 {
|
||||
// Simplified RDP accountant
|
||||
let sigma = self.config.noise_multiplier;
|
||||
let q = self.sample_rate;
|
||||
let alpha = 2.0; // RDP order
|
||||
|
||||
if sigma > 0.0 {
|
||||
// RDP bound with subsampling: q * alpha / (2 * sigma^2)
|
||||
let rdp = q * alpha / (2.0 * sigma.powi(2));
|
||||
// Convert to (epsilon, delta)-DP
|
||||
let delta = self.config.delta;
|
||||
rdp + (1.0 / delta).ln() / (alpha - 1.0)
|
||||
} else {
|
||||
f64::INFINITY
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute epsilon using Gaussian differential privacy.
|
||||
fn gdp_epsilon(&self, _num_clients: usize) -> f64 {
|
||||
// Simplified GDP accountant
|
||||
let sigma = self.config.noise_multiplier;
|
||||
|
||||
if sigma > 0.0 {
|
||||
// Central limit theorem approximation
|
||||
1.0 / sigma
|
||||
} else {
|
||||
f64::INFINITY
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute epsilon using privacy loss distribution.
|
||||
fn pld_epsilon(&self, _num_clients: usize) -> f64 {
|
||||
// Simplified PLD accountant (falls back to RDP for simplicity)
|
||||
self.rdp_epsilon(_num_clients)
|
||||
}
|
||||
|
||||
/// Get total epsilon spent.
|
||||
#[must_use]
|
||||
pub fn epsilon_spent(&self) -> f64 {
|
||||
self.epsilon_spent
|
||||
}
|
||||
|
||||
/// Get number of steps.
|
||||
#[must_use]
|
||||
pub fn num_steps(&self) -> usize {
|
||||
self.num_steps
|
||||
}
|
||||
|
||||
/// Check if privacy budget is exhausted.
|
||||
#[must_use]
|
||||
pub fn is_budget_exhausted(&self) -> bool {
|
||||
if let Some(target) = self.config.target_epsilon {
|
||||
self.epsilon_spent >= target
|
||||
} else {
|
||||
false
|
||||
}
|
||||
}
|
||||
|
||||
/// Get remaining privacy budget.
|
||||
#[must_use]
|
||||
pub fn remaining_budget(&self) -> Option<f64> {
|
||||
self.config
|
||||
.target_epsilon
|
||||
.map(|target| (target - self.epsilon_spent).max(0.0))
|
||||
}
|
||||
|
||||
/// Estimate number of remaining steps.
|
||||
#[must_use]
|
||||
pub fn estimate_remaining_steps(&self) -> Option<usize> {
|
||||
if self.num_steps == 0 || self.epsilon_spent <= 0.0 {
|
||||
return None;
|
||||
}
|
||||
|
||||
let epsilon_per_step = self.epsilon_spent / self.num_steps as f64;
|
||||
self.remaining_budget()
|
||||
.map(|remaining| (remaining / epsilon_per_step) as usize)
|
||||
}
|
||||
|
||||
/// Reset the accountant.
|
||||
pub fn reset(&mut self) {
|
||||
self.epsilon_spent = 0.0;
|
||||
self.num_steps = 0;
|
||||
}
|
||||
|
||||
/// Get the target epsilon.
|
||||
#[must_use]
|
||||
pub fn target_epsilon(&self) -> Option<f64> {
|
||||
self.config.target_epsilon
|
||||
}
|
||||
|
||||
/// Get delta.
|
||||
#[must_use]
|
||||
pub fn delta(&self) -> f64 {
|
||||
self.config.delta
|
||||
}
|
||||
}
|
||||
|
||||
/// Secure noise generation for distributed settings.
|
||||
#[derive(Debug)]
|
||||
pub struct SecureNoiseGenerator {
|
||||
/// Seed for reproducibility.
|
||||
seed: u64,
|
||||
/// Number of parties.
|
||||
num_parties: usize,
|
||||
}
|
||||
|
||||
impl SecureNoiseGenerator {
|
||||
/// Create a new secure noise generator.
|
||||
#[must_use]
|
||||
pub fn new(seed: u64, num_parties: usize) -> Self {
|
||||
Self { seed, num_parties }
|
||||
}
|
||||
|
||||
/// Generate correlated noise that sums to Gaussian noise.
|
||||
/// Each party generates their share independently.
|
||||
#[must_use]
|
||||
pub fn generate_noise_share(&self, party_id: usize, size: usize, sigma: f32) -> Vec<f32> {
|
||||
use rand::SeedableRng;
|
||||
let mut rng = rand::rngs::StdRng::seed_from_u64(self.seed.wrapping_add(party_id as u64));
|
||||
|
||||
(0..size)
|
||||
.map(|_| {
|
||||
let u1: f32 = rng.r#gen::<f32>();
|
||||
let u2: f32 = rng.r#gen::<f32>();
|
||||
let noise = (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos();
|
||||
noise * sigma / (self.num_parties as f32).sqrt()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Combine noise shares from all parties.
|
||||
#[must_use]
|
||||
pub fn combine_shares(&self, shares: &[Vec<f32>]) -> Vec<f32> {
|
||||
if shares.is_empty() {
|
||||
return vec![];
|
||||
}
|
||||
|
||||
let size = shares[0].len();
|
||||
let mut combined = vec![0.0_f32; size];
|
||||
|
||||
for share in shares {
|
||||
for (i, &s) in share.iter().enumerate() {
|
||||
if i < combined.len() {
|
||||
combined[i] += s;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
combined
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn sample_config() -> PrivacyConfig {
|
||||
PrivacyConfig {
|
||||
epsilon: 8.0,
|
||||
delta: 1e-5,
|
||||
clip_norm: 1.0,
|
||||
noise_multiplier: 1.1,
|
||||
local_dp: false,
|
||||
accountant: PrivacyAccountantType::RDP,
|
||||
target_epsilon: Some(10.0),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_differential_privacy() {
|
||||
let config = sample_config();
|
||||
let dp = DifferentialPrivacy::new(config);
|
||||
|
||||
let gradients = vec![0.1, 0.2, 0.3, 0.4, 0.5];
|
||||
let noisy = dp.add_noise(&gradients);
|
||||
|
||||
assert_eq!(noisy.len(), gradients.len());
|
||||
// Noisy gradients should be different from original
|
||||
assert!(
|
||||
noisy
|
||||
.iter()
|
||||
.zip(gradients.iter())
|
||||
.any(|(n, g)| (n - g).abs() > 0.0)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_local_dp() {
|
||||
let mut config = sample_config();
|
||||
config.local_dp = true;
|
||||
let dp = DifferentialPrivacy::new(config);
|
||||
|
||||
let gradients = vec![0.1, 0.2, 0.3];
|
||||
let noisy = dp.add_local_noise(&gradients);
|
||||
|
||||
assert_eq!(noisy.len(), gradients.len());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_gradient_clipping() {
|
||||
let clipper = GradientClipper::new(1.0);
|
||||
|
||||
// Large gradients should be clipped
|
||||
let large_grads = vec![10.0, 10.0, 10.0];
|
||||
let clipped = clipper.clip(&large_grads);
|
||||
|
||||
let clipped_norm = clipper.l2_norm(&clipped);
|
||||
assert!(clipped_norm <= 1.01); // Allow small floating point error
|
||||
|
||||
// Small gradients should not be affected
|
||||
let small_grads = vec![0.1, 0.1, 0.1];
|
||||
let not_clipped = clipper.clip(&small_grads);
|
||||
assert_eq!(not_clipped, small_grads);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_adaptive_clipping() {
|
||||
let clipper = GradientClipper::new(1.0);
|
||||
|
||||
let gradients = vec![0.1, 0.5, 1.0, 2.0, 5.0];
|
||||
let clipped = clipper.adaptive_clip(&gradients, 0.5);
|
||||
|
||||
// Median value should be the threshold
|
||||
assert!(clipped.iter().all(|&g| g.abs() <= 1.01));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_norms() {
|
||||
let clipper = GradientClipper::new(1.0);
|
||||
let gradients = vec![3.0, 4.0];
|
||||
|
||||
assert!((clipper.l2_norm(&gradients) - 5.0).abs() < 0.001);
|
||||
assert!((clipper.l1_norm(&gradients) - 7.0).abs() < 0.001);
|
||||
assert!((clipper.linf_norm(&gradients) - 4.0).abs() < 0.001);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_privacy_accountant() {
|
||||
let config = sample_config();
|
||||
let mut accountant = PrivacyAccountant::new(config);
|
||||
|
||||
assert_eq!(accountant.epsilon_spent(), 0.0);
|
||||
assert_eq!(accountant.num_steps(), 0);
|
||||
|
||||
let epsilon = accountant.step(5);
|
||||
assert!(epsilon > 0.0);
|
||||
assert_eq!(accountant.num_steps(), 1);
|
||||
assert!(accountant.epsilon_spent() > 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_budget_exhaustion() {
|
||||
let mut config = sample_config();
|
||||
config.target_epsilon = Some(0.1);
|
||||
config.noise_multiplier = 0.1; // Low noise = high epsilon per step
|
||||
|
||||
let mut accountant = PrivacyAccountant::new(config);
|
||||
|
||||
// Take many steps until budget is exhausted
|
||||
for _ in 0..1000 {
|
||||
accountant.step(5);
|
||||
if accountant.is_budget_exhausted() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
assert!(accountant.is_budget_exhausted());
|
||||
assert!(accountant.remaining_budget().unwrap() <= 0.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remaining_steps_estimate() {
|
||||
let config = sample_config();
|
||||
let mut accountant = PrivacyAccountant::new(config);
|
||||
|
||||
// Take some steps
|
||||
for _ in 0..10 {
|
||||
accountant.step(5);
|
||||
}
|
||||
|
||||
let remaining = accountant.estimate_remaining_steps();
|
||||
assert!(remaining.is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_accountant_reset() {
|
||||
let config = sample_config();
|
||||
let mut accountant = PrivacyAccountant::new(config);
|
||||
|
||||
accountant.step(5);
|
||||
assert!(accountant.epsilon_spent() > 0.0);
|
||||
|
||||
accountant.reset();
|
||||
assert_eq!(accountant.epsilon_spent(), 0.0);
|
||||
assert_eq!(accountant.num_steps(), 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_secure_noise_generator() {
|
||||
let generator = SecureNoiseGenerator::new(42, 3);
|
||||
|
||||
let shares: Vec<_> = (0..3)
|
||||
.map(|party| generator.generate_noise_share(party, 10, 1.0))
|
||||
.collect();
|
||||
|
||||
let combined = generator.combine_shares(&shares);
|
||||
assert_eq!(combined.len(), 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_different_accountant_types() {
|
||||
for accountant_type in [
|
||||
PrivacyAccountantType::Moments,
|
||||
PrivacyAccountantType::RDP,
|
||||
PrivacyAccountantType::GDP,
|
||||
PrivacyAccountantType::PLD,
|
||||
] {
|
||||
let mut config = sample_config();
|
||||
config.accountant = accountant_type;
|
||||
|
||||
let mut accountant = PrivacyAccountant::new(config);
|
||||
let epsilon = accountant.step(5);
|
||||
|
||||
assert!(epsilon > 0.0);
|
||||
assert!(epsilon < f64::INFINITY);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_sample_rate() {
|
||||
let config = sample_config();
|
||||
let mut accountant_full = PrivacyAccountant::new(config.clone());
|
||||
let mut accountant_sub = PrivacyAccountant::with_sample_rate(config, 0.1);
|
||||
|
||||
let eps_full = accountant_full.step(5);
|
||||
let eps_sub = accountant_sub.step(5);
|
||||
|
||||
// Subsampled should have lower epsilon (privacy amplification)
|
||||
assert!(eps_sub < eps_full);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_clipper_set_norm() {
|
||||
let mut clipper = GradientClipper::new(1.0);
|
||||
assert_eq!(clipper.clip_norm(), 1.0);
|
||||
|
||||
clipper.set_clip_norm(2.0);
|
||||
assert_eq!(clipper.clip_norm(), 2.0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_dp_set_noise_multiplier() {
|
||||
let config = sample_config();
|
||||
let mut dp = DifferentialPrivacy::new(config);
|
||||
|
||||
dp.set_noise_multiplier(2.0);
|
||||
assert_eq!(dp.config().noise_multiplier, 2.0);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user