Writing

Federated learning: the model was the easy part

In short: when sites can't share their data, the model code is the smallest part of the work. Most of the effort goes into making the data mean the same thing everywhere, coping with sites that drop…

published
read time
5 min
words
918
lang
en
filed under
Engineering

In short: when sites can't share their data, the model code is the smallest part of the work. Most of the effort goes into making the data mean the same thing everywhere, coping with sites that drop out, and coordinating people across time zones.

For a while I was a machine learning engineer on a research consortium across several institutions. The question was big and slow, the kind that follows people for years. The data was sensitive in every way data can be. Health records, mental health measures, family circumstances. Nobody was going to put it in a shared bucket, and nobody should.

So the plan was federated learning. Each site keeps its data. The model travels to the data, trains locally, and only the learned parameters come back to be combined. On a slide it takes three boxes and an arrow. In practice the model was the part I worried about least.

What federated learning actually moves

The idea is simple enough to draw. Every site runs the same training code on its own rows. It sends back a model update, not the data. A central aggregator combines the updates, usually a weighted average, and sends the new model out again. Repeat for a number of rounds.

site A site B site C, offline aggregator weights only raw rows never leave a site
Only model updates travel. The hard part is the box that goes quiet.

That part worked. Standard models, standard averaging. What did not work out of the box was everything around it.

The data didn't match

The cohorts in a consortium like this were not designed together. Each one was set up years apart, by different teams, with their own questionnaires and their own ways of coding answers. One site records an exposure as yes or no. Another records it on a scale. A third measured it at a different age. Missing values are blank in one file and a special code in another.

In a normal project you would notice this when you plotted the merged table. In a federated project there is no merged table. You can't look at another site's rows. So a column that means something slightly different at each site trains happily and gives you a model that averages apples with oranges, and nothing in the loss curve tells you.

What helped:

  • A shared data dictionary, agreed before any code ran. Every variable had a name, a type, a unit and an allowed range.
  • A local check script that each site ran first. It returned only counts, ranges and the share of missing values per variable. Summary statistics, not rows.
  • A dry run on public data. The public version of the toolbox was validated on an open-access dataset, so anyone could see the pipeline run end to end before touching the real thing.

A site that is always down

Every federated setup has one. A server behind a hospital firewall that only opens on certain days. A machine that needs one particular person to log in. A site waiting on an ethics amendment before it can run anything new.

The textbook answer is to let rounds proceed with whoever shows up. That is fine for the maths and bad for the science if you don't track it. A site with a very different population that drops out of half the rounds pulls the model toward everyone else.

CarefulA missing site changes your model without raising an error. Log which sites contributed to every round, and report results with that list next to them.

Security sat on top of all this. A model update is not the same as the data, but it isn't nothing either. Part of my work was the security of the models themselves, so that what left a site couldn't be used to reach the sensitive data behind it.

Meetings across time zones

The sites were spread across several time zones. There is no comfortable hour for a call across all of them. Someone is always on at breakfast and someone is always on after dinner. That shaped how the work ran more than any technical choice.

A round of training that needed a person at each site could take days, not minutes. So the goal became: make every step something a site can run alone, on its own schedule, with clear output it can send back. The rhythm that works is roughly this.

  1. Agree the variablesOne data dictionary, signed off by every site before any code is written.
  2. Ship one scriptThe same code and the same arguments at every site. No local edits.
  3. Check locallyEach site runs the summary check and sends back counts and ranges only.
  4. Train in roundsSites train when they can. The aggregator waits, logs who took part, and averages.
  5. Review togetherOne call per cycle, at an hour that rotates so the same people don't always lose their evening.

Written instructions beat calls. A short document that says exactly which command to run and what file to send back saved more time than any meeting.

If you're about to start one

Before you pick a framework, do these three things:

  1. Write the data dictionary and get every site to agree to it. If you can't, the project isn't ready for a model.
  2. Build the local check script first, and make its output small enough to read in an email.
  3. Run the full pipeline on a public dataset, so that every site has seen it work before their own data goes in.

Then write the model. It will be the easy part.

related

Keep reading