Information Flow Control in Machine Learning through Modular Model Architecture

TL;DR

Today’s ML models smear every training example into one set of weights, so anything the model saw, it can leak — a problem when different users are allowed to see different data (internal repos, paywalled content, project docs). This paper defines Information Flow Control (IFC) for ML using a clean security guarantee (non-interference: inaccessible data must have zero influence on the output), then shows how to actually achieve it. The trick: train one small expert (an adapter) per security domain so each domain’s data only touches its own parameters, then build gating and aggregation functions that only ever read information from the experts a user can access. The result enforces strict access control while improving accuracy by 38% (text) and 44–62% (code) over a public-only model, at just 1.9% latency overhead. It’s the access-control layer that RAG-only privacy stories are missing for the parametric (fine-tuned) half of a system.

Problem & Motivation

Concrete pain: a company wants an in-house code-completion model trained on its private repos, but engineers have access to different subsets of those repos. Two bad options today:

• Train on everything → the model can autocomplete code from repos an engineer was never cleared to see. That’s a leak.
• Train only on public/shared data → the model is mediocre because it can’t use the good proprietary data.

Same shape shows up for writing assistants over confidential project docs, and for Q/A over paywalled or classified corpora. The root issue: in a normal neural net any training example can influence any output. There is no notion of “this output is only allowed to depend on data X.” Access control exists for files (RBAC, ACLs) but evaporates the moment you fine-tune a model on those files — the weights become a single undifferentiated blob.

Why the obvious fixes fail:

• Fine-tune one model per access policy — there are up to 2^m policies for m domains. Exponential. Dead on arrival.
• Fine-tune one model per domain, pick one at query time — works, but a user with access to 10 domains gets the benefit of only one expert. Leaves accuracy on the table.
• Differential Privacy (DP) — protects individual samples and everyone equally; you can’t say “Alice benefits from domain 3, Bob doesn’t.” DP also only bounds leakage (ε > 0), never eliminates it, and degrades badly when you must protect a whole correlated group (a domain) at once.

What’s New (Core Contribution)

1. A formal definition of IFC for ML. Before: “privacy” in ML meant DP or vague hand-waving. Now: a precise non-interference guarantee borrowed from classic information-flow security — if two training sets agree on the data a user can access, the model’s output for that user must be identical, regardless of the inaccessible data. Zero leakage, not bounded leakage.
2. A modular architecture that actually satisfies it. Before: Mixture-of-Experts (MoE) used learned gating trained on all data — which itself leaks. Now: per-domain experts + secure gating + secure aggregation, every step provably touching only accessible domains.
3. Three real-time secure gating functions. gate-known (user supplies a domain label), gate-pairwise (infer relevant domains by embedding similarity), and gate-cluster (hierarchical version that scales to tens of thousands of domains). The hard constraint they all respect: the access policy arrives at query time and may differ per request, so gating must be fast and may never compute over inaccessible domains.
4. Two secure aggregation methods. Output ensembling (Bayesian-weighted combination of expert outputs) and parameter merging (average the top-k experts’ weights into one). Both training-free and leakage-free.
5. A generalization to ε-non-interference (future-work section) that allows bounded leakage à la DP, but with adjacency defined over whole domains instead of single examples.

The genuinely new part is #1 and the security-constrained redesign of gating/aggregation in #3/#4. Modular models existed; secure modular models that enforce non-interference did not.

How It Works (Technically)

The whole system has three phases. Keep this mental model: one domain → one expert; gating picks accessible experts; aggregation blends only those.

Phase 1 — Domain-aware training (parameter isolation). Start from a public pre-trained model (GPT-2 or OPT). For each security domain d_j, fine-tune only a small set of inserted parameters — an adapter block placed right before the MLP in each Transformer block — on that domain’s data alone. The frozen backbone is shared; the adapter is the “expert.” Because expert e_j = train(d_j) depends on d_j and nothing else, domain j’s data physically cannot influence any other expert. This is the parameter-isolation guarantee, enforced by construction rather than by hoping the optimizer behaves. Adapters are the sweet spot: 14.1M trainable params vs 117M for full fine-tuning, yet within 0.1% of full-FT accuracy.

Adapter, for the reader: a small bottleneck network (down-project → nonlinearity → up-project) inserted into a frozen model. You train only the adapter, so each “expert” is tiny and cheap to store/swap. This is the same family as LoRA — think “a few MB per domain,” not “another full model.”

Phase 2 — Secure gating (which experts fire). At inference the model gets: the user’s access policy a_i (set of allowed domain indices), the input tokens x, and optionally a domain label l. It must return the top-k most relevant accessible experts (default k=3).

• gate-known — Offline, build an all-to-all accuracy matrix M: row = domain of some held-out text, column = which expert evaluated it, value = perplexity. At query time, take the user-supplied label l, mask out every column not in a_i, and read off the best-performing accessible experts for domain l. (If l names an inaccessible domain, abort — even revealing “you can’t see that” is controlled.) Only accessible columns are ever read, so no leakage.

• gate-pairwise — No label given. Offline, compute a representative embedding vector v_j for each domain by running its text through frozen BERT and averaging the 768-dim token embeddings. At query time, embed the first c tokens of the user’s input into v', then compute cosine similarity φ_n = cos(v', v_n) only for accessible domains n ∈ a_i. Higher similarity = the user’s text looks like that domain = its expert is probably useful. Rank, take top-k. The key security move: similarities against inaccessible domains are never computed, so the ranking can’t encode anything about them.

• gate-cluster — Same idea, but for huge m. Offline, cluster the domain vectors into s clusters. At query time, compute each cluster’s center using only accessible domains, find the nearest cluster to v', then do pairwise similarity only within that cluster ∩ accessible. Turns an O(m) scan into a hierarchical search — the difference between “pairwise works up to ~tens of thousands of domains” and “cluster scales past that.”

The score adjustment (Equation 1). Raw cosine similarity ignores how well-trained an expert is — a domain with 900K samples is a better expert than one with 19K even if the topic match is weaker. So they nudge the score:

S_n = count(d_n) / Σ_{j∈a_i} count(d_j)     # n's share of accessible training data, in [0,1]
φ'_n = φ_n + λ · S_n                          # λ∈[0,1] tunes "trust bigger domains"

Plain English: final score = topic-match + λ × (how much data this expert was trained on, relative to other accessible domains). With λ=0.4, a well-trained-but-slightly-off-topic expert can still beat a tiny-but-on-topic one. Crucially S_n’s denominator sums only over accessible domains, so it stays leakage-free.

Phase 3 — Secure aggregation (blend the k experts into one output).

Output ensembling (the better performer). Run the input through all k experts, get k next-token distributions o_{t,1}...o_{t,k}, and combine them weighted by how likely each expert’s domain is, given the text so far. The weight is a Bayesian posterior (Equation 2/3):

P(d_t = j | x_<t)  =  P(x_<t | d_t=j) · P(d_t=j)
                      ---------------------------------------
                      Σ_{j'} P(x_<t | d_t=j') · P(d_t=j')

Translation: “given the text I’ve seen, what’s the probability it belongs to domain j?” Each expert already hands you P(x_<t | d_t=j) — that’s just the likelihood it assigns to the text so far (a well-fit expert assigns high likelihood to its own kind of text). Assume a uniform prior P(d_t=j), normalize across the k experts, and you get per-expert weights. Final output (Equation 4):

o_t = Σ_{j=1..k} o_{t,j} · P(d_t = j | x_t)

So the experts vote, weighted by how well each one explains the running context. No training, only accessible experts involved → non-interference holds.

Parameter merging (the cheaper alternative). Instead of k forward passes, average the k experts’ weights (normalized top-k scores as merge weights) into one merged expert, then run a single forward pass. Faster, but ~9 points worse at max (29% vs 38% improvement). Useful when latency/memory dominates.

One end-to-end trace. Engineer Bob (access to repos {Python-A, Python-B, JS-C}) types def parse_config(. → gate-pairwise embeds those tokens with BERT, scores cosine similarity against only the representative vectors of A, B, C, adjusts by data-size, picks top-3 (here all 3). → Each expert produces a next-token distribution; the Bayesian weighting favors whichever repo’s expert best explains def parse_config(. → Weighted-sum gives the suggestion. → Repos Bob can’t see were never embedded, scored, or run. Re-gate every r tokens to handle topic drift, in parallel so it adds no latency.

Architecture & data flow

flowchart LR
  subgraph Train["Phase 1: Domain-aware training (offline)"]
    D1[Domain d1] --> E1[Expert e1 - adapter]
    D2[Domain d2] --> E2[Expert e2 - adapter]
    Dm[Domain dm] --> Em[Expert em - adapter]
  end
  Q[User input x] --> G
  AP[Access policy a_i] --> G
  L[Optional domain label l] --> G
  G[Secure Gating<br/>known / pairwise / cluster] -->|top-k accessible experts| AGG
  E1 -. only if accessible .-> G
  E2 -. only if accessible .-> G
  Em -. only if accessible .-> G
  AGG[Secure Aggregation<br/>ensemble outputs OR merge params] --> O[Output o]
  note["Inaccessible experts are never read:<br/>no influence on gating or aggregation"]

Interactive: toggle which domains a user can access and watch gating score only the accessible experts (by topic similarity + data-size bonus), pick the top-k, then blend their votes. Inaccessible domains stay dark — they never enter any computation.

Interactive: the non-interference guarantee made visual. Flip the contents of an *inaccessible* domain and confirm the output distribution does not move at all (zero leakage) — versus a normal monolithic model where it shifts.

The algorithm, simplified

# Secure inference: pick accessible experts, blend only those.
# Stubs: embed(text)->768-vec (frozen BERT), expert(e, x)->next-token logprobs.

def secure_generate(x, access_policy, domain_vecs, experts, k=3, lam=0.4, sizes=None):
    # ---- GATING: only ever look at accessible domains ----
    v = embed(x[:c])                                  # vectorize sample of user input
    scored = []
    for n in access_policy:                           # NOTE: never iterate inaccessible domains
        phi = cosine(v, domain_vecs[n])              # topic similarity
        s_n = sizes[n] / sum(sizes[j] for j in access_policy)
        scored.append((phi + lam * s_n, n))          # Eq.1: similarity + data-size bonus
    topk = [n for _, n in sorted(scored, reverse=True)[:k]]

    # ---- AGGREGATION: Bayesian-weighted ensemble of accessible experts ----
    out = []
    for t in range(len(x)):
        logp = {n: expert(experts[n], x[:t+1]) for n in topk}      # each expert's distribution
        # posterior P(domain=n | context) from each expert's own likelihood, uniform prior
        like = {n: context_likelihood(logp[n], x[:t]) for n in topk}
        Z = sum(like.values())
        w = {n: like[n] / Z for n in topk}                         # Eq.2-3: normalized weights
        out.append(sum(w[n] * softmax(logp[n][-1]) for n in topk)) # Eq.4: weighted vote
    return out
# Security invariant: every loop ranges over access_policy / topk ⊆ access_policy,
# so inaccessible domains have provably zero influence on the output.

Built on Prior Work

Results & Evidence

Setup. GPT-2 on Pushshift.io (Reddit), 50 subreddits = 50 domains, 19K–500M tokens each. OPT on Codeparrot GitHub code, 79 MIT repos across Python/JS/C/C++/Java. Baselines: (a) pre-trained-only (secure but weak) and (b) fine-tuned-on-everything (insecure upper bound). Metric: perplexity (lower = better), reported normalized to pre-trained.

Headline numbers.

• Accuracy: when all domains accessible, perplexity improves 38% (text) and 44–62% (code) over public-only — versus the insecure full-FT ceiling of ~48% (text) and ~75% (code). So IFC captures most of the achievable gain while leaking nothing.
• Even with one accessible domain, text perplexity drops 11.5%.
• Latency overhead: ≤1.9% worst case (even at 10,000 domains), memory overhead ≤13% — provided you run the k expert forward passes in parallel.
• Adapters: within 0.1% of full fine-tuning accuracy at ~1/8 the trainable parameters.
• Gating ranking: gate-known > gate-pairwise > gate-cluster, but pairwise-vs-cluster differ only ~3% (text) / ~9% (code).
• Target-domain-excluded (user queries a domain they can’t see): with enough other accessible experts, accuracy is only 4–8 points worse than if the target were available — non-target experts compensate.
• Aggregation: ensembling (38%) beats parameter merging (29%).

What the evidence establishes: the architecture enforces strict non-interference by construction (it’s an architectural guarantee, not an empirical one) and the accuracy/latency costs of that guarantee are small on two realistic corpora.

What it does NOT establish / caveats:

• Only GPT-2 and OPT, only next-token perplexity. No modern instruction-tuned LLM, no downstream-task or human-eval results. Perplexity ≠ user-perceived quality.
• “Security domains” are simulated (subreddits, repos) — not real access-controlled enterprise data with real correlation structure.
• The 1.9% latency claim assumes k parallel GPUs/streams for the forward passes. Serially, k experts ≈ k× the forward-pass cost. That’s a real deployment cost the headline number hides.
• The non-interference guarantee is information-theoretic over the defined channels (gating + aggregation outputs). It does not analyze timing/cache side channels in the actual implementation.
• Per-domain adapters means storage grows linearly with domains; fine at 79, a question at millions.

How You’d Use It

This is the missing access-control layer for fine-tuned models — directly relevant if you sell AI systems into regulated or multi-tenant environments.

• Multi-tenant SaaS with per-customer fine-tuning. Today you either run a separate model per tenant (expensive) or one shared model (leaks tenant A’s data into tenant B’s completions). One shared backbone + one adapter per tenant + secure gating gives you a single deployable model that provably never crosses tenant boundaries. That’s a sellable compliance story (“your data only influences your outputs, guaranteed by architecture”).
• Enterprise coding/writing assistants over RBAC’d repos or docs. Map each repo/project/clearance level to a security domain. The assistant honors existing ACLs at inference — no model retraining when someone’s access changes, just flip which adapters their policy admits.
• Paywalled / licensed-content Q&A. Subscribers to different content tiers get a model whose answers can only draw on what they’ve paid for — a clean answer to publisher licensing concerns.
• Attribution & “right to be forgotten.” Because only k experts touched an output, you can tell a client which domains influenced a generation (attribution), and you can drop a domain by deleting its adapter — coarse-grained unlearning without retraining the whole model.
• In your MAS world: treat each agent/role’s knowledge base as a domain; agents compose only the experts their role is cleared for, with the same non-interference guarantee across agent boundaries.

Effort to stand up a demo: low-to-moderate. The pieces (PEFT adapters, BERT embeddings, cosine gating, ensemble) are all off-the-shelf. The “secure” discipline is mostly what you refuse to compute, which is cheap to enforce in code.

Build Your Own (Minimal Recipe)

Smallest version that captures ~80% of the value:

1. Backbone + adapters. Pick a small open model (e.g., a 1–3B). Use peft (LoRA adapters) and train one adapter per domain on that domain’s data only. Save each adapter (a few MB).
2. Domain vectors. Run each domain’s text through a frozen sentence/BERT encoder, average to one vector per domain. Store them.
3. Secure gating (start with gate-pairwise). At query time: embed the first ~c tokens, cosine-similarity against only accessible domain vectors, add the λ·data-size bonus, take top-k. Add the gate-known label path (read from a precomputed accuracy matrix) as a fast/accurate option when the caller knows the domain.
4. Aggregation (start with output ensembling). Run the k adapters, compute each one’s likelihood on the running context, normalize to weights, blend the next-token distributions. Parameter-merging is a v2 optimization.
5. Enforce the invariant in code. The entire security property reduces to: never iterate over, embed, score, or run an inaccessible domain. Make access_policy the only iterable the gating/aggregation loops ever see, and add a test that flips inaccessible-domain data and asserts byte-identical output.

The genuinely hard parts: (a) parallelizing the k forward passes so latency stays flat — without it your costs are k×; (b) scaling gating past a few thousand domains (that’s when you need gate-cluster); (c) tuning k and λ per use case (k=3, λ=0.4 are their defaults, not laws).

Reach for: transformers + peft (adapters/LoRA), sentence-transformers for domain vectors, FAISS if you go to many domains and want fast nearest-cluster.

How to Improve It

1. Modern LLMs + real tasks. Replicate on an instruction-tuned 7B+ model and evaluate on downstream tasks and human prefs, not just perplexity. This is the biggest open question for production credibility.
2. Learned-but-secure gating. The gating is hand-built (cosine + heuristics) to stay leakage-free. Could you train a gating policy using only per-domain-isolated signals (e.g., one gating head per domain, composed at runtime) to beat cosine while preserving non-interference? Test against gate-known as the ceiling.
3. Implement and audit the side channels. The math guarantees no leakage through gating/aggregation values; a real deployment leaks through timing and memory (how many experts ran, which cluster was hit). Add constant-work gating (always do k forward passes, pad cluster search) and measure the latency tax.
4. Hierarchical / shared experts to fight linear storage. One adapter per domain is fine at hundreds, painful at millions. Explore a tree of shared sub-experts where a domain is a path, keeping isolation while sub-linear in storage — without reintroducing cross-domain leakage.
5. Ship ε-non-interference for real. The paper sketches bounded leakage (ε-NI) but leaves the mechanism for future work. Design and benchmark a noised-aggregation scheme that trades a tiny, quantified leak for big accuracy gains on out-of-policy queries — the practical knob most clients will actually want.
6. Combine with Federated Learning (their own suggestion): train each domain’s expert via FL so raw data never centralizes (train-time privacy) and compose with IFC (inference-time access control) — a strong dual-guarantee offering.

Glossary

• Information Flow Control (IFC) — Ensuring a computation’s output only depends on inputs the requester is authorized to see; classic in OS/security, here applied to ML training data.
• Non-interference (NI) — The formal guarantee: change only the inaccessible data and the output must not change at all. Zero leakage.
• ε-non-interference (ε-NI) — Relaxed version allowing a bounded (ε) amount of influence from inaccessible data, in the spirit of differential privacy.
• Security domain — A partition of the training data that has its own access permission (e.g., one repo, one subreddit, one project).
• Expert — A small set of fine-tuned parameters (here, an adapter) trained on exactly one domain’s data; the unit of isolation.
• Adapter — A small bottleneck module inserted into a frozen Transformer; you train only it, so each expert is cheap (cf. LoRA).
• Access policy (a_i) — The set of domains a given user is allowed to access; supplied per query at inference time.
• Gating — Choosing which experts to activate for a query; here it must read only accessible domains.
• Aggregation — Combining the chosen experts into one output (output ensembling or parameter merging).
• Top-k — Activating only the k most relevant experts (default k=3) rather than all accessible ones, for speed and accuracy.
• Perplexity — Standard language-model accuracy metric; exp of the average negative log-likelihood. Lower is better.
• Mixture-of-Experts (MoE) — Architecture with many expert sub-networks and a (usually learned) router; the insecure ancestor of this design.
• Differential Privacy (DP) — Adds noise so the output is nearly unchanged whether or not any single sample was present; uniform and per-sample, unlike IFC’s selective, per-user, zero-leak guarantee.
• Bayesian posterior weighting — Weighting each expert by P(domain | context-so-far), computed from the expert’s own likelihood on the text.
• Parameter merging — Averaging multiple experts’ weights into one model to avoid running several forward passes.
• Cosine similarity — Angle-based closeness of two embedding vectors; the core relevance score in pairwise/cluster gating.