Private spectral whitening front-end / private_spectral_whitening.py

Unverified

Raw ⬇ ZIP
  1"""MVP for a locally-private spectral whitening front-end.
  2
  3The Gaussian noise mechanism here is the implementation-plan provisional
  4mechanism, not the paper's optimal Gaussian-specific LDP estimator.
  5"""
  6import json
  7from pathlib import Path
  8import numpy as np
  9
 10SEED = 145
 11rng = np.random.default_rng(SEED)
 12
 13
 14def ar1(n, t, phi=0.85, rng=None):
 15    rng = np.random.default_rng() if rng is None else rng
 16    x = np.zeros((n, t))
 17    x[:, 0] = rng.normal(size=n)
 18    e = rng.normal(size=(n, t))
 19    for k in range(1, t):
 20        x[:, k] = phi * x[:, k - 1] + np.sqrt(1 - phi * phi) * e[:, k]
 21    return x
 22
 23
 24def autocov_features(x, max_lag, B):
 25    """Clipped, centered per-sequence autocovariances, lags 0..max_lag."""
 26    z = np.clip(x, -B, B)
 27    z = z - z.mean(axis=1, keepdims=True)
 28    n, t = z.shape
 29    out = np.empty((n, max_lag + 1))
 30    for h in range(max_lag + 1):
 31        out[:, h] = np.mean(z[:, :t-h] * z[:, h:], axis=1)
 32    return out
 33
 34
 35def estimate_spectrum(x, alpha=None, B=2.5, smooth=3, max_lag=None, rng=None):
 36    """Estimate gamma then FFT it; alpha=None is the nonprivate control."""
 37    rng = np.random.default_rng() if rng is None else rng
 38    n, t = x.shape
 39    max_lag = t // 2 if max_lag is None else max_lag
 40    g = autocov_features(x, max_lag, B)
 41    if alpha is not None:
 42        # Clipping makes each feature bounded by B^2. This is a deliberately
 43        # simple Gaussian-noise local mechanism, as proposed for the MVP.
 44        sigma = 2.0 * B * B / alpha
 45        g = g + rng.normal(scale=sigma, size=g.shape)
 46    gamma = np.zeros(t)
 47    gamma[:max_lag + 1] = g.mean(axis=0)
 48    # A real even covariance sequence is required for a real nonnegative PSD.
 49    for h in range(1, max_lag + 1):
 50        if h < t:
 51            gamma[-h] = gamma[h]
 52    f = np.real(np.fft.fft(gamma))
 53    if smooth > 1:
 54        pad = np.r_[f[-smooth//2:], f, f[:smooth//2]]
 55        f = np.convolve(pad, np.ones(smooth) / smooth, mode='valid')[:t]
 56    f = 0.5 * (f + f[::-1])  # numerical symmetrization (FFT-bin reversal)
 57    return np.maximum(f, 1e-3)
 58
 59
 60def whiten(x, f, eps=0.05):
 61    z = x - x.mean(axis=1, keepdims=True)
 62    w = 1.0 / np.sqrt(np.maximum(f, eps))
 63    y = np.real(np.fft.ifft(np.fft.fft(z, axis=1) * w[None, :], axis=1))
 64    return y
 65
 66
 67def periodogram(x):
 68    z = x - x.mean(axis=1, keepdims=True)
 69    return np.mean(np.abs(np.fft.fft(z, axis=1)) ** 2, axis=0) / x.shape[1]
 70
 71
 72def autocorr_energy(x, max_lag=10):
 73    z = x - x.mean(axis=1, keepdims=True)
 74    vals = []
 75    for h in range(1, max_lag + 1):
 76        a = np.mean(z[:, :z.shape[1]-h] * z[:, h:])
 77        v = np.mean(z[:, :z.shape[1]-h] ** 2)
 78        vals.append(float(a / max(v, 1e-12)))
 79    return float(np.mean(np.square(vals))), vals
 80
 81
 82def verify_math():
 83    x = ar1(1, 128, phi=.7, rng=np.random.default_rng(7))[0]
 84    z = x - x.mean()
 85    # Wiener-Khinchin: circular autocorrelation FFT equals |FFT|^2.
 86    lhs = np.real(np.fft.ifft(np.abs(np.fft.fft(z)) ** 2)) / len(z)
 87    rhs = np.array([np.mean(z * np.roll(z, -h)) for h in range(len(z))])
 88    identity_err = float(np.max(np.abs(lhs - rhs)))
 89    f = np.maximum(np.real(np.fft.fft(rhs)), 1e-5)
 90    y = np.real(np.fft.ifft(np.fft.fft(z) / np.sqrt(f)))
 91    fy = np.abs(np.fft.fft(y)) ** 2 / len(y)
 92    # With the exact spectrum, each occupied Fourier bin is flattened.
 93    flatten_ratio = float(np.std(fy) / np.mean(fy))
 94    return identity_err, flatten_ratio
 95
 96
 97def alpha_probe():
 98    # Repeated independent datasets estimate the same AR spectrum. The
 99    # private variance should worsen rapidly as alpha decreases; fit slope.
100    t, n, reps = 48, 256, 28
101    true = np.array([0.85 ** h for h in range(t // 2 + 1)])
102    trueg = np.zeros(t); trueg[:len(true)] = true
103    for h in range(1, len(true)): trueg[-h] = true[h]
104    truef = np.maximum(np.real(np.fft.fft(trueg)), 1e-3)
105    out = {}
106    for a in [0.5, 1.0, 2.0, 4.0]:
107        errs = []
108        for r in range(reps):
109            xx = ar1(n, t, rng=np.random.default_rng(1000 + r))
110            est = estimate_spectrum(xx, alpha=a, B=2.5, smooth=3,
111                                    rng=np.random.default_rng(2000 + r))
112            errs.append(np.mean((est - truef) ** 2))
113        out[str(a)] = float(np.mean(errs))
114    # compare alpha^-4 endpoints, a descriptive check rather than a theorem test
115    ratio = out['0.5'] / out['4.0']
116    return out, float(ratio)
117
118
119def main():
120    identity_err, exact_flatten = verify_math()
121    t, ntrain, ntest = 64, 512, 512
122    train = ar1(ntrain, t, rng=np.random.default_rng(11))
123    test = ar1(ntest, t, rng=np.random.default_rng(12))
124    rows = {}
125    for name, alpha in [('raw', None), ('nonprivate', None), ('private_alpha1', 1.0), ('private_alpha4', 4.0)]:
126        if name == 'raw':
127            y = test
128        else:
129            f = estimate_spectrum(train, alpha=alpha, B=2.5, smooth=5,
130                                  rng=np.random.default_rng(300))
131            y = whiten(test, f, eps=.05)
132        energy, ac = autocorr_energy(y)
133        p = periodogram(y)
134        rows[name] = {'lag_autocorr_energy': energy,
135                      'mean_abs_lag1_10': float(np.mean(np.abs(ac))),
136                      'periodogram_cv': float(np.std(p) / max(np.mean(p), 1e-12))}
137    alpha_mse, endpoint_ratio = alpha_probe()
138    result = {'math': {'fft_identity_max_error': identity_err,
139                       'exact_filter_spectrum_cv': exact_flatten},
140              'heldout_ar1': rows,
141              'alpha_probe_mse': alpha_mse,
142              'alpha_0.5_to_alpha_4_mse_ratio': endpoint_ratio,
143              'config': {'T': t, 'train_sequences': ntrain, 'test_sequences': ntest,
144                         'phi': .85, 'B': 2.5, 'epsilon': .05, 'seed': SEED,
145                         'noise': 'Gaussian sigma=2*B^2/alpha'}}
146    Path('results.json').write_text(json.dumps(result, indent=2))
147    print(json.dumps(result, indent=2))
148
149
150if __name__ == '__main__':
151    main()