Implementing SecAgg Masking for Spatial Gradients

Secure Aggregation (SecAgg) lets a federated server learn the sum of many clients’ spatial gradients while learning nothing about any individual update, and the mechanism that makes that possible is pairwise masking: every ordered pair of clients agrees on a shared random mask that one client adds and the other subtracts, so the masks vanish in the server’s sum and leave only the aggregate. This page is the implementation deep-dive that sits beneath the secure aggregation protocols guide — where the protocol overview, the multi-round handshake, and the graph-of-clients topology are covered end to end — inside the wider Federated Learning Workflows for Geospatial Data architecture. Here the scope is deliberately narrow: the concrete construction of one masking round over a quantized spatial gradient vector, from Diffie–Hellman seed agreement through PRG expansion to modular addition in $\mathbb{Z}_R$. Because a coordinate-regression or heatmap model trained across regions produces gradients whose per-client distribution is skewed by geography — the same skew that motivates handling non-IID geospatial data in federated learning — the masking layer must be exact to the bit, and this page shows how to get it there.

Parameter Configuration and Calibration

SecAgg masking has four knobs, and each one is a correctness or privacy constraint rather than a performance dial. The masks live in a finite ring $\mathbb{Z}_R$; the gradients must be quantized into that same ring before any mask is applied. Get the ring size wrong and the true sum wraps around silently; get the clip bound wrong and one client’s gradient dominates the aggregate.

The ring constraint is the one that bites in production. Each quantized element has magnitude at most $S \cdot C$; summing $n$ clients gives at most $n\,S\,C$, and because we encode signed values as residues we need headroom on both sides, so $R > 2\,n\,S\,C$. With $n = 500$ mobile mapping nodes, $S = 10^3$, and $C = 1$, the sum reaches $5\times10^5$ — three orders of magnitude below $2^{32}$, which is why $R = 2^{32}$ is a safe default for gradient vectors up to a few thousand clients. Push $n$ into the tens of thousands or raise $S$ for fine-grained coordinate regression and you move to $R = 2^{64}$.

The pairwise seed itself comes from a Diffie–Hellman key agreement: client $i$ holds a secret $sk_i$ and publishes $pk_i = g^{sk_i} \bmod p$; the shared value for pair $(i,j)$ is $g^{sk_i sk_j} \bmod p$, which both parties compute independently and no one else can. That symmetry is the whole trick — it lets two clients derive the identical mask without ever exchanging it, so the additive inverse cancels perfectly. This is the same additive-inverse principle that underpins additive secret sharing for coordinates in Python, applied here to gradient vectors instead of raw coordinates.

Client 2 quantize g₂ ∈ Z_R y₂ = g₂ −m₁₂ +m₂₃ (+self-mask)

Client 3 quantize g₃ ∈ Z_R y₃ = g₃ −m₁₃ −m₂₃ (+self-mask)

Reference Implementation

The three functions below implement one masking round with pure standard library plus NumPy for the vector arithmetic. derive_pairwise_seed runs the Diffie–Hellman agreement, prg_mask expands a seed into a mask vector over $\mathbb{Z}_R$, and mask_update assembles the update a single client uploads. The self-mask is added on top of the pairwise masks: its seed is Shamir-shared with the other clients so that if this client drops out mid-round, a surviving quorum can reconstruct the seed and subtract the self-mask — the dropout-recovery path that keeps the round from stalling.

from __future__ import annotations

import hashlib
import secrets
from typing import Dict

import numpy as np

# --- Ring and quantization parameters (see Parameter Configuration) ---
R: int = 1 << 32                    # gradient ring Z_R; need R > 2 * n * S * C
DH_PRIME: int = (1 << 256) - 189   # DEMO prime only. A real round MUST use a
DH_GEN: int = 5                     # 2048-bit MODP group AND authenticated DH.


def derive_pairwise_seed(sk_i: int, pk_j: int, prime: int = DH_PRIME) -> int:
    """Diffie-Hellman shared secret between clients i and j -> a 256-bit PRG seed.

    Both clients compute the SAME group element: pk_j**sk_i == g**(sk_i*sk_j)
    == pk_i**sk_j, so neither has to transmit the mask.
    SECURITY: this raw DH is UNAUTHENTICATED. In production the public keys must be
    signed and exchanged over authenticated channels, or a man-in-the-middle server
    can learn every pairwise seed and strip the masks off each update.
    """
    shared: int = pow(pk_j, sk_i, prime)                    # g^{sk_i * sk_j} mod p
    # Hash the group element to a uniform seed; the raw DH value is not uniform.
    return int.from_bytes(hashlib.sha256(str(shared).encode()).digest(), "big")


def prg_mask(seed: int, length: int, modulus: int = R) -> np.ndarray:
    """Expand a seed into a deterministic mask vector in Z_modulus^length.

    SHA-256 in counter mode is the pseudo-random generator: identical seed -> identical
    mask on both clients, which is exactly what makes +mask and -mask cancel. A plain
    non-crypto RNG here would let the server predict masks and unmask a lone update.
    """
    seed_bytes: bytes = seed.to_bytes(32, "big")            # 256-bit seed
    out = np.empty(length, dtype=np.int64)
    filled, counter = 0, 0
    while filled < length:
        block = hashlib.sha256(seed_bytes + counter.to_bytes(8, "big")).digest()
        for k in range(0, 32, 4):                           # 8 x 4-byte words / block
            if filled >= length:
                break
            out[filled] = int.from_bytes(block[k:k + 4], "big") % modulus
            filled += 1
        counter += 1
    return out


def quantize(grad: np.ndarray, clip: float, scale: int) -> np.ndarray:
    """Fixed-point encode a real gradient into non-negative residues of Z_R.

    Clipping to [-clip, clip] bounds each client's contribution so one skewed
    regional gradient cannot dominate the masked sum, then scaling + mod R lands
    the value in the ring where masks live.
    """
    clipped = np.clip(grad, -clip, clip)
    q = np.round(clipped * scale).astype(np.int64)
    return np.mod(q, R)                                     # signed value -> residue


def mask_update(
    grad_q: np.ndarray,
    my_id: int,
    my_sk: int,
    peer_pks: Dict[int, int],
    self_seed: int,
) -> np.ndarray:
    """Build the masked vector client `my_id` uploads for one SecAgg round.

    For each peer j: derive s_ij by DH, expand to a mask, and add +mask if my_id < j
    else -mask. Because client j does the mirror operation, the two masks are additive
    inverses mod R and cancel in the server's sum. The self-mask on top is removed only
    if a quorum reveals this client's secret-shared self_seed (dropout recovery).
    """
    length = int(grad_q.shape[0])
    masked = grad_q.astype(np.int64).copy()
    for j, pk_j in peer_pks.items():
        if j == my_id:
            continue
        m_ij = prg_mask(derive_pairwise_seed(my_sk, pk_j), length)
        if my_id < j:
            masked = np.mod(masked + m_ij, R)              # lower id adds +m
        else:
            masked = np.mod(masked - m_ij, R)              # higher id adds -m
    masked = np.mod(masked + prg_mask(self_seed, length), R)   # per-client self-mask
    return masked

Validation Checkpoint

The invariant that must hold is exact: after the server sums every masked upload and removes the self-masks of surviving clients, the result must equal the plain modular sum of the quantized gradients — bit for bit, not approximately. A single mismatched mask leaves a large pseudo-random residue in the aggregate, so this check fails loudly rather than degrading quietly. The harness builds four clients, masks their quantized gradients, and asserts the cancellation.

def _validate() -> None:
    rng = np.random.default_rng(7)
    n, dim = 4, 16
    clip, scale = 1.0, 1_000

    # DH keypairs (UNAUTHENTICATED demo) and secret-shared self-seeds per client.
    sk = {i: secrets.randbelow(DH_PRIME - 2) + 1 for i in range(1, n + 1)}
    pk = {i: pow(DH_GEN, sk[i], DH_PRIME) for i in range(1, n + 1)}
    self_seed = {i: int.from_bytes(secrets.token_bytes(32), "big")
                 for i in range(1, n + 1)}

    # Each client quantizes its (geographically skewed) gradient.
    grad_q = {i: quantize(rng.normal(0.0, 0.3, dim), clip, scale)
              for i in range(1, n + 1)}

    # Each client masks and uploads.
    uploads = {}
    for i in range(1, n + 1):
        peers = {j: pk[j] for j in range(1, n + 1) if j != i}
        uploads[i] = mask_update(grad_q[i], i, sk[i], peers, self_seed[i])

    # Server: sum masked uploads, then subtract self-masks of surviving clients
    # (their self_seeds are recovered from Shamir shares; no dropouts here).
    agg = np.zeros(dim, dtype=np.int64)
    for i in range(1, n + 1):
        agg = np.mod(agg + uploads[i], R)
    for i in range(1, n + 1):
        agg = np.mod(agg - prg_mask(self_seed[i], dim), R)

    # Ground truth: plaintext modular sum of the quantized gradients.
    truth = np.zeros(dim, dtype=np.int64)
    for i in range(1, n + 1):
        truth = np.mod(truth + grad_q[i], R)

    assert np.array_equal(agg, truth), "pairwise masks failed to cancel"
    # Sanity: an individual masked upload must NOT equal its plaintext gradient.
    assert not np.array_equal(uploads[1], grad_q[1]), "update was left unmasked"
    print("SecAgg masking invariant holds: unmasked sum == plaintext gradient sum")


if __name__ == "__main__":
    _validate()

The second assertion matters as much as the first: it confirms the upload the server actually receives is not the raw gradient. If both hold, the masks are strong on the wire yet perfectly transparent in aggregate.

Incident Response and Edge Cases

SecAgg masking fails in ways that either stall a round or, worse, corrupt the aggregate without any error. The scenarios below are the ones that surface on live federated rounds over spatial data.

• Client dropout mid-round. A mobile mapping node uploads its masked vector, then goes offline before the server can collect the reveal messages. Its pairwise masks with every surviving peer are now stuck in the sum with no matching inverse. Remediation: the surviving clients reveal their Shamir shares of the dropped client’s DH secret so the server reconstructs the missing pairwise masks and subtracts them; for clients that survived, they instead reveal the self-mask seed. The protocol must never let the server obtain both the self-seed and the pairwise seeds of the same client, or it can unmask that client’s update.
• Quantization overflow past $R$. With too many clients or too large a scale $S$, the true sum exceeds $R$ and wraps, so the recovered aggregate is off by a multiple of $R$ on some coordinates — a silent, catastrophic gradient. Detection: assert $R > 2\,n\,S\,C$ at round setup and log the observed max element. Remediation: raise $R$ to $2^{64}$ or lower $S$, and clip more aggressively so the per-client magnitude bound $C$ actually holds.
• Mask reuse across rounds. If a client reuses the same DH keypair and self-seed for a new round, an attacker who saw last round’s reveal can subtract the stale masks from this round’s upload and recover the fresh gradient. Remediation: derive per-round seeds by hashing the DH element together with a monotonic round counter, rotate DH keys each round, and reject any upload carrying a stale round tag — the same round-nonce discipline used across async execution patterns.
• Non-canonical vector length. Two clients quantize gradients of different lengths (a model-version skew), so their pairwise masks are expanded to different lengths and no longer cancel. Remediation: pin the model architecture hash into the round setup and refuse to admit a client whose gradient dimension does not match the round’s declared dim.

Frequently Asked Questions

Related

• Secure Aggregation Protocols — the parent guide covering the full SecAgg handshake, client graph, and multi-round protocol this page drills into.
• Handling Non-IID Geospatial Data in Federated Learning — why regional gradient skew makes exact, unbiased aggregation essential.
• Additive Secret Sharing for Coordinates in Python — the additive-inverse cancellation principle applied to raw coordinates instead of gradients.
• Async Execution Patterns — the round-nonce and dropout handling that keeps masking rounds from stalling under churn.

Up: Secure Aggregation Protocols · Federated Learning Workflows for Geospatial Data