The Model Comes to the Data: Federated Learning for Private AI

Your model gets better the more real documents it sees. Your clients have the documents. Their compliance team will never let them leave the building.

Federated learning is the answer that doesn't involve pretending. Instead of moving the data to the model, you move the model to the data.

The whole idea

The shared model is sent to each participant. Each one trains it locally, on their own hardware, on their own data. Each sends back only a summary of what changed — not the documents, not the records, just the adjustment. The server averages those adjustments into a better shared model, and sends it out again.

A model shaped by everybody's data, where nobody's data moved.

The idea dates to 2016. It's becoming practical now for three boring reasons: regulators keep making data transfer more expensive, the interesting data increasingly sits at the edge rather than in anyone's data lake, and efficient fine-tuning finally made local training possible on hardware clients already own.

Three concepts make it work. Federated averaging combines everyone's local result into a weighted mean. Low-rank adaptation (LoRA) freezes the original model and trains a tiny add-on instead, so the thing you transmit is megabytes rather than gigabytes. Secure aggregation and differential privacy close the remaining gaps. Secure aggregation hides individual updates from the server, and differential privacy stops anyone reverse-engineering a single record out of the aggregate.

★ The maths, for anyone who wants it

Skip this section if equations aren't your thing. Everything after it reads perfectly well without it — nothing later in the post depends on the formulas below. It's here for the engineers who reasonably want to see that the mechanism is real and not hand-waving. The two charts at the end of the section are worth a look either way.

Federated averaging, in two lines

There are only two moving parts. First, each participant does ordinary gradient descent on its own data. It's the same update any model training loop performs:

w _k^t+1 ​ = w^t − η ⋅ ∇ L_k ​ ( w^t )

Local step, on participant k

Then the server combines everyone's result into a single new global model, weighting each participant by how much data it holds:

Server aggregation across all participants

w ^t

the shared global model at round t — what gets sent out to everyone

η

the learning rate. Exactly the same knob as in any other gradient step; nothing federated about it

∇L _k

the gradient of participant k's loss on its own data. This is the only place raw data is ever touched, and it happens inside that participant's network

n_k ⁄ n

participant k's share of the total training examples. A bank with 40% of the documents gets 40% of the vote

That's the whole algorithm. The reason it works is a fact about averages: if every participant's data were drawn from the same underlying distribution, the average of their gradients would be an unbiased estimate of the gradient you'd have computed on the pooled dataset. Federated averaging is a way of computing that pooled gradient without ever pooling the data.

The reason it gets hard is that the "same underlying distribution" assumption is false in practice. That's the non-IID problem, and it has a precise shape: the further each participant's local optimum sits from the others, the more the weighted average lands somewhere that's worse for everybody than their own local solution. Mitigations like FedProx add a proximal penalty — a term that charges each participant for drifting too far from the global model during its local steps:

FedProx — the drift penalty, tuned by μ

Set μ to zero and you're back to plain federated averaging. Turn it up and each participant stays closer to the shared model at the cost of fitting its own data less well. It's a dial, not a fix.

Why the update is small: low-rank adaptation

A weight matrix in a transformer might be 8,192 × 8,192. Fine-tuning it normally means learning a full-size change to it. LoRA's observation is that the useful change is usually far simpler than the matrix it's changing — it can be approximated by the product of two thin matrices:

Low-rank adaptation

W_O

the original pretrained weights. Frozen — never modified, never transmitted, identical on every participant

B, A

the two thin matrices whose product approximates the needed change. These are the only things trained, and the only things sent back

r

the rank. A small integer, typically 8 to 64, that sets how much capacity the adapter has — and directly, how many bytes cross the network

The arithmetic is what makes this the load-bearing trick. Replacing a d×d update with two d×r matrices cuts the parameter count from d² to 2dr — for d = 8,192 and r = 16, that's a reduction of about 250×. Applied across a 31-billion-parameter vision model like the one we run, the difference is roughly this: the full weights at 16-bit precision are on the order of 60 GB, while a rank-16 adapter over the attention projections lands somewhere around a hundred megabytes. Nearly three orders of magnitude. One of those numbers you can send over a bank's internet connection every few hours; the other you cannot send at all.

Combining the two ideas — federated averaging applied to LoRA adapters rather than full weights — is sometimes called FedIT, and it's the only version of this that is practical for models of the size we actually run.

What the privacy budget actually costs

Differential privacy makes a formal promise: an observer looking at the output cannot tell whether any single record was in the training set. The guarantee is parameterised by ε, and it's bought by adding noise to each update before it's sent:

Differentially private update, with gradient clipping

The clipping step bounds how much any one participant's update can move the model — which incidentally is also the main defence against a malicious participant trying to poison it. The noise term is the actual privacy mechanism. Because σ scales inversely with ε, a tighter privacy guarantee means proportionally more noise, and more noise means more rounds to reach the same accuracy. This is the tradeoff in one line: privacy is paid for in convergence speed.

What this looks like in the results

Two charts. The first is real, published data. The second is a diagram of the general pattern — the reason a participant joins in the first place.

The second chart is the one that explains the incentive. A participant who trains only on its own data gets a model that fits that data well and then stops improving — it has run out of variety to learn from. The federated model starts out worse, because averaging across mismatched participants is noisy, but it keeps climbing, because it keeps seeing new variety. Somewhere there is a crossover, and after it the shared model is better than anything the participant could have built alone.

The argument for joining isn't that federated beats centralised; it's that federated beats going it alone, and centralised isn't legally on the table.

End of the maths. Everything from here is plain English again.

What it costs

Federated learning is not a free upgrade. It buys a privacy guarantee and bills you for it:

• Participants don't have similar data. One bank skews toward letters of credit, another toward customs declarations. Averaging their models can drag both away from where each was heading. Mitigations exist; a solution doesn't.
• Convergence is slower and noisier. More rounds to hit the same accuracy. Budget for it rather than discovering it.
• Someone might cheat. In a consortium of competitors, a participant can try to poison the shared model. Getting rivals to behave, round after round, is a business problem wearing an engineering costume.

It already works elsewhere

Healthcare is the strongest precedent, because HIPAA and GDPR made centralising patient data a non-starter. The FeTS initiative, charted above, trained a brain-tumour segmentation model across 71 institutions on six continents — and beat what any single institution could build, without a single scan leaving its own network.

The pattern repeats in rare-disease genomics, in pharma (MELLODDY put ten competing drug companies on a shared molecule-screening model without anyone revealing their compound library), and in banking, where consortium fraud pilots let each bank keep its transaction data in-house and still catch rings that no single bank's model could see.

Where it fits for us

Our pipeline is a chain of model calls over one submitted document set: a classifier decides what each page is, type-specific extractors pull fields out of bills of lading and invoices and packing lists, a cross-check stage compares them against each other and the letter of credit, and a reviewing agent resolves the discrepancies against UCP 600.

The models we'd actually be fine-tuning

A 31B vision model is the sweet spot for this. Big enough that a client's document conventions genuinely confuse it; small enough that a LoRA adapter trains on one GPU the client already owns.

Stochastic doesn't mean unpredictable — it means shapeable

Every one of these models is stochastic: it doesn't return an answer, it returns a probability distribution over the next token, and something has to pick from it. We've pinned that picking as hard as it goes. On all three call paths we run temperature = 0.0, seed = 42, and a repetition penalty between 1.05 and 1.2. A compliance product cannot have the quantity on a bill of lading change between two runs of the same document.

But temperature zero doesn't remove the randomness. Greedy decoding takes the argmax of a distribution that may be genuinely torn. When a smudged fax makes "1,700" and "1,100" nearly tied, the model still commits, giving no outward sign that it nearly said something else. Our classifier at least admits it, with a confidence floor at 0.7. The extractors mostly just answer.

Fine-tuning doesn't change how we pick. It changes the distribution we're picking from.

That's the whole argument for doing this. We aren't trying to make the model creative or to loosen decoding — we want the greedy pick to be right more often, and the near-ties to get rarer. In practice, that means supervised fine-tuning on (page image → corrected field values) pairs, token-level cross-entropy on the corrected JSON. That loss is exactly the ∇L_k from the maths section, computed inside the client's network on the client's own documents. Sharpen the distribution on this bank's stamp placements and date formats, and temperature 0.0 starts giving better answers without us touching a single decoding parameter.

The labels already exist.

This is the part that makes it worth doing rather than merely interesting. Our pipeline generates supervised training data as a by-product of operating:

• Every update_lc_field call the judge makes is a decision with a documented resolution path. That's a labelled example of what the right answer was.
• Every classification that lands under the 0.7 confidence floor is a flagged hard case, pre-sorted for us.
• Every manual reviewer entry is ground truth. Our own rule is that manual always wins — which is another way of saying we already treat these as labels.

All of it is generated inside the client's environment, all of it is as confidential as the document it came from, and today all of it dies where it was born. Federated learning is the only mechanism that turns it into model improvement without asking a bank to email us its customers' paperwork.

What we wouldn't fine-tune

The judge, at least not first. Its mistakes are usually about UCP 600 interpretation, and our prompts are external, version-controlled text files, which makes prompt iteration a same-day fix with global reach. Fine-tuning is slow, per-federation, and hard to roll back. Reach for it only where prompting genuinely can't help, which is visual convention: layout, stamps, handwriting, scan quality. That's the extractor's problem, which is why the extractor goes first.

Only the adapter comes back to us. No documents, no extracted values, no corrections.

And the plumbing is mostly there. The pipeline already reaches models through a single provider interface with two endpoint slots, so pointing a client's deployment at a model inside their own network is configuration, not redesign.

How we'd start

One client, one model, one document type. The 31B vision extractor, on commercial invoices. Not a rollout.

• Freeze an eval harness first. A held-out set from that client, fixed seed, per-field accuracy, run identically every round. Training is not reproducible the way our inference is, and averaging adapters from several clients adds another source of drift. That's why the measuring stick has to be the one thing that doesn't move.
• An existing federated framework. Building the round-orchestration harness ourselves would be the least interesting part of the project.
• One edge container. Serves extraction and does the local LoRA training, inside the client's network. Rank 16, attention projections only, roughly a hundred megabytes on the wire.
• Decide cadence and failure handling up front. What happens when a client drops mid-round is a design decision, not a production surprise.
• Add secure aggregation before widening. With one participant there's nothing to aggregate away. With five competitors there is.

Definition of done, agreed before we start: per-field extraction accuracy versus today's baseline, adapter size and round duration, what fraction of rounds complete without a stalled client, and how many rounds before quality visibly improves.

One measurement worth adding, and it's the only place we'd deliberately raise the temperature. Sample the same page a handful of times at temperature = 0.7 and see how much a given field wobbles. High variance means the distribution is torn there — which is both a free confidence signal for the reviewer queue and the cleanest way to tell whether an adapter actually made the model more certain, or just moved which answer it was confidently wrong about.

One closing note. A federated pilot doesn't create new reliability problems so much as find the existing ones faster. Bounded concurrency, clean resumption after partial failure, a clear answer to where data goes when a primary endpoint fails over. All of it is worth fixing whether or not we ever run a round. Federated learning just gives us a deadline.

References

• Pati, S. et al. Federated learning enables big data for rare cancer boundary detection. Nature Communications 13, 7346 (2022).
• McMahan, H. B. et al. Communication-Efficient Learning of Deep Networks from Decentralized Data (2016) — federated averaging.
• Li, T. et al. Federated Optimization in Heterogeneous Networks (2018) — FedProx.
• Hu, E. J. et al. LoRA: Low-Rank Adaptation of Large Language Models (2021).
• MELLODDY — EU Innovative Medicines Initiative consortium on federated drug discovery (2019–2022).