Information Flow Control in Machine Learning through Modular Model Architecture

TL;DR

This work defines information flow control (IFC) for machine learning and formalizes non-interference (NI) in the ML context. It then presents a parametric IFC approach for Transformer language models that uses domain-specific experts, secure gating, and secure aggregation to ensure outputs depend only on accessible training data. The method yields substantial accuracy improvements (up to 38% for text and 44–62% for code) with minimal runtime overhead (about 1.9% latency) compared to non-IFC baselines. This enables secure training and inference over access-controlled data, with practical implications for domains like internal coding assistants and confidential document processing.

Abstract

In today's machine learning (ML) models, any part of the training data can affect the model output. This lack of control for information flow from training data to model output is a major obstacle in training models on sensitive data when access control only allows individual users to access a subset of data. To enable secure machine learning for access-controlled data, we propose the notion of information flow control for machine learning, and develop an extension to the Transformer language model architecture that strictly adheres to the IFC definition we propose. Our architecture controls information flow by limiting the influence of training data from each security domain to a single expert module, and only enables a subset of experts at inference time based on the access control policy.The evaluation using large text and code datasets show that our proposed parametric IFC architecture has minimal (1.9%) performance overhead and can significantly improve model accuracy (by 38% for the text dataset, and between 44%--62% for the code datasets) by enabling training on access-controlled data.

Figures (17)

• Figure 1: Information flow control (IFC) in machine learning.
• Figure 2: Pre-trained GPT-2 evaluated on 50 datasets from Pushshift.io consistently underperforms fine-tuned GPT-2.
• Figure 3: Parametric IFC through modular model architecture. The influence of each security domain is limited to one expert. A gating function, at inference time, decides which of the experts should be activated for a given input. The activated experts are aggregated either though ensembling their outputs or aggregating the parameters.
• Figure 4: A typical Transformer-based language model (LM) architecture (based on GPT-2).
• Figure 5: The overview of the proposed Parametric IFC scheme for a Transformer-based language model.