Federated Learning
by Parul Pandey
Interactive simulation of federated learning with hospitals training local models and FedAvg aggregation.
About this notebook
This notebook simulates federated learning with four hospitals as clients, built with PyTorch and marimo. Each hospital holds 60 synthetic two-dimensional points, and a different linear rule labels each hospital's points, so the local datasets differ. A small feed-forward network classifies the points. matplotlib heatmaps show the decision surface of each hospital's model and of the shared global model.
Three buttons control the simulation. The Local Training button copies the global weights to every hospital and then trains each local model for five SGD steps on its own data. The four heatmaps diverge because each hospital learns a different pattern. The Merge Models button averages the four models' weights with FedAvg and makes the average the new global model. Every hospital receives a copy, the round counter increases, and an Altair line chart gains a point on the accuracy curve. The Reset button replaces all models with a fresh untrained network and clears the history.
marimo's state primitive stores the models, the round counter, and the accuracy history, so the simulation persists between button clicks. The FedAvg step computes a per-parameter mean of the four models' weights. The accuracy curve follows a fixed formula of the round number, not a measured test score, so it illustrates the trend rather than a real evaluation.
