Private spectral whitening front-end / private_spectral_whitening.py
Unverified
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()