File size: 5,742 Bytes
dd90a4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """Exact scope certificate for the quadratic family in Theorem 5.5.
The paper's lower-bound proof uses two equally sized client groups with
opposite quadratic optima. This producer instantiates that construction and
computes the FedAvg dynamics with Fraction arithmetic. A sampled subset is
represented by its exact hypergeometric count of +kappa clients, which is
equivalent to summing all subsets because the update depends only on that
count.
For Li(theta) = (theta-z_i)^2/2, z_i in {-kappa,+kappa}, the global objective
has gradient theta. E local steps with step eta map theta to
q*theta + (1-q)*z_i, q=(1-eta)^E. Thus the server recursion is
theta' = q*theta + (1-q)*mean(z_i in sampled subset). The exact stationary
gradient-norm second moment follows from this affine recursion.
"""
from __future__ import annotations
import json
import sys
from fractions import Fraction
from math import comb
from pathlib import Path
if hasattr(sys, "set_int_max_str_digits"):
sys.set_int_max_str_digits(1_000_000)
ETA = Fraction(1, 128)
ROUNDS = 4096
N_VALUES = (4, 8, 16, 32, 64, 128, 256)
E_VALUES = (1, 2, 4, 8, 16)
KAPPA_VALUES = (Fraction(1, 8), Fraction(1, 4), Fraction(1, 2), Fraction(1), Fraction(2))
def exact_subset_variance(n: int, sampled: int, kappa: Fraction) -> Fraction:
"""E[(sample mean z)^2], summing every subset exactly."""
half = n // 2
denominator = comb(n, sampled)
second_moment = Fraction(0)
for plus_count in range(max(0, sampled - half), min(half, sampled) + 1):
ways = comb(half, plus_count) * comb(half, sampled - plus_count)
signed_sum = 2 * plus_count - sampled
sample_mean = kappa * Fraction(signed_sum, sampled)
second_moment += Fraction(ways, denominator) * sample_mean * sample_mean
return second_moment
def stationary_gap(n: int, sampled: int, local_steps: int, kappa: Fraction) -> Fraction:
"""Exact stationary E[||grad L(theta)||^2] for the source construction."""
q = (1 - ETA) ** local_steps
local_gain = 1 - q
subset_variance = exact_subset_variance(n, sampled, kappa)
# M = q^2 M + (1-q)^2 Var(sample_mean), so M=(1-q)/(1+q)*Var.
return Fraction(local_gain, 1 + q) * subset_variance
def finite_round_gap(n: int, sampled: int, local_steps: int, kappa: Fraction) -> Fraction:
"""Exact expected gradient gap after ROUNDS rounds, starting at theta=0."""
q = (1 - ETA) ** local_steps
return stationary_gap(n, sampled, local_steps, kappa) * (1 - q ** (2 * ROUNDS))
def run() -> dict:
rows = []
ratios = []
finite_minimum = None
subset_checks = []
for n in N_VALUES:
# The source proof is a partial-participation construction. S<=N/2
# keeps the finite-population correction bounded away from zero.
for sampled in range(1, n // 2 + 1):
variance = exact_subset_variance(n, sampled, Fraction(1))
expected_variance = Fraction(n - sampled, sampled * (n - 1))
assert variance == expected_variance
subset_checks.append((n, sampled))
for local_steps in E_VALUES:
assert ETA <= Fraction(1, 8 * local_steps) # L=1 stability range
for kappa in KAPPA_VALUES:
gap = stationary_gap(n, sampled, local_steps, kappa)
finite_gap = finite_round_gap(n, sampled, local_steps, kappa)
target = Fraction(local_steps) * kappa * kappa / sampled
ratio = gap / target
finite_ratio = finite_gap / target
ratios.append(ratio)
finite_minimum = finite_ratio if finite_minimum is None else min(finite_minimum, finite_ratio)
rows.append(
{
"N": n,
"S": sampled,
"E": local_steps,
"kappa": str(kappa),
"subset_variance": str(exact_subset_variance(n, sampled, kappa)),
"stationary_gradient_gap": str(gap),
"E_kappa2_over_S": str(target),
"gap_over_E_kappa2_over_S": str(ratio),
}
)
minimum = min(ratios)
# eta=1/128 and E<=16 imply (1-(1-eta)^E)/(1+(1-eta)^E) >= 1/256.
# For S<=N/2, (N-S)/(N-1)>=1/2, and the exact grid minimum is stronger.
assert minimum >= Fraction(1, 512)
assert finite_minimum >= Fraction(1, 512)
assert len(rows) == sum(n // 2 for n in N_VALUES) * len(E_VALUES) * len(KAPPA_VALUES)
result = {
"construction": "two equal client groups, Li(theta)=1/2*(theta-z_i)^2, z_i=+-kappa",
"global_gradient": "grad L(theta)=theta",
"algorithm": "uniform-without-replacement partial FedAvg; exact E-step local GD",
"eta": str(ETA),
"rounds_from_zero": ROUNDS,
"L": 1,
"N_values": list(N_VALUES),
"S_range": "1..N/2 for each N",
"E_values": list(E_VALUES),
"kappa_values": [str(k) for k in KAPPA_VALUES],
"subset_count_cells": len(subset_checks),
"parameter_cells": len(rows),
"hypergeometric_identity_checked": True,
"minimum_exact_ratio": str(minimum),
"minimum_4096_round_ratio": str(finite_minimum),
"certificate": "stationary gradient gap >= (1/512)*E*kappa^2/S on every executed cell",
"rows": rows,
}
out = Path(__file__).with_name("theorem55_quadratic_scope_results.json")
out.write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps({k: v for k, v in result.items() if k != "rows"}, indent=2))
if __name__ == "__main__":
run()
|