AI.Cruz - Learn AI Through Interactive Visualizations

FEDERATED LEARNING

FL Challenges: Non-IID Data and Communication

Non-IID distribution, communication overhead, stragglers, and poisoning attacks — the four hard problems that turn federated learning from a neat idea into a real engineering challenge.

1Non-IID Data: The Core Challenge

The fundamental assumption of classical ML — that data is independently and identically distributed (IID) — is violated in FL. A phone in São Paulo has text in Portuguese; one in Tokyo has Japanese characters. A hospital specialising in paediatrics has different patient distributions than one specialising in oncology. When data is non-IID, local models diverge in different directions, and averaging them produces a model that fits no client well.

The severity is measured by a Dirichlet parameter α: α→∞ means IID (uniform labels across clients); α→0 means maximally non-IID (each client has only one label). Li et al. (2020) showed that FedAvg accuracy on CIFAR-10 drops from 74% (IID) to 51% (α=0.1) — a 23-point penalty purely from data heterogeneity.

✓ Mitigations

✓ FedProx: proximal term limits client drift from global model
✓ SCAFFOLD: variance reduction corrects for drift using control variates
✓ FedNova: normalise updates by local steps to remove bias
✓ Share a small public dataset globally to anchor distributions

✗ Why it's hard to fix

✗ You can't see or inspect client data to characterise the distribution
✗ Optimal hyperparameters depend on how non-IID the data is — unknown
✗ Personalised FL helps locally but hurts global generalisability

2Communication Cost

A ResNet-50 model has 25 million parameters — at float32, that's 100MB per upload per client. With 1,000 clients per round, a single round requires 100GB of uploads. Google's Gboard FL system runs hundreds of rounds — the communication bill would be astronomical without compression.

Gradient compression
Quantise float32 → int8 (4× smaller). Top-k sparsification: send only the largest k% gradients (reduce 10–100×). Structured updates: restrict updates to subspace.

Gradient accumulation
Skip rounds when the connection is metered or slow. Accumulate local updates and send a batch. Works because FL tolerates asynchronous participants.

Knowledge distillation
Rather than sharing weights, clients share predictions on a public dataset. The server distils a global model from these predictions — much smaller payloads.

3Stragglers

In cross-device FL, the server cannot wait indefinitely for slow clients. If the round waits for the slowest 1% of phones, it could take hours. If it proceeds without them, those clients' data is never used. This is the straggler problem: synchronous FL tolerates it badly; asynchronous FL handles it naturally but introduces a bias towards faster clients.

Deadline-based aggregation
Server aggregates all updates received within a timeout window. Stragglers' updates are dropped. Fair if sampling is random each round.

Asynchronous FL
Server updates the global model whenever any client finishes, weighted by staleness. Faster but can diverge if stale updates dominate.

Client selection
Pre-select clients likely to complete (charged, on WiFi, idle). Google's FL system uses these signals to maximise round completion rate.

4Poisoning Attacks

Because the server aggregates updates from untrusted clients, a malicious participant can corrupt the global model by sending carefully crafted poisoned updates. Untargeted poisoning degrades overall accuracy. Backdoor attacks embed a hidden trigger: the model behaves normally on clean data but misclassifies anything with the trigger pattern (e.g. a specific pixel pattern always predicts the attacker's target class).

✓ Defences

✓ Norm clipping: reject updates with unusually large norms
✓ Byzantine-robust aggregation: Krum, Bulyan, Trimmed Mean — filter statistical outliers
✓ Anomaly detection: flag clients whose updates are far from the median
✓ DP noise makes backdoors harder (noise dilutes the trigger)

✗ Why defences are imperfect

✗ Clever attacks mimic normal update distributions — evade norm clipping
✗ Byzantine-robust methods degrade accuracy on benign non-IID data
✗ In cross-device FL, you cannot verify client identity or trustworthiness

Simulate non-IID data distribution across FL clients

import random

def dirichlet_split(n_clients, n_classes, alpha, n_per_class=100, seed=42):
    """
    Distribute data across clients using Dirichlet(alpha).
    alpha -> inf: IID; alpha -> 0: each client has one class only.
    Returns: list of class-count dicts per client.
    """
    rng = random.Random(seed)
    client_counts = [{c: 0 for c in range(n_classes)} for _ in range(n_clients)]

    for cls in range(n_classes):
        # Sample Dirichlet proportions using gamma sampling
        gammas = [rng.gammavariate(alpha, 1) for _ in range(n_clients)]
        total = sum(gammas)
        proportions = [g / total for g in gammas]

        # Assign samples
        remaining = n_per_class
        for i, p in enumerate(proportions[:-1]):
            assigned = round(p * n_per_class)
            client_counts[i][cls] = assigned
            remaining -= assigned
        client_counts[-1][cls] = max(0, remaining)

    return client_counts

def print_distribution(counts, n_classes):
    print(f"{'Client':<10}", end="")
    for c in range(n_classes):
        print(f"{'cls'+str(c):<8}", end="")
    print("  total")
    print("-" * (10 + 8 * n_classes + 8))
    for i, cnt in enumerate(counts):
        print(f"{'C'+str(i+1):<10}", end="")
        for c in range(n_classes):
            print(f"{cnt[c]:<8}", end="")
        print(f"  {sum(cnt.values())}")

print("=== IID (alpha=100) ===")
counts_iid = dirichlet_split(4, 5, alpha=100)
print_distribution(counts_iid, 5)

print()
print("=== Moderate non-IID (alpha=1.0) ===")
counts_mod = dirichlet_split(4, 5, alpha=1.0)
print_distribution(counts_mod, 5)

print()
print("=== Extreme non-IID (alpha=0.1) ===")
counts_extreme = dirichlet_split(4, 5, alpha=0.1)
print_distribution(counts_extreme, 5)