DistMoE: Stitching Private Experts Into One MLLM Without Sharing Any Data

Paper: 2608.09907 Authors: Mainak Singha, Niccolò Biondi, Elisa Ricci, Subhankar Roy Categories: cs.CV

The Gap

The standard recipe for adapting a multimodal LLM to a new visual-language domain is: collect the instruction data, pool it with everything else, and run joint instruction tuning (the LLaVA lineage, and every domain-specialized variant built on it — medical VQA, remote sensing, charts and documents, GUI agents). That recipe silently assumes one organization holds all the data at once.

Three families of prior work try to relax that assumption, and each one gives something up:

• Federated averaging of weights (FedAvg-style, including LoRA-federated variants): clients train locally, the server averages parameters. Averaging independently-trained nonlinear modules degrades sharply when client distributions differ; you get one blurred model, and you lose the ability to say “this capability came from client 3, revoke it.”
• Adapter/LoRA banks with post-hoc merging (task arithmetic, TIES, model soups): keeps modules separate, but merging assumes the modules live in a compatible coordinate system with comparable magnitudes. Independently trained modules typically do not.
• MoE upcycling of a dense FFN (MoE-LLaVA and relatives): works well, but the standard training loop assumes you can see all domains at once when learning the router. Routers are exactly the component that needs cross-domain contrast to calibrate.

So the specific boundary is: we know how to build a multi-domain MoE when data is centralized, and we know how to train locally when it is not, but we do not have a way to compose independently trained per-client experts without replaying other clients’ data. The blocker the paper names is concrete and mechanical — privately trained FFN experts drift to different output scales and magnitudes, so a router trained on one client’s activations cannot mix them sanely.

[P] Adapting an MLLM to many visual-language domains
    presumes centralized joint training over all data
              |
              v
[Constraint] Clients are private / permission-limited:
    no cross-client data movement, no cross-client replay
              |
              v
[A] Domain knowledge is largely FFN-local, so it can be
    parked in per-client experts .. IF their output
    scales/magnitudes are made comparable
              |
              v
[M] DistMoE
    stage 1: freeze public FFN, train private FFN per client
    stage 2: on (local data + PUBLIC data), train only
             router + tiny private projection adapter,
             with an isotropic regularization loss
              |
              v
[E] Diverse VL benchmarks: composed model competitive with
    joint training; experts reusable / pluggable; token-wise
    routing works with no domain label at test time
              |
              v
[C] Public data is a good enough shared anchor to replace
    cross-client rehearsal for expert composition

The Increment

One sentence: Before, combining independently trained per-client experts required either weight averaging (which blurs) or a joint calibration pass over pooled data (which is exactly what privacy forbids); after, you calibrate each expert against *public data plus a frozen public FFN, and the experts become mixable by a shared router without any client ever seeing another client’s samples.

Core Mechanism

The architecture is a per-layer two-slot MoE inside the language decoder. Every decoder layer keeps its original pretrained FFN — call it the public expert — and gains a private FFN expert owned by one client. In the first stage, each client trains only its own private FFN (plus whatever projector/visual glue the setup needs) on its local instruction data, with the public FFN frozen. This part is unremarkable and embarrassingly parallel: N clients, N private FFNs, zero communication. It is also where the problem is created. Each private FFN converges to whatever output statistics its own local loss preferred; one client’s expert may emit activations an order of magnitude larger than another’s, and the residual stream has no idea how to reconcile them.

The interesting stage is the second one, public-anchored expert composition. Here the private FFNs are frozen and only two lightweight things get trained: the per-layer routers and a small private projection adapter sitting on each private expert’s output. Crucially, the training mixture is *local client data plus public data — never another client’s data. Public data is the shared coordinate system: because every client’s composition step sees the same public distribution passing through the same frozen public FFN, the adapters are all being pulled toward a common reference frame rather than toward each other. That is what makes the composition rehearsal-free across clients: the usual fix for merging drift is to replay a buffer of other domains, and here public data plays that role without any privacy cost.

The regularizer is what turns “same reference frame” into “same scale.” An isotropic regularization loss pushes the private expert’s (post-adapter) representation toward having comparable energy in all directions instead of a few dominant, high-magnitude ones. Once each expert’s output has roughly whitened, comparable geometry, the router’s mixing weights mean the same thing across experts, and a weighted sum of public and private outputs stays inside the distribution the frozen decoder expects. At inference, routing is token-wise and label-free: each token’s hidden state hits the router, which decides how much public generalist and how much private specialist to draw on. A chart-reading token can pull from the chart client’s expert while a generic language token stays mostly public — with no one telling the model which domain the input came from.

                 h  (token hidden state, decoder layer L)
                 |
      +----------+----------+-----------------+
      |          |          |                 |
      v          v          v                 v
 [ Router R ] [Public   [Private FFN c1]  [Private FFN cN]
  (stage-2      FFN ]     (frozen after     (frozen after
   trained)    (frozen)    stage 1)          stage 1)
      |          |          |                 |
      |          |          v                 v
      |          |    [Proj Adapter 1]  [Proj Adapter N]
      |          |     (stage-2, tiny)   (stage-2, tiny)
      |          |          |                 |
      |          |          +-- isotropic reg |
      |          |              on outputs  --+
      |          |          |                 |
      |          v          v                 v
      +---> w_0 * y_pub + w_1 * y_1 + ... + w_N * y_N
                          |
                          v
                    h' (back to residual stream)

  stage 1 (per client, no comms):  train private FFN only
  stage 2 (per client, no comms):  train R + Adapter only,
           data = local client data + PUBLIC data,
           loss = task loss + isotropic regularization

The load-bearing metaphor: a remote recording session with no rehearsal.

Think of an album made by musicians who never meet. The public FFN is the backing track and click track — fixed, shipped to everyone, never re-recorded. Each client is a session musician in their own home studio recording one stem: the cellist in Berlin, the steel guitarist in Nashville. Stage 1 is everyone recording their part alone against the click. It works, and it needs no coordination.

Then you try to mix. The problem is exactly what every remote-collab producer hits: the stems come back at wildly different levels, with different room tone and different frequency emphasis. Sum them naively and the cello buries the guitar. This is the scale-and-magnitude drift.

Stage 2 is mastering-by-reference, done independently per stem. You do not send the cellist the guitarist’s stem (that is the cross-client rehearsal you cannot afford). Instead every musician masters against the same reference track — the public data. The projection adapter is the gain-and-EQ plugin on the channel strip: cheap, doesn’t touch the performance itself, just makes the stem sit correctly against the reference. The isotropic regularization is loudness and spectral normalization: flatten the freak peaks so no stem hogs the mix and so a fader position means the same thing on every channel. The router is the mixing engineer riding the faders — but riding them per bar rather than per song (token-wise, not sequence-wise), and doing it *by ear rather than from a cue sheet (no domain labels). Because every stem was mastered to the same reference, the engineer’s fader moves compose sensibly even though the musicians never played together.

The metaphor also predicts the failure mode, which is a good sign it’s the right one: if the reference track sounds nothing like the final song — public data far from a client’s domain — then mastering to it won’t make the stems sit together.

Key Concepts

• Mixture-of-Experts routing inside an FFN: A transformer layer has an attention block and then a feedforward block that is just “expand, nonlinearity, contract” — this is where most of the layer’s parameters and, empirically, most of its stored factual/domain knowledge live. MoE replaces that single block with several parallel blocks plus a tiny classifier (the router) that reads each token’s vector and decides which blocks to use for *that token. Concretely: for the token “ECG” the router might send 70% of the weight to the medical expert; for the token “the” it sends everything to the general expert. Two things worth internalizing: the decision is per token, not per sentence, and the router is a learned function of the hidden state, so it needs to have seen contrasting inputs to learn any contrast. That last point is precisely why distributed training breaks routers.

• Representation drift in scale and magnitude: Suppose two people are each told to write a number on a slip of paper indicating “how much I recommend this restaurant,” with no scale given. One writes 7 (out of 10), the other writes 850 (some personal scale). Both slips are informative; averaging them is nonsense. Neural modules trained independently do the same thing — nothing in a local loss pins down the *norm of the module’s output, only that it be useful to the rest of the local network. Since MoE combines experts by weighted sum into a residual stream, mismatched norms mean the loudest expert wins regardless of relevance. Fixing this requires a shared calibration signal, which is the whole ballgame here.

• Isotropic regularization: “Isotropic” means “the same in all directions.” Take the cloud of output vectors a module produces over many inputs and look at its shape. A drifted module produces a cigar: enormous variance along one or two directions, almost none elsewhere. An isotropic module produces a ball. Practically this is done by penalizing the spread of the representation’s singular values / off-diagonal covariance — a soft whitening. Why care? A ball-shaped output distribution has a well-defined characteristic magnitude, so scaling it by a router weight of 0.3 does something predictable; a cigar-shaped one does not, and it also wastes capacity because most of the output space is unused. It’s the neural-network version of normalizing every recording to the same LUFS before mixing.

• Rehearsal-free (in the cross-client sense): In continual learning, “rehearsal” means keeping a buffer of old data to replay so the model doesn’t forget. The natural fix for merging drifted experts is exactly that — replay a slice of every domain during calibration. Here that is forbidden, so the paper substitutes public data as a *proxy rehearsal set. The claim is not that no data is replayed; it’s that no client data crosses a boundary. Worth keeping straight, because it determines what the privacy guarantee actually is (a data-locality guarantee, not a formal one).

Framework Shift

Before (mainstream approach):          After (this paper):

  client A data                          [ Public FFN ]  frozen, shared
  client B data  --> POOL --+                 ^  anchor
  client C data             |                 |
                            v            +----+-------------------+
                    [ joint tuning ]     |    |         |         |
                            |          [PrivA][PrivB][PrivC]  .. (frozen)
                            v            |     |         |
                     one dense model     [ad] [ad]     [ad]  <- trained on
                     (no attribution,     \    |        /       local + PUBLIC
                      no removal)          \   |       /        w/ isotropic reg
                                            v  v      v
   ..or FedAvg:                             [ Router ] <- trained the same way
     w_A, w_B, w_C -> mean(w)                    |
              |                                  v
              v                        per-token mix, no domain label,
      averaged weights: drift          experts pluggable / removable
      collapses into blur

From pooling data to pooling modules: the core shift is replacing a shared training set with a shared reference distribution, so calibration — not data — becomes the thing that gets centralized.

Expert Assessment

A caveat on scope: I am working from the abstract and the framing here, not a line-by-line read of the tables. Treat the experimental section below as “what I would check and where the pressure points are,” not as a verified audit.

Problem choice: Real gap, and well-positioned. The trajectory is unmistakable — MoE upcycling of MLLMs became standard in 2024, modular adapter banks and merging became standard shortly after, and federated MLLM tuning has been the obvious next box to fill. The specific framing (composition without cross-client rehearsal) is sharper than the generic “federated multimodal tuning” papers because it names a mechanical failure — magnitude drift — rather than gesturing at “heterogeneity.” The modularity story also has genuine practical pull: per-client attribution and removal is a compliance feature, not just an ablation. Mild deduction because “clients” here is really “domains with a wall between them,” and the paper is closer to modular domain-incremental learning than to federated learning proper; some readers will feel the privacy framing is doing motivational work the method doesn’t cash out.

Method maturity: Clever, but the cleverness is in the assembly, not the parts. Public expert plus private expert per layer is upcycling; the projection adapter is a bottleneck adapter; isotropic/whitening regularizers have a decade of history in representation learning and have already been used for merging and for anti-collapse. The load-bearing idea is the *choice of anchor — using public data plus the frozen public FFN as the common coordinate system — and that idea is good and cheap. Two things nag. First, simpler baselines hover uncomfortably close: per-expert output RMS normalization, or fitting a scalar gain per expert on public data, might recover much of the isotropic loss’s benefit; I’d want that ablation explicitly, and I’d want it against a strong merging baseline (TIES / task arithmetic with rescaling) rather than plain averaging. Second, cost scaling: a private FFN per client per layer is a lot of parameters. With ten clients you are shipping something like a 10x-FFN model, and the abstract’s “lightweight” refers to the stage-2 adapters, not to the expert bank. Whether top-k routing keeps inference cost flat while memory blows up is the question a deployment reader will ask first.

Experimental integrity: The construction is the thing to scrutinize. Distributed-tuning papers get to choose their client split, and the difficulty of the whole problem is set by that choice — five visually and linguistically disjoint domains is a real test; five slices of the same instruction corpus is not. “Competitive performance” in the abstract is honest phrasing and I read it as slightly-below centralized joint training, which is the expected and acceptable outcome; if any table shows the distributed method *beating joint training, that needs a mechanism, not a victory lap. The other pressure points: (a) how much public data does stage 2 consume, and does performance degrade gracefully as that shrinks — because if it needs a lot of public data close to each client’s domain, the anchor assumption is load-bearing in a way that limits applicability; (b) does the router actually specialize, or does it collapse to near-uniform mixing (routing entropy per domain would settle this in one figure); (c) scaling in client count, since drift and router confusion should both worsen with N, and papers in this area frequently stop at three or four clients.

Writing quality: The abstract is well-constructed — it names the failure mode instead of hiding behind “heterogeneity,” which is rarer than it should be. The likely corner-cut is the justification of the isotropic loss: papers reach for whitening because it empirically works, and the gap between “different scales and magnitudes” and “therefore penalize anisotropy” needs an argument, ideally a measurement showing the drift is anisotropic rather than just a scalar mismatch. If I could have one section rewritten, it would be that: a short diagnostic — spectra of private expert outputs before and after composition, and routing weight distributions across clients — would convert the method from plausible engineering into an explained phenomenon, and it would preempt the “why not just normalize the norms?” reviewer.

Verdict: weak accept — a genuinely useful and well-scoped assembly of known parts around one good idea (public data as the merging anchor), whose acceptance should hinge on whether the ablations separate “isotropic regularization” from “any output normalization” and whether the client splits are hard enough to be interesting.

Takeaways

Things here that transfer beyond this paper:

• Use a public reference distribution as your merging coordinate system. This is the reusable trick. Any time you have independently trained modules that must later be combined — LoRA marketplaces, per-tenant adapters, per-team fine-tunes — you can make them compatible by having each one calibrated against a shared public batch through a shared frozen trunk, without ever moving private data. It’s an alternative to both weight averaging and joint retraining, and it costs one extra cheap stage.
• Separate “learn the knowledge” from “learn how to be combined.” Stage 1 trains capacity; stage 2 trains only interfaces (router plus adapter) with capacity frozen. Structuring modular systems this way means adding a new participant never touches existing experts and never requires their data. That’s a systems pattern, not a vision pattern.
• Anisotropy is a merge blocker, and it is measurable. Before you debug a bad model merge with hyperparameters, look at the singular value spectrum of each module’s outputs on a common probe set. A cigar-shaped output cloud tells you the modules aren’t speaking the same language and no mixing weight will fix it. Cheap diagnostic, widely applicable.
• Token-level routing removes the need for a domain classifier. If you have been maintaining a router that first classifies the input’s domain and then dispatches to a specialist, note that per-token mixing inside the FFN subsumes that, and handles inputs that straddle domains (a medical chart in a slide deck) rather than forcing a single choice.
• What not to take on faith: “rehearsal-free” means no *client data is replayed — public data still is. If your setting has no public data resembling your clients’ domains (genuinely idiosyncratic internal corpora), the anchor disappears and the method’s central assumption goes with it.