AIDA Technical Day
Federated Learning
1
2025-09-16, Erik Ylipää (erik.ylipaa@liu.se)
AIDA Technical Days
September 16 - Federated Learning
13:00 - 13:50 - Introduction to Federated Learning
13:50 - 14:00 - Break
14:00 - 14:50 - Federated Learning Example
14:50 - 15:00 - Break
15:00 - 16:00 - Federated Learning Challenges
2
AIDA & AIDA Data Hub
AIDA Community - medtech4health.se/aida
National collaboration arena for AI research and innovation in medical imaging diagnostics.
AIDA Data Hub - datahub.aida.scilifelab.se
E-infrastructure set up to support AIDA.
Vetenskapsrådet
Research funding agency
Government
Executive branch
Knut and Alice Wallenberg foundation
Private research funder
VINNOVA
Innovation agency
SciLifeLab
Life science research infrastructure/center
AIDA
Collaboration arena in Swedish medical imaging diagnostics AI innovation
AIDA Data Hub
e-infrastructure supporting AIDA
NBIS
Bioinformatics platform
AIDA mission
Bridge the gap between promising research results and actual patient benefit, through a clinic-native research agenda for innovation.
Clinical wilderness
Research sandbox
AIDA Community
Publicly funded collaboration arena for AI innovation in medical imaging diagnostics.
Healthcare
Academia
Industry
AIDA Data Hub
E-infrastructure for clinical innovation in data driven precision health.
Data services
Secure long term primary storage and compute
Support
AIDA Data Hub Team
Caroline Bivik Stadler
AIDA Arena Director
Varshith Konda
Systems development
Pontus Freyhult
IT Architect
Betul Eren
Data sharing lead
Erik Ylipää
Support lead
Emre Balsever
Systems development
Claes Lundström
AIDA Scientific Director
Joel Hedlund
AIDA Data Hub Lead
Federated Learning
Why? What? How?
9
The tenets of modern AI
10
Rich Sutton, The Bitter Lesson (2019) http://www.incompleteideas.net/IncIdeas/BitterLesson.html
“One thing that should be learned from the bitter lesson is the great power of general purpose methods, of methods that continue to scale with increased computation even as the available computation becomes very great. The two methods that seem to scale arbitrarily in this way are search and learning.”
Kaplan, Jared, et al. "Scaling laws for neural language models." arXiv preprint arXiv:2001.08361 (2020).
Statistical Learning Primer
11
World
Model
Data bottlenecks contemporary AI
12
Distributed Machine Learning
13
Large dataset
Machine 1
Machine 3
Machine 2
Partition dataset
Parameter server
Local parameter updates
Aggregated (global) parameters updates
Clinical data
While the legality of vacuuming the internet of data to train AI might be unclear, for health data it’s not.
Gathering all the world's health data to a supercomputer will not happen.
Alternative:
14
Share models instead of data
While the data itself might not be sharable, models derived from the data could be
A federation can be formed around this setup, and training models together becomes federated learning
15
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Share model with federation
Federation
Typically Centralized
Model is typically shared through a common server
Clients (organizations) share the local model they have trained on their private data with central server
Server aggregates local models into global model
16
Organization 1
Organization 3
Organization 2
Aggregator
Local parameter updates
Aggregated (global) parameters updates
Private data 1
Private data 2
Private data 3
Training is [often] synchronized
17
Aggregator
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Aggregator
Starts with global model, newly initialized or result of previous round
Train on private data
Train on private data
Train on private data
Share global model with federation
Share local models with aggregator
Aggregate local models to new global model
Adapt global model on private data into local model
One round of federated learning
Federated algorithms (how to aggregate)
There are many ways one can aggregate the models:
Federated Average (FedAvg) is a commonly used baseline: global model is the average of the local models
Easily extends into weighted variants (e.g. set weight inversely proportional to dataset size)
Other extensions involves server-side momentum-based methods (e.g. server has exponential moving average of aggregated gradients).
18
Aggregator
Aggregate local models to new global model
How long is a round?
19
Aggregator
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Aggregator
Starts with global model, newly initialized or result of previous round
Train on private data
Train on private data
Train on private data
Share global model with federation
Share local models with aggregator
Aggregate local models to new global model
Adapt global model on private data into local model
One round of federated learning
How long do the clients train on their local data?
Data non-IID
It is highly likely that data is not identically distributed across sites.
Local models are likely to diverge from each other, they learn different patterns from their respective datasets
The average of these diverged local models might perform poorly
The longer we train local models before synchronization, the larger the potential divergence
20
Zhu, Hangyu, et al. "Federated learning on non-IID data: A survey." Neurocomputing 465 (2021): 371-390.
An analogy: The average Gaussian fit of two independent Gaussian is a poor representation of either of them
Synchronize often?
Issues with divergence indicate we should synchronize often!
But, each synchronization requires us to upload/download the models to/from the server
If we synchronize too often, we mostly spend time waiting on transfers!
21
Zhu, Hangyu, et al. "Federated learning on non-IID data: A survey." Neurocomputing 465 (2021): 371-390.
Aggregator
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Aggregator
Starts with global model, newly initialized or result of previous round
Train on private data
Train on private data
Train on private data
Share global model with federation
Share local models with aggregator
Aggregate local models to new global model
Adapt global model on private data into local model
One round of federated learning
Communication has a cost!
How long is a round?
22
Aggregator
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Aggregator
Starts with global model, newly initialized or result of previous round
Train on private data
Train on private data
Train on private data
Share global model with federation
Share local models with aggregator
Aggregate local models to new global model
Adapt global model on private data into local model
One round of federated learning
How long do the clients train on their local data?
Run experiments!
Evaluation - client-side
23
Aggregator
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Aggregator
Starts with global model, newly initialized or result of previous round
Train on private data
Train on private data
Train on private data
Share global model with federation
Share local models with aggregator
Aggregate local models to new global model
Adapt global model on private data into local model
One round of federated learning
Aggregator
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Evaluation private dev data
Evaluation on private dev data
Evaluation on private dev data
Share global model with federation
Aggregate local performance metrics to global
Evaluate global model on private development data
AUCROC2
AUCROC3
AUCROC1
AUCROCglobal
Frameworks - ordered after github activity
24
What framework to choose
Pick one of:
25
NVIDIA Flare
26
NVIDIA Flare - Basic operation
27
Python program
How to determine what software to run?
28
Python program
What to run – implications
29
Python program
Execution setup
30
Flare server
Flare client 1
Private data 1
Flare client 2
Private data 2
Flare client 3
Private data 3
Prepares machine by setting up software and preparing data
Asks machine to run pre-installed program - gives out weights
Asks machine to run pre-installed program - gives out weights
Asks machine to run pre-installed program - gives out weights
Program reports updated weights back
Program reports updated weights back
Program reports updated weights back
Prepares machine by setting up software and preparing data
Prepares machine by setting up software and preparing data
Org 1
Org 2
Org 3
Aggregates weights
Questions?
31
Break
32
Federated Learning Example
33
Peter Lundberg, Erik Ylipää, Oskar Jerdhaf, Arne Jönsson, Emma Eneling
Community Council Östergötland, Linköping University, Center for Medical Imaging and Visualization
Federated health
34
Magnetic Resonance Imaging
35
Fantastic method of looking inside the body
Extremely strong magnet
Medical Implants
36
Deep Brain Stimulation (DBS):
‘pacemaker for the brain’
[P. Lundberg]
Medical implants and MRI can be a harmful combination
37
[P. Lundberg]
The Magnet is always on
MRI Safety concerns
38
[Lundberg P et al, Eur Radiol (2024)]
MRI Safety concerns
39
[Lundberg P et al, Eur Radiol (2024)]
The current clinical workflow
40
Extend the current flow – find implant terms
41
Available data
Clinical notes
Implant term list
42
Keyword search challenges
43
We only find the terms we know exactly, what about terms which means the same thing?
Automatically extract information from “natural language”
44
Natural language
Formal language
Text
summarization
Document classification
Automatic translation
Named Entity Recognition
What makes natural language hard?
45
“I shot an elephant in my pajamas”
No, he doesn’t have enough money to buy the car
There’s no car he doesn’t have enough money to buy
The car is beyond his budget
Superficial similarity does not mean similar meaning
Even complete superficial similarity can have different meanings
Two main approaches: Statistical and rule based
46
Rule based (grammars, structure)
Statistical (conditional probabilities, data, search)
David Bamman, Info 159/259, Lecture 12 - Syntax, https://people.ischool.berkeley.edu/~dbamman/nlp22.html
Statistical Learning Primer
47
World
Model
Our approach - semantic vector search
48
Word vectors - each word is encoded as a vector in a space, where different directions in space encode different semantics
Similar objects represented as vectors can be found by looking at how close they are in the semantic space
Semantic search with (BERT-style) LLM
49
XLM-Roberta
Langue modelling fine tuning
XLM-RoBERTa is a multilingual model trained on 2.5TB or text from 100 languages.
50
Experiment set-ups
2 main branches:
51
Flare server
Flare client 1
Private data 1
Flare client 2
Private data 2
Flare client 3
Private data 3
Prepares machine by setting up software and preparing data
Asks machine to run pre-installed program - gives out weights
Asks machine to run pre-installed program - gives out weights
Asks machine to run pre-installed program - gives out weights
Program reports updated weights back
Program reports updated weights back
Program reports updated weights back
Prepares machine by setting up software and preparing data
Prepares machine by setting up software and preparing data
Org 1
Org 2
Org 3
Aggregates weights
Federated (global) model
Flare client 1
Private data 1
Org 1
Flare client 2
Private data 2
Org 2
Flare client 3
Private data 3
Org 3
Local model org 1
Local model org 2
Local model org 3
Example evaluation (implant detection)
52
52
XLM-Roberta
XLM-Roberta
Federated (global) model
Local model org 1
Compare how well implants are found between global and local model
Status
Technical foundation in place, 4 sites has run a federated training on toy data (public domain texts).
4 Sites has performed local training: done masked language modelling on their own clinical datasets.
Technically, we just need to change what texts the clients point to.
53
Flare server
Flare client 1
Private data 1
Flare client 2
Private data 2
Flare client 3
Private data 3
Prepares machine by setting up software and preparing data
Asks machine to run pre-installed program - gives out weights
Asks machine to run pre-installed program - gives out weights
Asks machine to run pre-installed program - gives out weights
Program reports updated weights back
Program reports updated weights back
Program reports updated weights back
Prepares machine by setting up software and preparing data
Prepares machine by setting up software and preparing data
Org 1
Org 2
Org 3
Aggregates weights
Break
54
Federated Learning Challenges
Distributed machine learning is easy – Federated Learning is hard
55
Is it safe to share the model?
Some attack vectors for shared models:
56
Carlini, Nicholas, et al. "Extracting Training Data from Large Language Models." USENIX Security Symposium. Vol. 6. 2021.
Add noise to updates
We can disturb the model updates, by clipping gradients and adding noise
This makes inversion and inference attacks harder, but trades off performance. Differential Privacy gives the theoretical foundation for this and can guide you in how much noise is needed to bound the probability of revealing a data point
57
Aggregator
Organization 1
Private data 1
Organization 2
Private data 2
Organization 3
Private data 3
Aggregator
Starts with global model, newly initialized or result of previous round
Train on private data
Train on private data
Train on private data
Share global model with federation
Share local models with aggregator
Aggregate local models to new global model
Adapt global model on private data into local model
One round of federated learning
Add noise to model updates
Add noise to model updates
Add noise to model updates
Limit model capacity
Another mitigation is to limit the number of parameters the model has to learn the private data.
If the model is fine-tuned from a pre-trained foundation model (as is often the case), we can limit the number of parameters we train. Low Rank Adaptation (LoRA) is a common technique for this.
58
Other LoRA benefits
Apart from making inference and inversion attacks more difficult, LoRA also dramatically reduce the number of parameters (e.g. from 125M to 200k), drastically reducing communication overhead.
It also tends to have a regularizing effect (much fewer parameters to overfit)
59
Secure aggregation
60
Federated Machine Learning - just a special case of distributed machine learning?
61
61
Machine 1
Machine 3
Machine 2
Parameter server
Local parameter updates
Aggregated (global) parameters updates
Private data 1
Private data 2
Private data 3
Large dataset
Machine 1
Machine 3
Machine 2
Partition dataset
Parameter server
Local parameter updates
Aggregated (global) parameters updates
Distributed Machine Learning - Control
62
Large dataset
Machine 1
Machine 3
Machine 2
Partition dataset
Parameter server
Local parameter updates
Aggregated (global) parameters updates
Single organization
Physical
Software
Communication
Data
System
Access
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
What machines are used, where are they located
How are the machines configured
What software is being run. What model is trained, how is the training defined, what hyper parameters are used.
What fabric is used for communication. How do machines communicate with each other. What machines can communicate with each other.
Who can access the machine and data. Who can configure the system.
What does the data contain. How much data is there. What data will be used for what (e.g. train vs. test)
Federated Machine Learning - Coordination
63
Physical
Software
Communication
Data
System
Access
Machine 1
Machine 3
Machine 2
Parameter server
Local parameter updates
Aggregated (global) parameters updates
Private data 1
Private data 2
Private data 3
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 1
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 2
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 3
Controlled by each participant - minimal coordination needed
Controlled by each participant - some coordination needed
Controlled by the federation - much coordination needed
Controlled by the federation - much coordination needed
Controlled by the participant - some coordination needed
Controlled by each participant - some coordination needed
What machines are used, where are they located
How are the machines configured
What software is being run. What model is trained, how is the training defined, what hyper parameters are used.
What fabric is used for communication. How do machines communicate with each other. What machines can communicate with each other.
Who can access the machine and data. Who can configure the system.
What does the data contain. How much data is there. What data will be used for what (e.g. train vs. test)
NVIDIA Flare (and other frameworks)
64
Physical
Software
Communication
Data
System
Access
Controlled by each participant - minimal coordination needed
Controlled by each participant - some coordination needed
Controlled by the federation - much coordination needed
Controlled by the federation - much coordination needed
Controlled by the participant - some coordination needed
Controlled by each participant - some coordination needed
What machines are used, where are they located
How are the machines configured
What software is being run. What model is trained, how is the training defined, what hyper parameters are used.
What fabric is used for communication. How do machines communicate with each other. What machines can communicate with each other.
Who can access the machine and data. Who can configure the system.
What does the data contain. How much data is there. What data will be used for what (e.g. train vs. test)
Technologists are used to automate everything
In a single organization setup, the ML researcher can do most things at will; data exploration, data pre-processing, setting up experiments and evaluating experiments
Automation across a federation is much harder, most steps need coordination!
65
Large dataset
Machine 1
Machine 3
Machine 2
Partition dataset
Parameter server
Local parameter updates
Aggregated (global) parameters updates
Repeat: What to run – implications
66
Python program
What can you automate centrally?
67
What can you automate centrally?
68
You don’t know how the data is formatted, what idiosyncrasies it has and what needs to cleaned up.
What can you automate centrally?
69
Even if you knew how to filter, you can’t be allowed to (centrally limiting to small subsets would be a privacy risk).
What can you automate centrally?
70
Similar to filtering, but additionally you lack domain knowledge about the specific sites biases to make good selection for a test set.
What can you automate centrally?
71
Each site must be in control of what software will run. If centrally prepared, the sites still need to thoroughly inspect the environment. Probably better to inspect the build files (dockerfile, python requirements) and build your own.
What can you automate centrally?
72
Most setups disallow the central server to establish connections into the clients (only allowing the reverse). Some separate communication channel for initiating connections could be established with additional complexity. Often requires manual intervention (recover from crashes, network errors, power outage)
What can you automate centrally?
73
No
What can you automate centrally?
74
Probably ok, but needs to be pre-approved
What happens in case of drop-out?
Depends on task
In cross-silo training, losing a client is often cause to stop that training
Often
75
Federated learning is personnel intense
Many competences need to be replicated in each organisation.
76
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 2
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 1
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 3
Software
Communication
Data
System
Access
Federated learning is more about Project Management than distributed machine learning
77
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 2
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 1
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 3
Software
Communication
Data
System
Access
Questions
78
Thank you!
79
CCÖ Local+Federated training: Experiences
80
Software environments
Since we can’t just “push” software installation to clients, everything needs to be manually replicated at each site.
Differences in systems at sites makes software management more of a challenge.
81
Data preprocessing
Data has to be handles out-of-sight, difficult to look at concrete data examples from different sites.
For MLM, relatively easy, we just need text. For labeled tasks or multimodal learning this would be much more of a challenge.
82
Secure compute environment
Difficult to work with the local environment. Firewalls (a necessity) makes installing software more difficult.
Changes made to local software can’t be synchronized back to external repo.
Limited remote desktop, software can’t be directly uploaded.
83
Hardware limitations
Secure compute environments might have limited hardware (procured for the project), limitations in storage and compute compared to High Performance Compute (HPC) environments.
84
Coordination
Synchronizing between busy people is hard!
85
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 2
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 1
Principal investigator
Security expert
Machine learning expert
Domain expert
Network expert
Policy expert
System administrator
Data controller
Organization 3
Software
Communication
Data
System
Access