Federated Privacy-Preserving Typing Assistant

The challenge

The objective was to obtain useful next-word predictions without ever centralising raw keystroke data, and to prove the whole loop works end to end: local training on each client, differential privacy applied to every shared update, federated averaging into a global model, and live predictions served from that global model. Two training paths had to interoperate on the same weights , Python Flower clients consuming streamed data, and a fully in-browser trainer.

Approach

• Built a 2-layer Transformer encoder (~943K parameters) predicting the next word from a 31-word context over a 2,939-word vocabulary (Tiny Shakespeare), kept Opacus-compatible so identical weights load into either training path.
• Ran federated averaging with Flower over 15 rounds and 3 clients, each holding a distinct shard of the corpus and performing ~330 local steps per round, with aggregated weights published to Redis for any consumer.
• Applied per-client differential privacy with Opacus DP-SGD , gradient clipping plus calibrated noise , tracked by a formal privacy accountant so every round reports an explicit epsilon.
• Streamed keystroke events through Kafka into a native Scala Flink job that windows them into per-client word-sequence batches, with Redis persistence and OpenTelemetry/Grafana observability across the loop.
• Implemented an edge variant that trains entirely in the browser on TensorFlow.js's WASM backend (forward, backward, optimizer step, update-level clip and noise) and submits to a FastAPI gateway that only advances a round once every distinct client has contributed.
• Provisioned the full stack with Terraform on AWS , VPC, Kafka/Flink/observability hosts, FL server and clients , with admin access ranges and credentials required as explicit variables.

Outcome

After fixing the privacy-accountant calibration , it had been sized against the wrong data volume, which overstated the budget , and increasing the DP noise level to compensate, the privacy guarantee is now real and defensible: cumulative ε reaches approximately 1.12 after 15 rounds, a strong differential-privacy guarantee by industry standards. Training loss falls from 7.80 to 6.88 cross-entropy over 15 rounds, and the model reaches 7.0% top-1 and 19.0% top-5 next-word accuracy (perplexity 993). Top-1 is now essentially tied with the trivial "always predict the most frequent word" baseline (7.2%), while top-5 remains meaningfully above it. This is the expected privacy/utility trade-off: a meaningful DP guarantee costs some model quality, here slightly reducing top-5 accuracy while leaving top-1 roughly unchanged. The DP budget is no longer the bottleneck; closing the accuracy gap for a production typing assistant would require more training data/rounds or a larger model.