Designing a Privacy-Preserving AI Recommendation Engine with Federated Learning, PyTorch, and FastAPI

Why privacy matters in modern recommendation systems

Imagine a streaming platform that suggests movies based on a user’s watch history, yet never stores that history on a central server. In 2023, more than 70% of data‑driven companies reported at least one breach involving personal data, highlighting the urgency of privacy‑preserving techniques. A recommendation engine that respects user confidentiality can reduce legal risk, improve brand trust, and still deliver personalized experiences.

Federated learning fundamentals for recommendation tasks

Federated learning (FL) moves model training to the edge devices that hold the raw interaction logs. Instead of uploading clickstreams, each client computes gradients locally and sends only encrypted weight updates to a coordinating server. The server aggregates updates—often with a weighted average known as FedAvg—and produces a global model that improves for all participants without ever seeing raw data.

Key parameters that influence convergence include the number of local epochs (E), the client‑side batch size (B), and the participation rate (C). Research from Google in 2021 shows that with E=5, B=32, and C=0.1, a language model converged to 95% of the centralized baseline after 200 communication rounds.

Building the recommendation model with PyTorch

For a typical collaborative‑filtering engine, a simple matrix factorization model works well. PyTorch offers automatic differentiation and GPU support, making it ideal for both local client training and server‑side aggregation.

import torch
import torch.nn as nn

class MFModel(nn.Module):
    def __init__(self, num_users, num_items, embed_dim=64):
        super(MFModel, self).__init__()
        self.user_emb = nn.Embedding(num_users, embed_dim)
        self.item_emb = nn.Embedding(num_items, embed_dim)
        self.bias_user = nn.Embedding(num_users, 1)
        self.bias_item = nn.Embedding(num_items, 1)
        self.global_bias = nn.Parameter(torch.zeros(1))

    def forward(self, user_ids, item_ids):
        u = self.user_emb(user_ids)
        i = self.item_emb(item_ids)
        bu = self.bias_user(user_ids).squeeze()
        bi = self.bias_item(item_ids).squeeze()
        pred = (u * i).sum(dim=1) + bu + bi + self.global_bias
        return pred

The model can be instantiated on each client with local user/item vocabularies. When a new client joins, the server sends the latest global weights, and the client fine‑tunes for a few epochs before sending back the delta.

Exposing the federated service with FastAPI

FastAPI provides asynchronous endpoints that can handle model version queries, client registration, and weight aggregation. Its built‑in data validation (Pydantic) ensures that malformed updates are rejected early.

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
import json

app = FastAPI()

global_model = MFModel(num_users=100000, num_items=50000)

class UpdatePayload(BaseModel):
    client_id: str
    weights: dict  # serialized state_dict

@app.post("/upload_update")
async def upload_update(payload: UpdatePayload):
    try:
        client_state = {k: torch.tensor(v) for k, v in payload.weights.items()}
        # Simple FedAvg placeholder – in production use secure aggregation
        for name, param in global_model.state_dict().items():
            param.data = (param.data + client_state[name]) / 2
        return {"status": "ok"}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.get("/model")
async def get_model():
    return {"weights": {k: v.tolist() for k, v in global_model.state_dict().items()}}

The endpoint returns a JSON‑serializable version of the model so that lightweight devices can download it without additional libraries.

Orchestrating federated updates in practice

Typical FL pipelines involve three scripts: (1) a client trainer, (2) a server aggregator, and (3) a scheduler. Docker Compose can spin up the FastAPI service, a Redis message broker, and a simple Celery worker that triggers aggregation every hour.

docker network create fl-net

docker run -d --name redis --network fl-net redis:6-alpine

docker run -d --name api \
    -p 8000:80 \
    -e REDIS_HOST=redis \
    --network fl-net \
    myorg/federated-api:latest

docker run -d --name worker \
    -e REDIS_HOST=redis \
    --network fl-net \
    myorg/federated-worker:latest

Clients pull the latest model via GET /model, train locally for E=3 epochs on their interaction logs, then POST the weight delta to /upload_update. The server aggregates using a weighted average based on the number of samples each client processed.

Practical tips for production deployment

1. **Secure aggregation** – Use homomorphic encryption or secret‑sharing schemes to ensure the server never sees individual updates in clear text. 2. **Differential privacy** – Clip gradients and add Gaussian noise (σ≈1.0) before transmission; this limits the influence of any single user. 3. **Client heterogeneity** – Adjust learning rates per device; slower phones may need a reduced step size to avoid divergence. 4. **Monitoring** – Track validation loss on a hold‑out set at the server after each round; sudden spikes often indicate corrupted updates.

By combining these safeguards, a recommendation engine can stay both accurate and privacy‑aware, meeting GDPR or CCPA requirements without sacrificing user experience.

Sources

Google AI Blog – Federated Learning: Communication‑Efficient Strategies
PyTorch Official Documentation – torch.nn.Embedding
FastAPI Documentation – Building APIs with async support

Author: Mahmut Sarıkaya — sarikayadev.com