Based on Bagel Labs' CVPR 2026 paper on Heterogeneous Decentralized Diffusion Models: arxiv link

Decentralized Diffusion Models (DDMs) train independent experts on disjoint data partitions and combine them at inference time. Existing DDM frameworks assume all experts share the same training objective. We relax this constraint. In our setup, some experts train with DDPM (ε-prediction) and others with Flow Matching (velocity-prediction), then unify at inference through a deterministic conversion into a common velocity space. No retraining, no fine-tuning, no coordination during training.

Heterogeneous experts trained independently on single GPUs, unified at inference through velocity conversion.

Under aligned inference settings, this heterogeneous ensemble (2 DDPM + 6 FM experts) achieves better FID (11.88 vs. 12.45) and higher intra-prompt diversity (LPIPS 0.631 vs. 0.617) than a homogeneous ensemble of 8 FM experts. Relative to the training scale reported for prior DDM work, our framework reduces compute from 1176 to 72 A100-days (16×) and data from 158M to 11M images (14×), with each expert requiring only 24–48 GB VRAM on a single GPU, making decentralized diffusion training accessible on commodity hardware.

Background: Decentralized Flow Matching

DDMs decompose a generative model into K experts, each trained on a semantically coherent subset of the data. Following prior DDM work, the marginal velocity field is expressed as a weighted combination of per-expert conditional flows:

where uₜ⁽ᵏ⁾(xₜ) is the velocity predicted by expert k trained on cluster Sₖ, and pₜ(k | xₜ) is a posterior weight from a learned router.

Training

We partition the dataset into K semantic clusters using DINOv2 features (1024-dimensional representations with hierarchical k-means). This produces semantically coherent partitions, e.g. portraits, landscapes, architecture. Each expert θₖ trains exclusively on its assigned cluster Sₖ with zero communication between experts. No shared parameters, no gradient synchronization, no activation passing.

Routing

A lightweight router (DiT-B, 129M parameters) learns to predict cluster assignments from noisy inputs:

trained with cross-entropy loss against ground-truth cluster labels. At inference, the router dynamically selects and weights experts based on the current noisy state and timestep. We support three selection modes: Top-1 (single best expert), Top-K (weighted ensemble of K highest-probability experts), and Full Ensemble (all experts weighted by router probabilities). As shown in our prior work, Top-2 routing consistently outperforms Full Ensemble because sparse routing maintains expert-data alignment, selecting experts that are in-distribution for the current denoising state.

Heterogeneous Objectives

Previous DDM work requires all experts to share the same training objective. This is a coordination requirement that may be impractical when contributors operate independently. We remove this constraint.

Two objectives, different emphasis

We train n experts with Flow Matching loss and m experts with DDPM loss. DDPM experts predict the noise ε added during the forward process:

Flow Matching experts predict the velocity field directly:

Here x₀ ∈ Sₖ means expert k trains only on clean samples from its assigned cluster Sₖ, ε ~ N(0, I) is Gaussian noise, t is the sampled timestep, and αₜ, σₜ are the schedule coefficients controlling signal and noise strength, xₜ = (1-t)x₀ + tε is the linear interpolation between clean data and noise.

Both objectives model the same generative process through different parameterizations. But they weigh errors differently across timesteps, and this difference is the mechanism behind heterogeneous ensembles.

Complementary loss weighting

We can write both losses in terms of the squared clean-sample estimation error |x̂₀ - x₀|² (the detailed derivation can be read from variational diffusion models). Under a variance-preserving schedule (αₜ² + σₜ² = 1):

The ratio between them is:

Here w_ε(t) and wᵥ(t) are the effective per-timestep weights after rewriting each loss in terms of clean-sample estimation error. Since αₜ ≤ 1 and decays toward 0 at high noise, this ratio diverges. Velocity-prediction experts receive relatively stronger gradients at high-noise timesteps (global structure), while ε-prediction experts are relatively upweighted at low noise (fine details). In the paper, we derive this under a variance-preserving schedule and then note that linear interpolation recovers the same 1 / αₜ² structure. So the complementary weighting pattern holds both for the VP analysis and for the linear FM path used here.

DDPM concentrates learning signal near clean images. Flow Matching emphasizes high-noise timesteps. Their blind spots are complementary.

The implication is that each objective has a “blind spot” region where its effective weight is relatively low, and these blind spots are complementary. Where DDPM under-trains (high noise), FM trains hardest. Where FM under-trains (low noise), DDPM concentrates its signal. Mixing both in an ensemble lets each objective cover the other’s blind spots, providing more uniform coverage across the full denoising trajectory.

Inference-Time Unification

The central technical challenge is that DDPM experts output noise predictions ε_θ(xₜ, t) while FM experts output velocity predictions v_θ(xₜ, t). These live in different spaces. You cannot average them directly.

We unify all expert predictions into a common velocity space through a deterministic, schedule-aware conversion. The derivation proceeds in three steps.

The conversion pipeline: DDPM noise predictions are converted to velocity through deterministic algebra. FM predictions pass through unchanged.

Step 1: Recover the clean-image estimate

From the DDPM forward process xₜ = αₜ x₀ + σₜ ε, invert the linear map using the model’s noise prediction:

Step 2: Derive the velocity

Treating x̂₀ and ε_θ as fixed at their current-timestep values defines a deterministic path x̃ₜ = αₜ x̂₀ + σₜ ε_θ. Differentiating with respect to t:

Under linear interpolation (αₜ = 1-t, σₜ = t), the schedule derivatives are -1 and +1, so this simplifies to:

This is exactly the FM velocity target v = ε - x₀.

Step 3: Combine

FM experts already output velocity, so they pass through unchanged. All predictions are now in the same space. The router assigns per-expert weights, and we take a weighted combination to form the ensemble field uₜ, which drives a standard ODE integration step (x_{t−Δt} = xₜ − uₜ · Δt).

Numerical stability

The conversion requires dividing by αₜ, which approaches zero at high noise. We apply three safeguards: (1) clamp x̂₀ to [-20, 20] for VAE latents, (2) use α_safe = max(αₜ, 0.01) in the denominator, and (3) apply adaptive velocity scaling that dampens converted predictions at elevated noise levels where schedule derivatives become large. These are simple clamps with no learned parameters.

The entire conversion is closed-form algebra. No learned components, no fine-tuning, no additional training of any kind.

Efficient Training

Prior DDM work required 1176 A100-days on 158M images. We achieve competitive quality with 72 A100-days on 11M images, a 16× reduction in compute and 14× in data. Three techniques enable this.

Resource comparison: our approach requires a fraction of the compute and data of prior DDM work.

Pretrained checkpoint conversion

We initialize experts from ImageNet-pretrained DiT checkpoints. Patch embeddings, positional embeddings, and all transformer blocks are fully transferred. Only the final projection layer (which differs between ε- and velocity-prediction targets) and the text projection (new modality) are reinitialized. Class-conditional embeddings from ImageNet pretraining are removed.

A key technical detail is timestep compatibility. DiT models expect discrete timesteps t ∈ {0, 1, …, 999} while Flow Matching uses continuous t ∈ [0, 1]. We handle this with runtime conversion (t_DiT = round(999t)) rather than modifying pretrained weights, preserving the learned timestep embedding structure.

Converted checkpoints reach validation loss parity 1.2× faster than training from scratch.

Efficient architecture

Each expert uses DiT with PixArt-α’s AdaLN-Single conditioning. Rather than computing adaptive layer normalization parameters with per-block MLPs, a single global MLP produces all modulation signals:

reshaped into per-block slices plus learned per-block embeddings Eᵦ. This reduces parameters by 30% (891M to 605M for DiT-XL/2) while maintaining generation quality.

True isolation

Each expert trains on its own semantic cluster on a single GPU requiring 20–48GB VRAM. No gradient synchronization, no activation passing, no pipeline coordination, no parameter servers. This is not data parallelism with relaxed communication. There is literally zero inter-expert communication during training.

Experiments

We train on 11M LAION-Aesthetics images. For a high-quality subset of 3.9M images, we use LLaVA to generate improved captions. We evaluate using FID-50K on a held-out 50K test set. We train at two scales: DiT-B/2 (129M parameters per expert) and DiT-XL/2 (605M parameters per expert).

Our standard configuration uses K=8 experts. Experts 0 and 3 train with DDPM objectives (assigned to clusters containing high-fidelity subjects where ε-prediction excels at detail preservation). The remaining six use Flow Matching.

Monolithic versus decentralized

We first validate that decentralized training works. Using DiT-B/2 experts trained from scratch on LAION-Art (3.9M images), all with Flow Matching:

Top-2 achieves 22.60, outperforming the monolithic baseline by 23.7%. Full Ensemble underperforms dramatically (47.89), consistent with our prior finding that indiscriminate combination introduces prediction conflicts from out-of-distribution experts. Selective expert activation is essential.

Homogeneous versus heterogeneous

To isolate the effect of objective heterogeneity, we compare homogeneous and heterogeneous 8-expert DiT-XL/2 models under matched inference settings.

Under aligned settings (first and last rows), heterogeneous experts improve FID from 12.45 to 11.88. Scaling from 1 to 2 DDPM experts improves FID from 19.75 to 15.09 under the conversion setting, suggesting the optimal DDPM:FM ratio deserves careful tuning per domain.

For diversity, we measure intra-prompt LPIPS by generating 10 images per prompt for 100 held-out prompts. Heterogeneous experts achieve 0.631 (± 0.078) vs. homogeneous 0.617 (± 0.074). Objective heterogeneity produces more varied outputs for identical prompts.

Heterogeneous ensembles improve both image quality (lower FID) and output diversity (higher LPIPS) over homogeneous baselines.

Conversion quality

We evaluate the DDPM → FM conversion in isolation, using experts trained on the same data cluster (to isolate objective effects from data distribution differences). Both DDPM and FM experts use DiT-XL/2 with identical hyperparameters.

Three findings emerge. First, the conversion works. DDPM → FM improves over native DDPM (FID 25.61 vs. 27.04) while preserving semantic coherence (CLIP 0.319 vs. 0.316). The conversion is most valuable as a compatibility mechanism rather than a lossless objective replacement.

Second, combined experts achieve higher output diversity (LPIPS 0.782) than single FM (0.752), approaching native DDPM levels (0.787). Heterogeneous objectives create complementary generation patterns.

Third, using the same cosine schedule for both objectives yields marginally better FID than different schedules (32.67 vs. 33.29), suggesting schedule alignment facilitates smoother expert transitions. But both combinations show similar diversity gains, indicating that objective heterogeneity drives the primary benefit.

Routing threshold analysis

For combined DDPM+FM experts, a deterministic router switches between them at a threshold t: DDPM handles timesteps t’ ≤ t (low noise), FM handles t’ > t (high noise).

Impact of router threshold on generation quality. Different thresholds produce a clear quality-diversity trade-off.

Lower thresholds (0.2–0.3) favor quality (FM-dominated denoising, optimal FID). Mid-range thresholds (0.4–0.5) favor diversity (balanced workload, highest LPIPS). Extreme thresholds (0.7) degrade both metrics, confirming that both expert types contribute essential complementary capabilities.

Effects of expert ordering and router thresholds

Expert ordering also matters. In a 2-expert setup (1 converted DDPM + 1 FM), we vary the ordering — DDPM→FM versus FM→DDPM — and the switching threshold τ ∈ {0.3, 0.5, 0.7}.

Expert ordering and router threshold effects. FM→DDPM ordering produces more stable, coherent results, while DDPM→FM shows higher sensitivity to threshold selection.

The results reveal a striking asymmetry. FM→DDPM (bottom row) produces consistently cleaner images across all thresholds. DDPM→FM (top row) degrades at higher thresholds, with blocky artifacts and oversaturation. The reason: when DDPM operates first at high noise (αₜ → 0), the conversion x̂₀ = (xₜ - σₜ ε_θ)/αₜ is numerically unstable, and errors early in the trajectory get baked into the image structure. Letting FM handle high-noise timesteps first avoids this entirely — the converted DDPM expert then refines at low noise where conversion is stable.

Takeaway: DDPM-to-velocity conversion should be restricted to low-noise regimes (t < 0.5).

Discussion

Resource efficiency in context

Our results should be interpreted carefully relative to prior DDM work. The DDM FID range of 5.5–10.5 was achieved at substantially larger training scale (1176 A100-days, 158M images). Our numbers are not directly comparable in absolute FID terms. What they do show is that competitive generation quality is attainable at a fraction of the resources, and that heterogeneous objectives provide an additional quality gain at no extra training cost.

Limitations

We evaluate only a narrow set of DDPM-to-FM ratios (1:7 and 2:6). The ideal allocation likely depends on the data distribution and downstream requirements. The deterministic conversion relies on hand-tuned numerical safeguards; a more robust conversion mechanism that generalizes across arbitrary schedules would strengthen applicability. We consider only ε- and velocity-prediction; extending to x₀-prediction or consistency objectives could further diversify expert specialization but would require generalizing the conversion and routing mechanisms.

What this enables

The practical upshot is that decentralized diffusion training no longer requires coordinated infrastructure or agreement on training objectives. A contributor with a single GPU can train a DDPM expert on portraits. Another can train an FM expert on landscapes using different hardware. These experts combine at inference time without either contributor needing to know what the other was doing.

Citation

@inproceedings{jiang2026heterogeneous,
  title     = {Heterogeneous Decentralized Diffusion Models},
  author    = {Jiang, Zhiying and Seraj, Raihan and Villagra, Marcos and Roy, Bidhan},
  journal = {arXiv preprint arXiv:2603.06741},
  year    = {2026}
}

In Decentralized Diffusion Models (DDMs), denoising is routed through independently trained experts at inference time. These experts can strongly disagree in their denoising predictions. What actually governs the quality of generations in such a system? We present the first ever systematic interpretability study of this question.

The natural expectation is that minimizing denoising trajectory sensitivity — minimizing how perturbations amplify during sampling — should govern generation quality. It doesn’t. Full ensemble routing, which combines all expert predictions at each step, achieves the most stable sampling dynamics and the best numerical convergence. It also produces the worst generation quality (FID 47.9 vs. 22.6 for sparse Top-2 routing). We call this the stability–quality paradox.

Instead, we identify expert-data alignment as the governing principle. Generation quality depends on routing inputs to experts whose training distribution covers the current denoising state, even when doing so makes the trajectory unstable. For DDM deployment, routing should prioritize expert-data alignment over the usually used numerical stability metrics.

Decentralized Diffusion Models

DDMs aren’t Mixture-of-Experts layers. They’re ensembles of independently trained models. In standard MoE architectures, experts are FFN layers within a shared backbone, trained jointly with load balancing losses, and routed at the token level. DDM experts are complete diffusion models, and routing occurs at the input level (entire noisy generations) rather than token level.

Training

Each expert trains in isolation on a disjoint data partition. The training data is partitioned into K clusters (e.g., using k-means on DINOv2 embeddings), and each expert sees only its assigned cluster. So one might train exclusively on landscapes, another on portraits. No shared parameters or gradient communication. Experts only collaborate at inference time.

MoE experts, therefore, operate within a shared representational framework that limits their differences. In contrast, DDM experts are unrestricted and can generate vastly different outputs from the same input.

Routing

At inference time, a lightweight router predicts weights

at each denoising step. The routed velocity field is:

The critical design question is how many experts should contribute at each step. Three natural strategies emerge.

Top-1 commits fully to the single most relevant expert. Every prediction comes from one model. If the router picks poorly, there’s no fallback.

Top-2 blends the two most relevant experts after renormalizing their weights. This allows experts to cross-check each other while still filtering out the majority.

Full Ensemble weights all experts by their router probability. Every expert contributes to every step, giving the mathematically complete combination.

Standard numerical analysis would suggest that including more experts should help. Averaging reduces variance, smooths the velocity field, and stabilizes the ODE integration. We tested whether this intuition holds.

The Stability Paradox

Numerical stability has been the default lens for optimizing diffusion sampling. The foundational probability-flow ODE formulation frames Lipschitz constants and discretization error as determining solver accuracy. Recent work develops stabilized Runge-Kutta methods for stiff diffusion ODEs, studies Lipschitz singularities and their effects on sampling, and the entire solver design space from Euler to Heun to DPM++ is organized around stability-accuracy tradeoffs. These analyses target single-model diffusion, and DDM-specific stability analysis did not exist. We provide the first systematic test of whether this framework transfers to distributed training systems.

We evaluated on Paris, the world’s first publicly released DDM, comprising 8 experts trained on LAION-Aesthetics. We tracked trajectory sensitivity, measuring how strongly the velocity field responds to input perturbations, and step-refinement disagreement, the difference between images generated with N and 2N steps. The results reveal a stability–quality paradox.

Full Ensemble achieves the lowest trajectory sensitivity. It converges the most cleanly. And it produces the worst images, with FID nearly double that of Top-2.

The cumulative IQR measures variability across denoising trajectories. Full Ensemble shows the lowest variability, indicating consistent, well-behaved numerical integration. Top-2 shows higher variability, yet produces superior images. For full experimental details, see the full paper.

Why does averaging more experts — which stabilizes the denoising trajectory degrade image quality?

Expert-Data Alignment Is The Governing Principle

The answer is not in how smooth the path is, but where it leads. Full ensemble averaging reduces Jacobian spectral norms by cancelling variance across expert predictions, which is exactly why it wins on stability metrics. But each expert is trained on a disjoint data cluster. When all experts contribute to every input, most of them are producing velocity predictions for inputs that lie far outside their training distribution. The resulting velocity field is smooth because the out-of-distribution errors partially cancel, but it points toward a weighted average of all cluster centers rather than the data manifold.

In single-model diffusion, a smoother velocity field means cleaner integration means better samples. DDMs break this. The smoothing that lowers trajectory sensitivity is not coming from a better-conditioned ODE. It is coming from averaging contradictory predictions, which suppresses variance at the cost of introducing systematic bias away from any individual data cluster’s learned distribution. Sparse routing avoids this by selecting only the experts whose training data clusters are close to the current input in embedding space, keeping each active expert within its training distribution. The velocity field is noisier but points in the right direction.

We call this governing principle expert-data alignment. Generation quality depends on routing inputs to experts whose training clusters cover the current denoising state, even when doing so reduces the trajectory stability.

If expert-data alignment governs quality, 3 predictions should hold. Sparse routing should achieve higher alignment (selected experts have lower data-cluster distance). Selected experts should produce superior velocity predictions. And expert disagreement should correlate with quality degradation. We tested all three.

Experimental Validation

Cluster Distance Analysis

Does sparse routing actually select experts whose training distribution match the input?

Using the Paris DDM (K=8 experts, DiT-XL/2 architecture with ~606M parameters each), we extracted DINOv2-ViT-L/14 embeddings at timesteps t in {0.3, 0.5, 0.7} during sampling for 500 samples. For each state, we computed the Euclidean distance from the embedding to each of the 8 cluster centroids used during expert training, then ranked the experts by this distance.

The router consistently selects experts whose training distribution are closest to the current denoising state. Top-1 achieves a mean rank of 1.54 (near-optimal selection), while Top-2 maintains strong alignment at 1.96. Full Ensemble, by construction, averages across all experts regardless of relevance.

The router is picky, and that’s the point. Sparse routing filters out experts that are statistically likely to produce poor predictions because they’ve never seen similar data.

Per-Expert Prediction Quality

Do selected experts actually produce better predictions?

For each step in Top-2 generation, we computed velocity vectors from all experts and measured angular deviation from the final blended velocity (which successfully guides the image to completion).

Selected experts consistently produce predictions that align more closely with the successful trajectory. The gap widens with specialization: the MNIST system (each expert trained on a single digit) shows a 43% difference versus 29% for Paris. When experts are highly specialized, the cost of including an inappropriate expert increases. A landscape expert might still contribute useful texture when generating a portrait. A “zero” expert offers actively harmful gradients when drawing a “seven.”

Expert Disagreement Analysis

Does expert disagreement predict quality degradation?

We computed trajectory-integrated disagreement: the average pairwise Euclidean distance between expert velocity predictions, summed over the denoising trajectory. We sorted generated images into quartiles by disagreement and measured LPIPS (perceptual distance to reference).

Images in the high-disagreement quartile (Q4) exhibit worse LPIPS scores than those in the low-disagreement quartile (Q1). The relationship is monotonic.

This is the Full Ensemble failure mode. When experts agree, the average produces reasonable results. When experts disagree (which happens frequently because most experts are out-of-distribution), the average becomes an incoherent compromise. It’s like asking a portrait painter and a landscape painter to collaborate on a cityscape. Sparse routing avoids this by silencing experts that are likely to disagree.

Is Stability Still Useful?

If numerical stability doesn’t govern quality, is it still useful? Yes, but it just measures the wrong thing for generation quality.

Stability Measures Convergence, Not Correctness

Step-refinement disagreement measures whether the solver converges consistently. Full ensemble achieves excellent convergence with step-refinement disagreement approximately 0.020 - doubling the number of sampling steps barely changes the output. Top-2 exhibits more numerical noise with step-refinement disagreement approximately approximately 0.051.

But it’s possible to converge perfectly to a blurry, incoherent average. Stability metrics indicate how easily the solver finds a solution, not whether that solution is good.

Within-Strategy Diagnostics

Trajectory sensitivity may still work as a within-strategy diagnostic. If practitioners are using Top-2 routing, a sudden spike in sensitivity for a specific input might flag a “hard” sample that needs more inference steps, even though sensitivity doesn’t predict quality across routing strategies.

Discussion

Limitations

We validate our hypothesis on two DDM systems (Paris with 8 experts, MNIST with 10 experts). The pattern is consistent across both, but additional systems would strengthen the conclusions - in progress at Bagel Labs. The relationship between cluster distance and prediction quality could be confounded by other factors, such as experts trained on larger clusters being more robust. We control for this by using the same embedding space (DINOv2) that was used during expert training.

Implications

For practitioners building DDM systems:

1. Routing should prioritize alignment over stability. Standard numerical stability metrics don’t indicate system health in DDMs. A “smooth” sampler may simply be averaging away useful signal.

2. Sparse routing is preferable. Top-2 routing achieves a favorable tradeoff: it maintains expert-data alignment while allowing experts to cross-reference each other. Top-1 may be too aggressive; Full Ensemble destroys alignment.

3. Monitor expert-data alignment directly. Track data-cluster distance ranks and expert disagreement during development, not just final FID scores.

Conclusion

Numerical stability doesn’t govern generation quality in Decentralized Diffusion Models. Expert-data alignment does. It means routing inputs to experts trained on similar data.

DDM systems should be evaluated and optimized for alignment rather than stability. Sparse routing succeeds not because it produces stable trajectories, but because it ensures each active expert operates within its domain of competence.

This post presents findings from our research. For full experimental details, methodology, and additional analyses, see:

Expert-Data Alignment Governs Generation Quality in Decentralized Diffusion Models

@misc{villagra_expertdataalignment_2026,
  author       = {Marcos Villagra and Bidhan Roy and Raihan Seraj and Zhiying Jiang},
  title        = {{Expert-Data Alignment Governs Generation Quality in Decentralized Diffusion Models}},
  howpublished = {\url{https://arxiv.org/abs/2602.02685}},
  note         = {arXiv:2602.02685 • accessed DD Mon YYYY},
  year         = {2026}
}

We named it Paris after the city that has always been a refuge for those creating without permission. Two remarkable facts that makes Paris first of it’s kind,

1. It’s a combination of smaller expert diffusion models pre-trained from scratch across different continents in complete isolation. The experts required zero gradient, parameter, or intermediate activation synchronization among each other during training.

2. This zero communication protocol achieves comparable quality to SOTA distributed approaches using 14× less data and 16× less compute.

How? Full technical report and model weights below.

Full Technical Report : https://github.com/bageldotcom/paris/blob/main/paper.pdf
Model Weights : https://huggingface.co/bageldotcom/paris

We believe we can scale this approach to global state-of-the-art results. But that requires solving some more really, really hard problems. If you’re an ML researcher or engineer interested in helping us achieve this while doing the best open-source work of your career, come work with us: jobs.bagel.com.

Bagel Labs launching Tiny Tool Use, an intentionally tiny but production grade open-source library designed to simplify the process of training open-source LLMs for tool use.

Tool-aware LLMs turn text into real-world actions. Which unlocks autonomous decision making for robotics and general infrastructure use.

Tiny Tool Use distills the latest advances in tool-use RL, SFT and evaluation into easy to use templates. Letting teams train and evaluate tool-calling models without extra scaffolding.

It is fully open source : https://github.com/bagel-org/bagel-RL

Tiny Tool Use ships with:

• Interchangeable Training Algorithms – swap SFT, Direct Preference Optimization (DPO), synthetic teacher signals and more with a single config change.

• Configuration‑only workflows – every experiment, tool schema, and hyper‑parameter lives in a JSON file as a result performing training with different configuration is easy.

• First‑class evaluation support – TensorBoard dashboards for training visualization and integration with Berkeley Function Calling Leaderboard scripts.

• Dataset flexibility – plug in real data, generate synthetic traces, or compose both without touching core code.

Training Example Using Qwen3 Models

We now provide an example of using the library to train Qwen3 models. We use Low Rank Adaptation (LoRA) to customize Qwen3 models on ToolBench dataset. The library ships with the example configuration provided in configs/sft_toolbench_config.json which downloads the data, extracts it and uses the processed data for training.

To run the training code with Qwen3—0.6B model, use the following command

python train.py --config configs/sft_toolbench_config.json --output-dir lora_sft_qwen3/

The script will start downloading the ToolBench dataset and unzipping, which will several minutes considering the size of ToolBench.

The above code starts the training procedure, using lora adapters. The configuration file can be edited for full training instead of lora adapters. Furthermore, the configuration of the adapters can also be changed accordingly.

Evaluation and Benchmarking

Beyond its capabilities, the tiny tool use library offers a robust framework for evaluating the general tool-use capabilities of an adapted model. This includes the ability to compare evaluation results directly with established benchmarks, such as the Berkeley Function Calling Leaderboard.

The training statistics can be visualized by running tensorboard with the following command

tensorboard --logdir lora_sft_qwen3/ 

The performance of the adapted model on ToolBench data as training progresses. The evaluation data is displayed in the TensorBoard dashboard for the Qwen3-0.6B model, demonstrating that the Tiny Tool Use library offers clear and interpretable training and evaluation metrics along with improved model capability for function calling.

When the training is complete, the adapters can be merged and saved using the following command

Python save_merge_model.py \
--base_model Qwen/Qwen3-0.6B \
--adapter_path\ lora_sft_qwen3/ \
--output_dir merged_model/ \ 
--trust_remote_code 

A model adapted with a subset of Toolbench data can be obtained from the following link: Qwen3-0.6B-ToolBench

BFCL Leaderboard

The BFCL evaluation on the model can be performed using the following commands, which will generate model response on different test cases.

export BFCL_PROJECT_ROOT=/path/to/your/desired/project/directory

bfcl generate --model Qwen/Qwen3-0.6B --local-model-path merged_model/ \
--test-category simple,parallel,multiple,multi_turn

Finally to obtain the score on the generated model response, the following code is executed, which will save the scores as a csv file.

bfcl evaluate --model Qwen/Qwen3-0.6B \
--test-category simple,parallel,multiple,multi_turn

Bagel Labs team will continue to improve the library to adapt it for broader tool use, with an emphasis on distributed learning algorithms. We welcome contributions, feature requests, and issues on our fully open-source repository: https://github.com/bagel-org/bagel-RL

Can a single number line up every milestone so far and hint at the next leap?

We think maybe. We call that number Return on Experience (RoE). It’s an early phase benchmark.

What you read below is a lab notebook, handed down as working scripture. Feel free to share your peer review, or join the loop (what’s loop? more on that below).

New experience costs more than extra GPU time

When power is cheap, we can keep feeding a language model massive amounts of internet text at little extra cost.

But a surgical robot, a self-driving car, or an ICU-triage policy cannot just “scrape more experience”.

Every new trial burns tire, takes staff time, or waits for legal clearance.

In these settings the scarce resource is fresh experience - each environment step that has to happen in the real world.

Sample-efficiency is the art of squeezing the most learning from the fewest such experiences, and RoE is a meter for that exchange rate.

Recent AI breakthroughs point in an interesting direction

DeepSeek-R1-Zero improves reasoning capability through millions of GRPO self-play prompt-episodes. Meaning, the model practically “argue with itself” millions of times, turning each prompt-response into a tiny lesson, without requiring tons of human text.

OpenAI o3 tool use agent was trained to ask whether the next tool call (browser, code runner, or image) will actually move the answer forward. Each tool use carries an explicit cost during training, so the model invokes a tool only when the expected reward of that experience outweighs the cost.

DreamerV3 learns a world-model first, then practices inside that imagined space, harvesting thousands of virtual trials for every real one.

Different domains, same theme. Progress comes from squeezing more value out of each new experience. The next question is, how much result do we get per unit of experience? That ratio is RoE.

RoE compresses sample-efficiency into a single number that rises roughly logarithmically with progress.

The Formula

Let’s turn that idea into one formal equation.

Return on Experience (RoE) tells us how much “win” an RL agent buys per interaction it pays for. In other words, if you spend one unit of real world experience, how far does your performance score climb?

One step, one score, divided by weighted cost. Where,

headline score - the score a paper or a model brags about - Atari score, math-test accuracy, etc.

count - how many times the RL agent interacted with something while learning.

weight - Not all interactions are equal. Weight adds a rough price tag for each kind of interaction because some are costly and some are cheap.

• real-world action = 1 (full price)

• calling an external tool (e.g. OpenAI o3) ≈ 0.01

• self-play or a step inside a learned simulator (e.g. DeepSeek-R1-zero) ≈ 0.001

We think of the weights as discounts. Because some synthetic/simulated steps cost a thousandth of a real-world step, tool calls cost about two hundredth.

This way, RoE lets us put a language model, a robot arm, and a video-game agent on the same chart. High RoE means the agent climbs the leaderboard quickly and keeps the cost for fresh experience low. Low RoE means it wastes a lot of real-world practice for each bump in score.

A short history of RoE Loops

Below we cast some historical milestones of Reinforcement Learning into the RoE benchmark, to show RoE could predict the performance of them.

We call each big RoE jump a loop - one full turn of experience → insight → back to model/agent. Like the closed-loop theorem (or a bagel ring).

But first, some info on the classic benchmarks used to calculate RoE,

• Atari - 57 classic video games used as a standard obstacle course for RL.

  • Human-normalised score takes a human play-through on each game, sets that to 1.0, then shows how far the agent climbs above it.

  • Median is the middle game after ranking all 57 scores; a quick “overall” health check.

  • Atari-80k / Atari-100k mean looks at the average score but limits the agent to only 80k or 100k game frames, so we can see who learns fastest, not just who learns best.

• GSM8K - 8.5k grade-school word problems. For each problem an RL agent writes one answer. Pass@1 is simply the percentage of problems it gets right on that very first try.

Here is a worked example of RoE calculation with OpenAI o3,

For o3, headline_score = 0.982, weight = 0.01, interaction count = 2M.

So, according to our RoE equation above,

Observations,

• Rainbow quadrupled RoE over DQN by upgrading loss functions and replay, proving that better optimization alone can unlock large efficiency from same amount of experiences.

• EfficientZero learned to plan inside a latent world model instead of in the real world. With most interactions being synthetic, RoE passed 10.

• DreamerV3 pushed the same idea further. Over 90% of the experiences were simulated. Real-world performance climbed with very little new real-world experience.

• DeepSeek-R1 brought model self play into the limelight. Model debates with itself on each problem, boosting GSM8K accuracy while keeping the new and cost heavy experience ledger small.

• Cost-aware tool use has been another RL win of 2025. OpenAI o3 trained the RL agent to accrue a cost while calling a tool. That optimized RoE to ~50.

Prediction for the next 5 years

RoE has grown on a log curve - 0.023 for Rainbow (2018) up to about 50 for OpenAI o3 (2025). There’s no sign that the slope is flattening. World-model research, sample-efficient policy optimization and a coming surge in cheap compute all point the same way.

Below is a 5 year prediction on RoE numbers, and examples of what that could mean, inspired from peer-reviewed academic work,

(We count the past jumps - Rainbow through o3 - as Loops 1-5.)

Loop 6 - Year 2026 - RoE ≈ 150

World-model research will slip from Atari into the operating room. Hospitals will train surgical robots almost fully in world-model simulators such as Surgical Gym and ORBIT-Surgical. A bot will practise thousands of virtual stitches for before doing a real one.

Loop 7 - Year 2027 - RoE ≈ 1000

Scaling labs will run foundation world-models such as Genie on ~30x cheaper compute hardware. Agents will write and test entire NeurIPS papers inside those generated 3D sandboxes, requiring only high-level edits from humans. RoE will move past the four-digit mark.

Loop 8 - Year 2028 - RoE ≈ 3000

Self-play on synthetic molecule sets will reduce wet-lab screening from tens of thousands of assays to a few hundred. Drug-discovery will be massively accelerated. RoE will keep its log-pace and enter “several thousand” territory.

Loop 9 - Year 2029 - RoE ≈ 8000

Regulators will approve the first cargo aircraft whose flight-control laws are proved largely in simulation-based certification workflows. One hour of real flight data will be enough to sign off.

Loop 10 - Year 2030 - RoE ≈ 12000

Companies will field “digital graduates” - ~30B parameter reasoning models trained with Dr. GRPO style length-penalised objectives. Each will master a new technical field from a few annotated pages, then draft patents, close mergers and rewrite tax law overnight. One human auditor will review the output of a thousand such agents.

But, it is 2025. The year of AGI. And, we have made it possible.

Today, we’re open sourcing a frontier research, ZKLoRA. A zero knowledge protocol that allows verification of LoRA fine-tuning of open source AI models, in 1-2 seconds.

And not only for toy models, but for current state of the art open source models like llama 3.3 etc. with tens or hundreds of billions of parameters.

Want to try it yourself? Here’s the code : https://github.com/bagel-org/ZKLoRA

Want to see the benchmarks and curious how it works? Read the full research paper here : https://github.com/bagel-org/ZKLoRA/blob/main/paper.pdf

ZKLoRA Paper

331KB ∙ PDF file

Download

Download

We are Bagel Labs, a distributed machine learning research lab.

We believe ZKLoRA kickstarts a new era for verifiable model training across untrusted networks.

We have had radical technological advancements in recent history. Social media, augmented reality, platform shifts like web, mobile. But AI is way more significant of a technology. It is as significant as the discovery of fire. It has the potential to change the trajectory of the evolution of our species.

One of the holy grails of unlocking this potential of AI is to build systems that can reason like humans. By improving AI's, Large Language Models in particular, ability to break down complex problems and apply logical steps.

Bagel's research team has been exploring this problem. Analyzing LLM building techniques, especially fine-tuning techniques, to allow Large Language Models to evolve from pattern-recognizing prediction agents to true cognitive agents. Our deep research spanned three major types of reasoning, aka intelligence: arithmetic, commonsense, and symbolic.

Today, we're sharing our findings. This research targets the core of what we believe to be the ultimate AI evolution, human-level reasoning. Or beyond (God level?).

We have explored techniques for the training and fine-tuning phases of model development. We have also ventured into the absolutely fascinating world of inference-time reasoning. This is where LLMs can be built or fine-tuned to generate novel solutions during inference, even if the solutions aren't part of their training dataset.

Dive in. And if you're in a rush, we have a TLDR at the end.

Types of Reasoning

Varied types of reasoning tasks stretch AI's abilities. First, let's understand how they're defined.

Arithmetic reasoning pushes machine learning to test problem-solving in a clear way. It forces models to break down problems. Choose from many strategies. Connect steps to find solutions. This makes math different. It shows exactly how well models can grasp details. And use the right solution steps in order.

Commonsense reasoning upends our expectations. Models must understand the strange logic of everyday life. The challenges emerge when systems face the quirks of human interactions. The implicit rules we take for granted. For example, a door opens before you walk through. Time flows forward not backward. Water makes things wet. These obvious truths become complex puzzles for artificial systems to unravel.

Symbolic reasoning flips the script on traditional machine learning. While neural networks excel at fuzzy pattern matching, symbols demand precision. Models must follow strict rules. Manipulate abstract concepts. Chain logical steps. Like a careful mathematician rather than an intuitive artist. The symbols hold no inherent meaning. Yet through them, we build towers of logic that reach toward human-level reasoning.

Beyond these core types, reasoning takes many forms. Logical deduction draws rigid conclusions while induction makes creative leaps. Causal reasoning traces the hidden threads between actions and consequences. Multimodal reasoning juggles text, images, and data in a complex combination of understanding. Knowledge graphs map the relationship of facts and relationships. Yet all serve one goal - moving AI from pattern matching toward true comprehension. From memorized responses to novel insights. From prediction to understanding.

Below, we look into training time and inference time approaches to enhance these types of reasoning.

1. Training Time Approaches

1.1. Fine-Tuning Approaches

Parameter Efficient Fine-Tuning (PEFT)

How it works: PEFT reverses traditional model adaptation (Hu et al. 2023). Four methods reveal new techniques.

Prompt-based learning embeds adjustable signals into frozen models. Prefix-tuning and P-tuning introduce small changes. These changes alter outputs without altering the main model.

Reparametrization methods like LoRA simplify complex weight matrices. They turn large updates into efficient low-rank forms. LoRA captures patterns from high-dimensional spaces with minimal adjustment.

Adapters create extra neural pathways. Series adapters stack, each layer adjusting outputs gradually. Parallel adapters develop side skills, keeping the base intact.

Adapter placement is key. Series adapters fit after MLP layers. Parallel adapters excel within them. LoRA touches both attention and MLP layers. Each method targets the right spot.

Why it's useful: PEFT reduces resource demands. Large models gain new abilities without major changes. PEFT preserves the base while adding specialized skills. Hardware that struggled with fine-tuning now handles complex updates.

Tradeoffs: Not all tasks fit PEFT. Some models need deeper changes. Base model limitations still exist. Combining methods is tricky. PEFT may struggle with very complex tasks.

WizardMath

How it works: WizardMath learns in three distinct steps (Luo et al., 2023).

First is supervised fine-tuning. Here, the model picks up raw mathematical patterns. It starts recognizing basic structures. Patterns get mapped to solutions. This step builds intuition for common operations. The foundation is set.

Next, instruction reward models refine the process. They judge both answers and methods. These models look for efficiency. They guide the model toward elegant solutions. The focus shifts from correctness to quality.

Finally, PPO-based reinforcement learning enhances problem-solving. The model tests ideas, adapts, and improves. Evol-Instruct feedback loops refine its logic with each run (Xu et al. 2023). It gets better at selecting strategies.

Why it's Useful: Most models just match patterns. WizardMath thinks in logical steps. It breaks down problems like a mathematician. It selects methods based on understanding, not memory. This leads to solutions that are both effective and precise.

Tradeoffs: Training WizardMath takes heavy computational resources. Its deep math focus limits general use. Low-quality data can introduce errors. Practical solutions can sometimes lose to elegant ones.

Divergent Chain of Thought (DCoT)

How it works: DCoT breaks the single-path approach (Puerto et al. 2024). Multiple paths form at once. Each one tackles the problem differently. Yet all conclude in a single inference.

Zero-shot generation creates diverse solutions. Every path seeks the truth. Each follows its own course. Some are direct. Others are more complex. All valid. The model acts like a group of experts. Each path offers a different view.

These paths then interact. Strong strategies merge. Weak points become clear. The model learns to assess its own reasoning. It compares methods. It blends insights. All this happens without extra training.

Why it's Useful: Multiple paths offer built-in validation. When paths align, certainty rises. When they don’t, issues appear. Different views reveal hidden details. Diversity deepens understanding.

Tradeoffs: More paths need more computation. Balancing variety and consistency is tricky. Conflicting paths need resolving. For simple tasks, it's overkill. A group isn't always better than one.

1.2. Pre-training and Knowledge Transfer

Continued Pre-training

How it works: Models like Galactica (Taylor et al. 2022) and MINERVA (Lewkowycz et al. 2022) go beyond standard training. They learn from over 100 billion tokens of scientific data. This includes mathematical papers, scientific articles, and technical documentation. Raw data is converted into structured knowledge.

Galactica includes tokens for specific scientific terms. It treats citations as part of the vocabulary. Chemical formulas become meaningful. Mathematical symbols are treated like tools. It learns the language of science.

MINERVA focuses on quantitative reasoning. It answers natural language questions in physics, chemistry, and economics. Converts questions into math formulas. Uses LaTeX to present detailed solutions. It performs the calculations on its own.

Why it's useful: Smaller models can surpass larger ones in specific fields. They grasp complex math. Work with technical notation naturally. The gap between general models and experts shrinks.

Tradeoffs: Training costs rise. Each field requires massive new data. As new knowledge grows, old knowledge fades. Balancing focus and breadth is hard. It might be great at physics but weak in other areas.

Curriculum Learning

How it Works: Learning transforms from random sampling to structured progression (Adyasha & Maharana 2022). Like evolution, but guided. Deliberate. Purposeful.

A teacher network ranks training samples. Easy concepts come first. Complex ideas build on simple ones. The pacing function controls the flow of knowledge. Sometimes fixed. Sometimes adaptive. Responds to the model's growing understanding.

Three methods measure sample difficulty. Question Answering Probability tracks how often the model succeeds. Model Variability watches for consistent responses. Energy-based scoring identifies outliers and edge cases. The curriculum adapts based on these signals.

Why it's Useful: Models learn more efficiently. They build strong foundations before tackling complexity. Understanding grows naturally. Organically. Each concept reinforces the last. Difficult ideas become manageable when approached in sequence.

Tradeoffs: Designing effective curricula challenges even experts. Learning time stretches longer. Some concepts resist ordered progression. The path from simple to complex isn't always clear. Sometimes chaos teaches better than order.

CoT Knowledge Distillation

How it Works: Large models become teachers. Small models become students. Knowledge transfers through carefully curated examples (Magister et al. 2023).

The process splits into two phases. First, generate Chain of Thought data. Large models solve problems step by step. Show their work. Create a roadmap of reasoning. Only correct solutions make the cut. Quality matters more than quantity.

Then comes student fine-tuning. Small models learn from these examples. They see not just answers but thinking processes. The target answer guides early steps. This prevents small errors from derailing entire solutions. Teacher forcing ensures the student stays on track.

Why it's Useful: Advanced reasoning becomes accessible to smaller models. Complex problem-solving skills transfer efficiently. Small models learn to think clearly with limited resources. They gain the wisdom of larger models without the computational burden.

Tradeoffs: Some sophistication gets lost in translation. Students never quite match their teachers. The distillation process demands careful curation. Bad examples can teach bad habits. The balance between compression and comprehension remains delicate.

Join Bagel Community

2. Inference Time Approaches

2.1. Chain-Based Methods

Chain of Thought (CoT)

How it works: Wei et al. redefined reasoning with their 2022 paper. They introduced language models to step-by-step problem-solving using just eight examples. These guided models unlock hidden potential.

With precise prompts, models show their internal reasoning. No need for new training or changes to the model. This latent capability is accessed by using strategic examples.

The models learn to break down problems into logical steps that mimic human thinking. Each step becomes clear. The internal thought process shifts from a black box to a visible sequence.

This approach scaled well. PaLM, with chain-of-thought prompting, hit 75.6% on StrategyQA. Even sports questions saw 95.4% accuracy, surpassing human experts. Complex math problems were solved with clear, step-by-step reasoning. In commonsense tasks, hidden assumptions surfaced in natural language. Symbolic problems became easy to follow.

Why it's useful: Wei et al.'s work showed breakthroughs across fields. LaMDA 137B demonstrated this by generating 96% correct answers with sound reasoning. Problem-solving became transparent. Larger models produced more coherent explanations.

Tradeoffs: Reasoning sometimes fails. Models can get confused. Wei’s research showed that 46% of wrong answers had minor mistakes, while 54% had major logical errors. Sequential reasoning can hit barriers. Complex tasks push models to their limits.

Program of Thought (PoT)

How it works: Chen et al.'s 2022 work changed how models approach math. They turned natural language into executable programs that solve complex problems with machine-level precision.

The process is seamless. Word problems convert directly into Python code. Variables capture key details from the text. Functions embody solution strategies. Algorithms emerge from simple descriptions. The model coordinates external tools with precision.

PoT set new records, improving math benchmarks by 8% in few-shot settings. In zero-shot, the gains were 12%. The code tells a story with structured logic. Control flows mirror human thought. Programs serve as both solution and explanation.

PAL expanded on this. Gao et al. in 2023 showed how models could use Python interpreters for better reasoning. Complex calculations became sharper. Formal math operations translated into natural expression.

Why it's useful: Precision dominates. Math problems flow into code. The model combines high-level reasoning with computational accuracy. It's like a mathematician working alongside a supercomputer.

Tradeoffs: Some problems don't translate well into code. Executing programs raises security concerns. The model must handle both natural language and code, increasing the risk of errors.

2.2. Consistency and Verification Methods

Self-Consistency (SC)

How it works: Wang et al. introduced SC in 2022, shifting from greedy decoding to statistical sampling. This method changes inference entirely.

Instead of one solution, each step produces multiple paths. SC explores various reasoning attempts at once. The decoder samples different trajectories in the probability space. Errors are reduced by repeating steps, leading to validation through sampling.

SC’s statistical foundation is strong. It marginalizes over samples to minimize errors in individual paths. Think of it like quantum mechanics: multiple paths exist, and truth emerges from the statistical patterns.

Their approach was groundbreaking. The decoder generates n unique reasoning chains, each following a different probability path. Final answers come from majority voting, but the process goes beyond simple counts.

Wang's team tested models from UL2-20B to PaLM-540B. Accuracy increased across the board. Smaller models showed the most improvement, indicating SC unlocks hidden potential in models of all sizes.

Why it's useful: Numbers tell the story. Multiple paths automatically validate answers. Different paths catch edge cases. Robustness increases as more paths are explored. Quantity becomes quality.

Tradeoffs: Computation grows costly. Each path demands resources. Memory use spikes. Contradictory paths sometimes arise. Resolving these conflicts adds complexity.

Self-Endorsement (SE)

How it works: Wang et al.’s 2024 paper introduced SE, a new verification method. The system generates diverse responses and then analyzes them. Facts are extracted, labeled, and compared. Cross-response validation assigns endorsement scores to each fact.

SE uses advanced fact extraction algorithms. Neural retrieval identifies key claims, and automatic cross-referencing helps the model distinguish strong facts from weaker ones. This statistical validation process drives the system.

High-scoring facts shape future outputs, while low-scoring ones lead to re-evaluation. Each pass refines the model’s response through consistency.

The fact extraction pipeline is highly technical. Named entity recognition identifies key elements, and relation extraction maps connections. All of this occurs without human input.

Why it's useful: Accuracy improves. Hallucinations drop. The system validates its own facts. Confidence scores make responses more reliable.

Tradeoffs: Processing takes longer. Fact extraction sometimes fails. Complex statements resist simple validation. Some valid facts get rejected if they don’t fit the statistical pattern.

Least-to-Most Prompting (LM)

How it works: Zhou et al. introduced LM in 2022, a system that breaks tasks into smaller parts and solves them step by step.

The process follows phases. First, the model analyzes the input. Next, it identifies sub-tasks. Then, it solves each part. Finally, it combines the results. Each phase builds on the previous one.

For example, in the last-letter task with "cat dog bird," the model processes each word separately. It finds ‘t’ from "cat," ‘g’ from "dog," and ‘d’ from "bird." Then, it combines them into "tgd." The model achieved 94% accuracy with four words and 74% even with twelve.

Errors are predictable. Sometimes letters drop during connection. Sometimes extras appear. But it rarely confuses the final letter of each word.

Why it's useful: LM is highly efficient. It only needs two examples to work well. It uses less tokens than traditional methods, achieving equal or better results.

Scaling is impressive. The model handles sequences four times longer than its training examples without losing accuracy. Standard methods fail on long sequences, scoring 31.8% on twelve-word tests. LM hits 74%, with a growing advantage on harder tasks.

Tradeoffs: Some tasks don't split easily. Certain problems need a different approach. The method requires more steps, which adds processing time.

Technical limits arise. The model must track partial solutions, and memory usage grows with longer sequences. Some tasks need several attempts to find the best split.

Careful planning is essential. The order of sub-tasks affects accuracy, and managing information efficiently becomes critical. The system must adapt its splitting strategy for different problems.

How to Test for Reasoning

Cognitive sciences have studied human reasoning since experimental psychology emerged in the late 19th century. This field has been crucial for technological development, education improvement, cognitive disorder treatment, and better decision-making. Scientists use various tools to study reasoning, including problem-solving tasks, computational models, brain imaging (fMRI and EEG), and behavioral measurements like eye-tracking. These combined methods help researchers understand how humans reason.

Analogously, AI researchers have invented reasoning tasks to test the reasoning capabilities of LLMs in the form of special datasets. Being AI more of an engineering field and computer science field, these datasets provide a rigorous benchmark to test AI systems. This allows researchers to measure a model’s accuracy and identify areas where it may be falling short.

Datasets for testing AI reasoning on one type of reasoning should be diverse around that type of reasoning in order to test various complexities and nuances in tasks. For example, to evaluate language models' common sense capabilities, a dataset like ARC is used. The figure below we shows a ranking of the best LLMs for the ARC-challenge dataset taken from different sources.

Inference-time techniques appear in green, training-time techniques appear in orange, and standard base models appear in blue.

In the image above, the best performing techniques correspond to inference-time approaches, in particular, SC has a clear advantage over standard CoT. The fine-tuning approaches cannot match the inference-time approaches where they show a clear advantage.

TLDR

Our research focuses on strengthening reasoning in Large Language Models (LLMs) in three ways. First, arithmetic reasoning - approaching math problems logically. Next, commonsense reasoning - grasping everyday situations and drawing conclusions. Finally, symbolic reasoning - handling abstract symbols by strict logic.

We explore two strategies to push these areas forward. First are training-time methods. These adjust AI’s learning process, adjusting it for specific tasks but needing time and computing power. For example, WizardMath teaches detailed problem-solving for math, while PEFT (Parameter Efficient Fine-Tuning) builds skills without huge resources. DCoT (Divergent Chain of Thought) allows AI to consider multiple solutions simultaneously.

The second approach is Inference-time methods. These enhance existing models without retraining, bringing quick improvements, though sometimes with less depth. Chain of Thought (CoT) prompts AI to explain each step it takes. Program of Thought (PoT) has AI write and run code to boost accuracy. Self-Consistency (SC) checks multiple paths to ensure reliable answers.

The below table is a summary of our findings.

By open-sourcing our research on AI reasoning, our team at Bagel aims to collaborate with the Open Source AI community to forge humanity's next chapter.

It's the absolute cutting edge of AI model safety and data privacy.

Microsoft researchers showcased this capability by removing Harry Potter references from META’s Llama2 Model. OpenAI researchers applied unlearning to eliminate biased content, enhancing model safety.

Unlearning addresses GDPR's "right to be forgotten" mandate. For example, Google held a machine unlearning competition focusing on removing specific facial images from an AI model trained to predict age from images.

For cybersecurity, unlearning mitigates data poisoning risks. JPMorgan Chase applied unlearning to consumer systems, preserving customer privacy without compromising effectiveness.

The future of scalable AI systems lie in selective forgetting, not indiscriminate data accumulation. This breakthrough ushers in a new paradigm of AI privacy, where machines learn, unlearn, and relearn with unprecedented flexibility.

As we stand on this new frontier, we have researched different machine unlearning techniques, that will lead us to a world where machines possess human-like memory malleability.

And if you're in a rush, we have a TLDR at the end.

What is Machine Unlearning

Machine unlearning erases specific data from pre-trained models without full retraining. This process requires embedding unlearning mechanisms during initial model training.

We have broken down the techniques into four major ones: Exact, Approximate, Prompt-Based, Decentralized.

1. Exact methods remove specific data precisely from AI models. They divide data or transform algorithms to target unwanted information. This approach deletes or updates only necessary parts, avoiding full model retrains.

2. Approximate methods take a broader approach. They fine-tune models on slightly altered data to reduce targeted information's influence. Some add controlled noise during training, limiting each data point's impact from the start.

3. Prompt-Based techniques guide AI language models to "forget" without changing their core. They use smart word tricks to create examples that steer the AI away from unwanted info. The AI receives new instructions to ignore certain knowledge, without erasing its memory.

4. Decentralized approaches erase data from AI systems spread across multiple devices. They remove information when a participant exits the network. These methods use compact models to efficiently spread the impact of this exit. This eliminates the leaving party's influence without retraining the entire system. The process maintains collaborative learning while enabling selective data removal.

As these techniques evolve, machine unlearning becomes more effective, enabling models to forget specific information while maintaining strong performance.

1. Exact Methods for Machine Unlearning

Exact unlearning removes specific data from machine learning models. It targets a subset called the "forget set" within the main dataset. The process creates a new model that performs nearly identically to the original. This technique deletes chosen data points rapidly, without the need of full retraining.

The updated model maintains high accuracy on remaining data. Think of it as selectively erasing memories from a human brain while preserving overall knowledge and functionality.

We have broken down the exact methods in two major types.

• SISA.

• Statistical Query.

1.1. SISA

SISA (Sharding, Isolation, Slicing, and Aggregation) is an advanced algorithmic framework designed to accelerate machine unlearning processes (Bourtoule et al. 2019). This technique strategically limits the influence of individual data points during model training, facilitating efficient removal of specific data without necessitating complete model retraining.

SISA Architecture

The diagram shows a layered machine learning system. Training data forms the base, split into 'n' shards (Shard 1 to Shard n). Each shard is an isolated subset of the dataset. Shards divide into 'R' slices (slice 1 to slice R), enabling incremental training and unlearning.

Models 1 through n, shown as connected nodes, train on corresponding shards. These models constitute the SISA (Sharded, Isolated, Sliced, Aggregated) system core. This structure allows parallel processing and targeted updates.

An aggregation layer tops the system. It combines outputs from all models to generate the final prediction. This method leverages insights from different data subsets.

The design enhances data management and system flexibility. It processes large datasets efficiently while allowing precise adjustments. The architecture balances scalability with fine-tuned control.

SISA Training

The SISA training process operates systematically. The main dataset splits into 'n' shards, each training a separate model. For each shard:

1. The first slice initializes a new model, training for a preset number of epochs.

2. Subsequent slices continue training from the model's previous state.

3. The last slice finalizes training, producing the shard's constituent model.

This slice-wise approach enables incremental learning. Models update efficiently as new data arrives. It also facilitates data removal without full retraining. The process repeats for all shards.

The final step assembles the complete SISA model by combining all trained constituent models. This structure allows for parallel processing and targeted updates.

SISA Inference

SISA inference runs input through all models (Model 1 to n). Each model produces a prediction. The system aggregates these predictions via majority voting. It selects the most common label as output.

This method reduces individual model biases. The aggregation layer functions as an ensemble. It combines insights from all data shards.

Voting can be weighted for more reliable models. This allows inference fine-tuning. The system can track confidence by measuring model agreement.

This enables parallel processing during inference and can reduce latency in large-scale applications.

SISA Unlearning Procedure

SISA unlearning locates the target data to be deleted, in a specific shard and slice, shown as the red box in Shard 2, slice 2. The system then jumps back to a point just before this data was added. This is shown in the green outline around slice 2 in Shard 2.

From this past state, SISA begins selective retraining. Model 2 Retrain, in green, shows this process. It retrains using all data except the unlearn target.

Once complete, the retrained Model 2 replaces its original version. The Aggregation layer at the top then updates to use this new model's input.

This selective retraining approach allows SISA to precisely remove specific data. It does this without disturbing other parts of the system, as Model 1 and Model n remain unchanged.

SISA Advantages

1. Computational Efficiency: SISA significantly reduces the computational overhead of unlearning by limiting retraining to a single shard and subset of slices.

2. Scalability: The sharding approach allows for parallel processing and efficient handling of large datasets.

3. Isolation: By confining data points to specific shards and slices, SISA minimizes the impact of individual data removals on the overall model.

4. Adaptability: The framework can be extended to various machine learning algorithms and data types.

SISA-based Methods

Several adaptations of SISA have been developed for specific use cases.

• DaRE Forest: A random forest variant using a two-level approach with random and greedy nodes, enabling efficient updates of subtrees upon data removal (Brophy & Lowd 2021).

• HedgeCut: Focuses on low-latency unlearning in extremely randomized trees, introducing split robustness concepts (Schelter et al. 2021).

• GraphEraser: Tailored for Graph Neural Networks (GNNs), incorporating balanced graph partition and optimized shard model importance scoring (Chen et al. 2021).

• RecEraser: Specialized for recommendation tasks, employing adaptive aggregation methods to combine sub-model predictions (Chen et al. 2022).

In conclusion, SISA represents a flexible and efficient framework for managing data removal requests in machine learning models. Its architecture and process flow, as illustrated in the provided images, demonstrate a sophisticated approach to balancing model performance with data privacy and right-to-be-forgotten requirements.

1.2 Statistical Query (SQ) Approach

Cao & Yang (2015) introduced an innovative approach to machine unlearning, distinct from the SISA framework. This method uses statistical queries and summation forms to enable efficient unlearning without complete model retraining.

SQ Core Concepts

1. Summation Form Transformation: The key idea is transforming learning algorithms into a summation form. In this form, the algorithm relies on a small number of summations, each representing an efficiently computable transformation of the training data samples. These summations are saved alongside the trained model.

2. Statistical Query (SQ) Learning: To achieve the summation form, Cao & Yang employ Statistical Query (SQ) learning, introduced by Kearns (1993) as a refinement of the PAC learning model. In SQ learning, the algorithm interacts with a statistical query oracle instead of directly examining individual examples. The learning algorithm queries the oracle with a function and a tolerance parameter, receiving an estimate of the expected value over a distribution of labeled examples.

3. Formal Definition: More formally, the learning algorithm queries an oracle with a function ξ:X x {0,1} → {0,1} and a tolerance parameter τ. The oracle responds with an estimate of the expected value of ξ over a distribution D of labeled examples such that the estimate of E[ξ(x,y)] is within an additive τ. The training algorithm only has access to these statistics; it does not directly access the input data.

How SQ Works

The image below illustrates the process flow.

The image shows a sophisticated approach to machine learning and unlearning using Statistical Query (SQ) learning. Here's a technical breakdown:

• Data Points: The left side shows input data points x₁, x₂, x₃. These represent individual training examples from the dataset D.

• Statistical Queries: Dashed lines connect data points to q₁ and q₂, representing statistical query functions ξ:X x {0,1} → {0,1}. These queries transform raw data into relevant statistics.

• Summations: The Σ symbols denote summation operators. They aggregate query results across the dataset, computing E[ξ(x,y)] for each query function.

• Tolerance: While not explicitly shown, each summation has an associated tolerance parameter τ, defining the acceptable error margin for the computed statistic.

• Neural Network: The orange box on the right depicts a neural network architecture. It contains multiple layers (Input, Hidden, Output) interconnected by edges.

• SQ to Model Pipeline: Solid arrows from summations to the neural network illustrate how aggregated statistics directly inform model parameters or training.

• Unlearning Process:

  • To forget a data point, the system updates stored summations by subtracting that point's contributions.

  • This involves simple arithmetic operations on the summations.

  • The model is then efficiently recomputed using these updated summations.

• Unlearning Efficiency: This structure allows for efficient unlearning. Removing a data point only requires updating the summations, not retraining the entire model.

• Scalability: The approach scales well with large datasets, as the model depends on a fixed number of summations rather than individual data points.

This SQ learning framework transforms complex learning algorithms into a summation form, enabling rapid updates and unlearning while maintaining model accuracy.

SQ Technical Considerations

The effectiveness of this method depends on expressing the learning algorithm in terms of statistical queries. The choice of queries and summations can significantly impact unlearning efficiency and model performance. There's a balance to strike between the granularity of summations and the efficiency of unlearning.

This non-SISA approach provides a novel perspective on machine unlearning, focusing on data representation and model dependency rather than architectural modifications. Its application in various machine learning contexts remains an active area of research and development.

2. Approximate Methods for Machine Unlearning

While exact methods for machine unlearning are effective, they often require significant computational resources and storage overhead. Xu et al. (2024) present approximate methods as a more efficient alternative, especially for large datasets and complex models.

General Approach (Xu et al. 2024)

The process, as illustrated in the image, consists of four key steps:

1. Influence Computation: The process begins with the "Dataset" shown on the left side of the image. From this dataset, we identify "data to forget" - specific training examples we want the model to unlearn. The "Influence" step analyzes how these data points impact the model's predictions. This analysis employs influence functions, a technique from robust statistics. For a given data point z, its influence I(z) is calculated as:

Where:

• H_θ is the Hessian matrix of the loss function L,

• θ represents the model parameters,

• ∇_θ L(z, θ) is the gradient of the loss with respect to θ.

Computing this for large models can be challenging due to the Hessian inverse. Practical implementations often use approximation techniques like the conjugate gradient method to estimate H_θ^-1∇_θL efficiently.

2. Model Parameter Adjustment: The image shows a transition from an initial neural network to an adjusted one. This step modifies the model parameters to counteract the influence of the data to be forgotten. The adjustment typically follows the direction opposite to the computed influence:

where ε is a small step size. In practice, this update may be applied iteratively or combined with optimization techniques like Adam or RMSprop for stability.

3. Noise Addition: The "noise" element above the neural network in the image represents the introduction of controlled noise to the model. This step implements differential privacy techniques to prevent inference of recently removed data. The noise is calibrated based on the sensitivity of the unlearning operation. For Gaussian noise, we might add noise drawn from N(0, σ²) to each parameter, where σ is calculated as:

Here, c is a constant, Δf is the sensitivity of the unlearning operation, and ε is the desired privacy parameter. The choice between Gaussian and Laplacian noise depends on the specific privacy guarantees required.

4. Updated Model Validation: The final neural network in the image, labeled "New model with influence minimized" represents the outcome of the unlearning process. This step assesses the performance of the new model to ensure effective unlearning while maintaining overall functionality.

Validation employs metrics such as:

1. Accuracy: (TP + TN) / (TP + TN + FP + FN)

2. F1 Score: 2 * (Precision * Recall) / (Precision + Recall)

3. AUC-ROC: Area under the Receiver Operating Characteristic curve

These metrics are compared against both the original model and an unlearned model.

Additionally, specific tests may be conducted to ensure the removed data cannot be inferred, such as membership inference attacks or analyzing the model's uncertainty on the forgotten examples.

This process provides a balance between unlearning efficiency and effectiveness, making it applicable to a wide range of machine learning models and datasets.

Below, we will go in-depth into two approximate methods of machine unlearning: Certified Removal and Unlearning in Spiking Neural Networks.

2.1 Certified Removal

Guo et al. (2019) unveiled certified removal, a big step in approximate machine unlearning. This method offers theoretical safeguards against adversarial extraction of removed training data. The image provided illustrates the key components and process flow of certified removal.

How Certified Removal Works

Certified removal operates by strategically adjusting the model to "forget" specific data points. Here's a simplified breakdown of the process.

1. Noise Injection: During the initial training, a controlled amount of noise is added to the model. This noise acts as a form of protection, making it harder for an adversary to extract information about individual data points.

2. Influence Calculation: For each piece of data in the forget set, the method calculates its "influence" on the model. This influence represents how much that specific data point contributed to the model's current state.

3. Parameter Adjustment: Using the calculated influence, the model's internal parameters are carefully adjusted. This adjustment effectively cancels out the impact of the forgotten data points.

The result is a new model that behaves as if it had never seen the forgotten data in the first place.

Why Certified Removal Matters

Guo et al. (2019) demonstrated the efficiency of certified removal compared to full retraining. In their experiments with a linear model trained on the MNIST dataset:

• Certified removal took 0.04 seconds to remove a data point.

• Full retraining of the model took 15.6 seconds.

These results highlight the significant time savings offered by certified removal, especially when dealing with large datasets or frequent removal requests.

2.2 Unlearning in Spiking Neural Networks

Spiking Neural Networks (SNNs) mimic the behavior of biological neurons in the human brain (Maass 1997). As shown in the figure below, SNNs consist of interconnected neuron-like units that communicate through a processing network.

In SNNs, each neuron has a membrane potential that changes over time. When this potential reaches a specific threshold, the neuron "fires" or "spikes," sending a signal along its axon (represented by orange lines in the image). This spiking behavior is a key characteristic that distinguishes SNNs from traditional artificial neural networks.

Wang et al. (2023) introduced a technique for unlearning in SNNs, addressing the need for privacy mechanisms and unlearning techniques in these biology inspired models. SNNs have shown success in pattern recognition, particularly in speech (Wang et al. 2023) and image recognition (Su et al. 2023). They have also demonstrated potential in medical applications, such as constructing stimulation systems for Parkinson's patients (Geiger 2023).

How SNN Unlearning works

It happens in three phases, which can be understood in the context of the image above.

1. Selective Retraining: This phase identifies neurons (the branching structures) and synapses (the connections between neurons, represented by arrows in the central network) responsible for the information to be forgotten. It estimates the correlation between a neuron's spike train (the pattern of signals sent along the orange lines) and the targeted data. Synapses are selected based on their weight change due to learning the targeted data. The weights of these synapses are then adjusted using a modified learning rule.

2. Synaptic Pruning: This step aims to remove synapses whose weight change surpasses a given threshold. In the context of the image, this would involve selectively removing some of the neuron connections in the interconnected network. All synapses in the network are verified and removed if necessary, effectively eliminating traces of the targeted data.

3. Adaptive Thresholding: In this phase, neuron firing thresholds are dynamically modified based on their activity in relation to the targeted data. This would affect how easily the neurons "fire" or send signals along the orange lines. This reduces the neuron's response to stimuli linked with the data to be unlearned.

The image shows input neurons on the left (Input 1, Input 2), hidden layers in the center (H1₁, H1₂, H1₃, H2₁, H2₂, H2₃), and output neurons on the right (Output 1, Output 2). The unlearning process would involve modifying the connections and behaviors of these components to "forget" specific information.

Why SNN Unlearning matters

Wang et al. (2023) tested the unlearning effectiveness on two datasets: UCI HAR and MNIST. The results showed:

• A decrease in performance metrics (accuracy, precision, recall, etc.) for both datasets after unlearning. This indicates that the network successfully "forgot" the targeted information.

• The UCI HAR dataset exhibited a bigger drop in performance compared to MNIST, indicating dataset-specific resilience to the unlearning process. This suggests that the complexity of the data and the way it's encoded in the network can affect the unlearning process.

• Retraining after unlearning recovered a considerable portion of the lost performance, especially for the MNIST dataset. This demonstrates the network's ability to relearn and adapt after the unlearning process.

• A clear trade-off emerged between the percentage of samples unlearned and accuracy loss. As more samples were removed, accuracy decreased. This highlights the balance between preserving overall performance and effectively removing specific information.

This method focuses on minimizing the impact of neurons and synapses highly correlated with the forget set. By targeting only these elements, as represented by the specific connections and nodes in the central network of the image, the approach avoids full retraining. This makes it an approximate method for unlearning in SNNs, offering a balance between effective information removal and computational efficiency.

2.3 Other Approximate Methods

The field of approximate machine unlearning continues to evolve beyond the work of Guo et al. (2019). Several researchers have proposed innovative approaches.

Sekhari et al. (2021) introduced an algorithm utilizing cheap-to-store data statistics. This method is great for convex loss functions and can unlearn a significant number of samples without full dataset access.

Suriyakumar & Wilson (2022) developed an online unlearning algorithm based on the infinitesimal jackknife method. Their approach reduces computational overhead by inverting the Hessian matrix only once.

Mehta et al. (2022) proposed a parameter selection technique using conditional independence tests. Their L-CODEC and L-FOCI algorithms identify relevant model parameters for unlearning, avoiding full Hessian matrix inversion.

Wu et al. (2022) presented the Performance Unchanged Model Augmentation (PUMA) method. PUMA removes unique characteristics of marked data points while preserving model performance through influence function calculations.

Tanno et al. (2022) introduced a "predictive approach" for identifying causes of model failures in medical imaging. They adapted the Elastic Weight Consolidation method to compute training example influence on failure sets.

Warnecke et al. (2021) proposed a framework for unlearning features and labels instead of entire data points. Their method leverages influence functions to perform closed-form updates on model parameters.

These advancements demonstrate the ongoing efforts to improve machine unlearning efficiency and applicability across various domains and model types.

Subscribe now

3. Prompt-Based Unlearning

Prompt-based unlearning techniques aim to make language models "forget" specific information without directly modifying their parameters. These methods are particularly relevant for SOTA large language models (LLMs) where direct access to model parameters is often restricted.

It's important to stress that prompt-based methods do not facilitate machine unlearning as per our earlier definition, they provide a way to "pretend" to forget information.

How Prompt-Based methods work

Guardrails

Thaker et al. (2024) explored the effectiveness of simple guardrail approaches for unlearning in LLMs.

1. Prompting: Crafting specific prompts to guide the model's behavior.

2. Input/Output Filtering: Screening inputs and outputs to prevent unwanted information.

These techniques were tested on three different benchmarks:

• Who's Harry Potter?

• TOFU (TOpic-based Forgetting Using model Update)

• WMDP (Wikipedia Movie Distance Plot)

In-Context Unlearning (ICUL)

Inspired by in-context learning (Brown et al. 2020), ICUL (Pawelczyk et al., 2024) involves constructing a specific context in a prompt that includes both correctly labeled and mislabeled examples. The process involves three steps.

1. Relabeling forget points: Choose a number of data points you want the model to forget. For each of these points, change its original label to a different, incorrect label. This creates a list of items, each with its content and a new, incorrect label.

2. Adding correct examples: Select a number of correctly labeled examples from your dataset. Add these examples to the list created in step 1. Now you have a longer list that includes both the relabeled "forget" points and some correctly labeled examples.

3. Creating the final prompt: Take the list from step 2 and add your actual query or question at the end. When you submit this to the language model, set the temperature to 0, which makes the model's responses more deterministic (less random).

Why Prompt-Based methods matter

Guardrails (Thaker et al., 2024)

• Achieved unlearning performance comparable to more complex fine-tuning approaches.

• Offered a simple, resource-efficient approach to unlearning in LLMs.

In-Context Unlearning (ICUL) (Pawelczyk et al., 2024)

• Demonstrated performance equal to or better than some leading unlearning methods that require access to model parameters.

• Effectively eliminated the influence of a training point on a model's output in text classification and question-answering tasks.

• Used significantly less memory compared to traditional unlearning methods such as gradient ascent.

Both methods address the challenge of unlearning in black-box LLMs without requiring direct access to model parameters. This is crucial given the widespread use of LLMs in various professional and personal contexts, where avoiding outputs like hate speech, toxic behavior, or hallucinations due to data poisoning is essential.

3. Decentralized Machine Unlearning

Decentralized machine unlearning addresses the challenge of removing specific information from AI systems geographically distributed across interconnected devices.

The HDUS (Heterogeneous Decentralized Unlearning framework with Seed model distillation) method tackles this complex issue when knowledge disseminates through a network of collaborating machines (Ye et al. 2023).

This approach is vital in a connected world where decentralized systems like edge computing expand, and data privacy laws demand deletion rights.

How HDUS Works

HDUS uses a dual-model approach for greater precision and efficiency.

1. Main Model: A high-capacity neural network trained on local data, functioning as the core intelligence.

2. Seed Model: A lightweight version that distills insights from the main model, enabling peer collaboration.

The seed model, trained on a reference (or synthetic) data sample, mimics the main model’s behavior without revealing sensitive data. In collaborative training, peers share seed model parameters. For inference, each peer generates output using its main model and an ensemble of neighbors' seed models.

The above image shows how client a₁ processes data using its main model, extracts key insights through a seed model, and integrates inputs from neighboring clients (b₁, b₂, …, bₖ) to generate a final output. The diagram highlights the complex connections between processing units, storage, and data paths, reflecting the intricate nature of the HDUS framework.

When a device exits, it triggers unlearning by notifying its neighbors. They remove its contribution from their submodels, ensuring its influence is erased without needing full model retraining.

Why HDUS Matters

HDUS offers critical advantages by leveraging decentralized machine unlearning:

1. Efficiency: Fast unlearning without full model retraining, using lightweight seed models for updates.

2. Compatibility: Works across diverse device networks, supporting varied model architectures.

3. Complete removal: Ensures total erasure of a departing device's contribution, ensuring privacy.

4. Exact unlearning: HDUS achieves precise data removal in a decentralized setting, maintaining unlearning integrity across distributed nodes.

HDUS overcomes the challenge of propagating unlearning requests in networks where knowledge has already spread across peers. It maintains integrity across dynamic network topologies, vital for edge computing and federated learning environments.

The HDUS approach removes data associated with specific devices, not individual data points. This shift is crucial for privacy management in distributed AI systems, especially in environments like mobile edge computing and IoT, where devices frequently join or leave.

Unlike other approaches, HDUS focuses on removing a party's entire contribution, advancing privacy and data rights in decentralized AI.

Join Hack-AI-Thon by Bagel

TLDR

Machine unlearning transforms AI privacy by erasing data from trained models without full retraining. Companies like META and JPMorgan Chase leverage this innovation.

Unlearning specific data points is often exponentially faster than full retraining or distillation. Cutting computational power and storage needs. It removes targeted information, preserves model performance, and enables selective forgetting without losing other critical knowledge.

At Bagel, we fuse breakthrough technology like machine unlearning into an evolving Machine Learning ecosystem, redefining privacy and collaboration in AI.

Four major methods drive machine unlearning: Exact (SISA, SQ), Approximate, Prompt-Based, and Decentralized (HDUS). Each offers unique approaches to efficient data removal, with out full model retraining.

Method comparison:

Machine unlearning enables AI to adapt, correct, and protect privacy at scale.

As a leader in privacy-preserving machine learning (PPML), Bagel's lab has been advancing the field significantly. Building upon our research of PPML approaches using privacy-enhancing technologies (PETs) like differential privacy, federated learning, ZKML, and FHE, we have deeply investigated two additional technologies gaining adoption: trusted execution environments (TEEs) and secure multiparty computation (MPC).

See here for the part 1 of this series.

In this article, through a blend of theoretical insights and real-world examples, we will demonstrate how TEEs and MPCs can be harnessed to build robust, secure, and scalable PPML solutions while navigating the complex regulatory landscape of AI safety and privacy. We will discuss how these techniques can be utilized for your use case, compare them to the discussed privacy technologies in our previous research post, and analyze their pros and cons.

And if you're in a rush, we have a TLDR at the end.

Trusted Execution Environments

Trusted Execution Environments (TEEs) are technologies for hardware-assisted confidential computing. TEEs enable the execution of isolated and verifiable code inside protected memory, also known as enclaves or secure worlds. TEEs ensure that code and data within enclaves are protected from external threats, including malicious software and unauthorized access. This isolation is critical for maintaining the integrity of machine learning models, especially when deployed in untrusted environments such as edge devices or cloud platforms. By using TEEs, developers can safeguard the intellectual property embedded in their models and be sure that the models produce reliable and tamper-proof results.

General View of TEE Components

In a TEE, the user that wants to use the enclave first establishes a secure connection (attestation) using a key generated by a secure processor. The user can verify the integrity of the enclave through an attestation process that is certified by the manufacturer. If the attestation is successful, the user can trust the enclave to execute code and store data securely.

The overall security also depends on the underlying hardware security features, such as memory encryption and isolation. Hardware-based isolation mechanisms ensure that the code and data inside the enclave are protected from any unauthorized access that resides outside the enclave. The chip containing the enclave is designed to resist tampering, and any hardware tampering will make the enclave non-functional.

Software Attestation Process

1. Manufacturer Root Key: The manufacturer has a unique root key embedded in the hardware.

2. Attestation Keys: Each device processor has an embedded attestation key (private and public) from the manufacturer.

3. Endorsement Certificate: The manufacturer certifies the public attestation key, indicating it is bound to a tamper-resistant hardware chip.

4. Verifier Interaction: The verifier sends a message to the TEE, which includes the required attestation.

5. Attestation Response: The TEE responds with the attestation key and the signed message.

6. Verification: The verifier checks the attestation by verifying the hash (key+message) with the attestation key. If accepted, the verifier trusts the attestation.

Chain of Trust in Software Attestation

The measurement in the context of TEEs is a cryptographic hash produced from the code and data inside the enclave. The measurement is included in the attestation signature to prove that the enclave was not tampered with. The verifier can then check the measurement against a known good value to ensure the enclave's integrity. Finally, the TEE produces the attestation signature containing the measurement plus other attestation data and sends it back to the verifier. The verifier can then use this attestation signature to verify the integrity of the enclave. In some cases, the verifier shares a secret key that can be used to establish a secure communication channel between the TEE and the verifier.

Main Hardware-Assisted TEEs

• Intel SGX

• AMD TrustZone

• Arm TrustZone

• Apple Secure Enclave

• NVIDIA Backed Authentication

Attack Vectors

Despite the robust security guarantees provided by Trusted Execution Environments (TEEs) through isolation and attestation from a trusted manufacturer, research has shown that TEEs are still vulnerable to several types of attacks:

1. Side-Channel Attacks: Side-channel attacks exploit indirect information leakage from the TEE, such as timing, power consumption, electromagnetic emissions, or memory access patterns. Examples include:

   

   1. Timing Side-Channel Attacks: These attacks measure the time taken to execute certain operations to infer sensitive information. Timing fluctuations caused by operations like multiplication and division can be exploited to obtain encryption keys. Studies have shown that timing attacks are one of the most damaging side-channel attacks, particularly in cryptographic implementations. Research by Fei et al. (2021) provides detailed insights into these vulnerabilities.

   2. Memory Side-Channel Attacks: These attacks monitor memory access patterns to deduce the data being processed. Memory-based side-channel attacks observe events on shared resources in the memory hierarchy, such as cache hits and misses, to infer secret-dependent memory access patterns. Research has demonstrated the effectiveness of these attacks on both CPUs and GPUs. Studies by Sasy et al. (2017) highlight the risks associated with these attacks.

   3. Network Side-Channel Attacks: These attacks analyze network traffic patterns to extract confidential information. Network side-channel attacks exploit the correlation between network traffic and the operations being performed within the TEE. Studies have highlighted the risks associated with these attacks in various networked environments. Research by Costan and Devadas (2016) discusses these vulnerabilities in detail.

2. Replay Attacks: Replay attacks involve intercepting and retransmitting valid data to create unauthorized effects. For instance, an attacker might replace a newer message with an older one from the same sender, effectively resetting the state of a computation. Research by Bellare et al. (2000) and Mo et al. (2023) provide comprehensive studies on these attacks.

3. Host-Based Attacks: Host-based attacks exploit vulnerabilities arising from the interactions between the host operating system (OS), user-space processes, and the TEEs running on the same system. Examples include:

   1. Data Poisoning: Malicious data is injected into the system to corrupt the training process of machine learning models. Research by Mo et al. (2023) and Liu et al. (2021) provide detailed analyses of these attacks.

   2. Adversarial Examples: Inputs are crafted to deceive machine learning models into making incorrect predictions. Studies by Mo et al. (2023) and Liu et al. (2021) discuss these vulnerabilities.

4. Access Pattern Exploitation: Previous research has demonstrated methods to exploit access patterns to TEEs to classify encrypted inputs with high accuracy. This type of attack can reveal sensitive information about the data being processed within the TEE. Grover et al. (2018) and Zhang et al. (2021) provide detailed insights into these vulnerabilities.

These vulnerabilities highlight the need for continuous research and development to get the security of TEEs ready for untrusted peer-to-peer networks. Below, we discuss how TEEs can be used for privacy-preserving machine learning (PPML), especially for inference or model training.

Inference in TEEs

In order to perform inference inside TEEs, the data providers and the model provider first need to prepare their corresponding dataset and model. They encrypt this information, to keep it safe during transit.

Next, a remote attestation ceremony takes place. The data provider, model provider, and the TEE engage in a key exchange process described above. They verify each other's identities, exchange keys and establish a secure communication channel.

With trust established, the encrypted dataset and model are transferred into the TEE. Inside the TEE they are decrypted.

Now, the inference begins. The TEE executes the inference process, analyzing the data using the model while ensuring secure isolation. The inference results, computed labels, are encrypted before being returned to the data provider. The model, having served its purpose, can be discarded. The figure below shows the entire inference process.

Some implementations of this ML inference process include:

• Occlumency. A system that leverages Intel SGX to preserve the confidentiality and integrity of user data throughout the entire deep learning inference process (Lee et al. 2019).

• Branchy-TEE. A framework that dynamically loads the inference network into the TEE on-demand, based on an early-exit mechanism to break the hardware performance bottleneck of the TEE (Wang et al. 2019).

• Origami. A system that provides privacy-preserving inference for large deep neural network models through a combination of enclave cryptographic blinding and accelerator-based computation (Giri Nara et al. 2019).

Training in TEEs

In order to train a model, a data provider needs to load its dataset and an initial model describing the architecture of the network that is going to be trained.

Similarly as for inference, first we encrypt the dataset and the initial model. Then, remote attestation and key exchange takes place between the data provider and the TEE and a secure communications channel is created. Through this channel, the dataset, the initial model and the training algorithm are loaded into the TEE. Then, the TEE executes the training algorithm isolated from any external influence. After a given number of epochs, the weights are returned to the data owner through the secure channel.

Some implementations of ML training are:

• TensorSCONE. A secure TensorFlow framework using Intel SGX, which enables the training and usage of TensorFlow models within a TEE (Kunkel et al. 2019).

• Graphcore IPUs. Graphcore's IPU Trusted Extensions (ITX) provide a TEE for AI accelerators, ensuring strong confidentiality and integrity guarantees with lower performance overheads for training (Vaswani et al. 2022).

• Citadel. This is an ML system that protects both data and model privacy using Intel SGX. Citadel performs distributed training across multiple training enclaves for the data owner, and an aggregator enclave for the model owner. It uses Zero-sum masking for secure aggregation between training enclaves, which is a technique where data owners collectively generate masks and apply them to their individual updates before sending them to the aggregator (Bonawitz et al. 2017). The authors of Citadel show that increasing the number of training enclaves from 1 to 32 results in increased throughput of 4.7X - 19.6X (Zhang et al. 2021).

Pros of TEEs

The advantages of using TEE with neural networks are:

• Confidential Execution. The ML inference and training process is executed within the TEE. The TEE ensures that the data and model are protected from any external access or tampering during the inference process.

• Isolation. The TEE isolates the inference and training process from the rest of the system, ensuring that even if the host system is compromised, the data and model remain secure.

Cons of TEEs

The disadvantages of using TEE with neural networks are:

• Requires trust on the manufacturer. In the attestation process, a certificate of the public key from the manufacture is required who works as a certificate authority. In a peer-to-peer network setting, trust is an undesirable property (Mo et al. 2023).

• Restricted resources. TEEs offer restricted computation resources inside their secure enclave. This forces ML developers to find clever ways to implement ML tasks, for example, efficient partition of ML processes (cf. Mo et al. (2021); Liu et al. (2021)).

• ML inside TEEs are still begin attacked. The ML process is still vulnerable. Previous research conducted by Grover et al. (2018) demonstrated methods to exploit access patterns to TEEs in order to classify encrypted inputs with high accuracy.

• Vulnerable interactions with the outside. This refers to a type of security attack known as "Host-based Attacks" in the context of TEEs. These attacks exploit vulnerabilities arising from the interactions between the host operating system (OS), user-space processes, and the TEEs running on the same system. It includes data poisoning and adversarial examples (Mo et al. 2023).

• SDK support. Most TEEs only provide basic low-level SDK, and therefore, it is hard to port all ML dependencies into TEEs. Hence, porting an ML code still requires a lot of code refactoring (Mo et al. 2023).

• Large TCB size due to libOS use. The Trusted Computing Base (TCB) is the set of hardware and software components that are critical to the security of a TEE system. A library operating system (libOS) is a lightweight operating system that provides application-level abstractions without the need for a full operating system kernel. A libOS can be used to easily port ML applications into a TEE, but such port results in a huge TCB size, which can lead to a reduced security (Mo et al. 2023).

Secure Multiparty Computation (MPC)

Secure Multiparty Computation (MPC) is a cryptographic technique that allows multiple parties to jointly compute a function over their private inputs without revealing those inputs to each other. The key idea behind MPC is to divide the computation into smaller steps, where each party performs a portion of the computation (Lindell 2020).

How It Works

The process usually involves the following steps:

1. Function Agreement: All parties agree on the function f to be computed, as shown in the image: f(x1, x2, ..., xn).

2. Input Sharing: Each party provides its private input (x1, x2, x3, x4, x5 in the image) but keeps it hidden from others.

3. Protocol Execution: The parties interact through a carefully designed protocol to compute the output. This protocol ensures that no party learns anything about the other parties' inputs beyond what is revealed by the output y = f(x1, x2, ..., xn).

4. Output Computation: The final output y is computed and revealed to all parties.

The key properties of MPC protocols are:

• Privacy: No information about any party's input x is revealed beyond what the output y inherently reveals.

• Correctness: Each party is guaranteed that the output is correctly computed from the inputs according to the function f, even if some parties deviate from the protocol.

• Independence of inputs: Parties choose inputs freely.

• Guaranteed output delivery: All parties get the output.

• Fairness: Bad actors can't gain an edge (Lindell 2020).

Privacy and correctness are two basic properties that any MPC protocol should support. For stronger security, protocols should also ensure independence of inputs, guaranteed output delivery, and fairness to all parties.

Common Schemes in MPC

• Secret Sharing: One common technique used in MPC is secret sharing. Discovered independently by Shamir (1979) and Blakley (1979). In secret sharing, a secret is divided into multiple shares, and each share is distributed to different parties. No single share reveals any information about the secret, but when combined, the shares can reconstruct the secret. Shamir's Secret Sharing is a well-known method where a secret is represented as a polynomial, and shares are points on this polynomial. The secret can be reconstructed using Lagrange interpolation.

• Garbled Circuits: Another technique is Garbled Circuits, discovered by Yao (1986), which is used for two-party computations. In this method, one party (the garbler) creates an encrypted version of the function (the garbled circuit) and sends it to the other party (the evaluator). The evaluator uses its input and the garbled circuit to compute the output without learning the other party's input. This method ensures that the computation is secure and private.

• Oblivious Transfer: Oblivious transfer is a fundamental building block in many MPC protocols. It allows a sender to send one of many possible messages to a receiver, but the sender does not know which message was received. This technique is crucial for ensuring that parties do not learn more information than they are supposed to.

Common Applications

• Data Analytics: MPC enables joint data analysis while keeping inputs private. This is useful in healthcare, finance, and marketing where privacy is crucial. Companies can perform secure statistical analysis on combined datasets without exposing individual records (Zhou et al. 2019).

• Auctions and Biddings: MPC can conduct auctions where bids remain private, only revealing the winning bid price. This enhances fairness and privacy in auctions (Bogetoft et al. 2006).

• Fraud Detection: By securely analyzing combined transaction data from multiple sources using MPC, fraud detection algorithms can operate on a broader dataset while preserving individual transaction privacy (Nielsen 2021).

Below, we will review how to use MPC to perform inference and training in machine learning in a privacy preserving manner using secret-sharing. The main idea is that data and model are shared and computations occur on shares, not raw data.

MPC in Inference

In order to explain how inference is done using MPC, we need to introduce some basics on secret sharing. We present here a definition by Evans et al. (2017). A (t,n)-secret sharing scheme splits a secret s into n shares in a way that t-1 shares reveal no information about the secret s, while t or more shares allow reconstruction of s. In a two-party secret sharing scheme we have t=n=2. In the following we will explain a simple approach for inference given by Li et al. (2022) in the context of transformer models.

Similarly to zkML and FHE, where all operations of a neural network must be translated into arithmetic operations, we use a secret sharing scheme for computing additions and multiplications. Let us consider two parties, a data provider with an data input x, and a model provider with a model input y consisting of all weights in the model.

To compute additions, the data provider splits x into two shares x1 and x2 with x=x1+x2 and the model provider splits y into y1 and y2 with y=y1+y2. The data provider keeps x1 and y1 and the model provider keeps x2 and y2. This way, neither can reconstruct the original inputs x and y. Then the data provider computes D=x1+y1 and the model provider computes M=x2+y2. The reconstruction process then results in D+M=x+y.

To compute multiplications, the parties can use a secure protocol based on Beaver triples. Both parties compute the shares x1,x2,y1,y2 and the data provider gets x1,y1 and the model provider gets x2,y2. Now a Beaver triple is generated c=ab using oblivious transfer or homomorphic encryption and a,b,c are secret shared with the data provider receiving a1,b1,c1 and the model provider receiving a2,b2,c2. Now the data provider computes ϵ1=x1-a1 and δ1=y1-b1 and the model provider computes ϵ2=x2-a2 and δ2=y2-b2. They communicate these numbers and reconstruct ϵ=ϵ1+ϵ2 and δ=δ1+δ2. Then the data provider can compute r1=c1+ϵ·b1+δ·a1+δ·ϵ and the model provider can compute r2=c2+ϵ·b2+δ·a2+δ·ϵ. The multiplication result is then given by xy=r1+r2, which is an addition operation that can be jointly computed using the protocol of the previous paragraph.

The implementation of more complex operations, like comparisons, can benefit from 3-party MPC as shown by Dong et al. (2023).

MPC in Training

For the case of training with MPC we assume we have two or more data providers that would like to keep their inputs secret and they would like to jointly compute the weights of a model. A secret-sharing scheme was used by Liu et al. (2021) where they showed how to construct an efficient n-party protocol for secure neural network training that can provide security for all honest participants even when a majority of the parties are malicious.

Similarly as in the case for inference, all operations in the architecture of the neural network must be performed using arithmetic operations, that is, additions and multiplications.

In the case of a 2-party protocol, for each private value x, the data provider must split it into shares x1,x2 and it gives x2 to the other data provider. Similarly as in the inference case, a Beaver triple is computed and shared between each data provider. In general, for n data providers, each input is partitioned into n shares.

At the end of training, a model is learned and shared between the data provider and the model trainer.

Pros of MPC

The advantages of using MPC with neural networks are:

• Joint privacy-preserving data analysis. MPC allows multiple parties to collaboratively analyze data without revealing their individual datasets. This is particularly valuable in industries like healthcare, finance, and marketing, where sensitive data must be protected.

• Secure Machine Learning Model Training. MPC can be used to train machine learning models on data from multiple sources without exposing the individual datasets. This is crucial in scenarios where data owners are reluctant to share their data due to privacy concerns or legal restrictions.

Cons of MPC

The disadvantages of using MPC with neural networks are:

• Trust. While MPC protocols aim to provide trustless computation, their deployment in real-world systems introduces additional trust assumptions and practical considerations (Lindell et al. 2023).

• Complexity. Developing and deploying MPC solutions can be complex. It requires expertise in cryptography, secure protocol design, and distributed systems. Ensuring the correctness and security of the MPC protocol itself is essential, as any vulnerabilities could compromise the privacy of the data being processed (Liu et al. 2021).

TLDR

In Part 1, we discussed the four main privacy-preserving machine learning techniques: differential privacy (DP), zero-knowledge machine learning (ZKML), federated learning (FL), and fully homomorphic encryption (FHE). We assessed these techniques based on data privacy, model algorithm privacy, model weights privacy, and verifiability.

In this Part 2, we looked into two other popular privacy-enhancing technologies: trusted execution environments (TEEs) and secure multiparty computation (MPC).

TEEs allow for isolated and verifiable code execution in protected memory, safeguarding code and data from external threats. However, TEEs require trust in the manufacturer and introduce a single point of failure, which is not suitable for trustless and peer-to-peer (P2P) networks.

MPC enables multiple parties to compute on shared data while keeping inputs private. It facilitates joint privacy-preserving data analysis and secure machine learning model training without exposing individual datasets. However, their deployment in real-world systems introduces trust assumptions.

The table below summarizes our findings - see part 1 for the definitions of each property.

Consider privacy. Industries like finance, drug discovery, medical imaging can work around strict data sharing restrictions through synthetic data. It mimics real world data without exposing personal information, allowing for secure, regulation compliant AI development.

Training AI for rare, extreme scenarios is also vital, but often impossible with real data. Synthetic data solves this problem. It can simulate events like uncommon car accidents to train autonomous vehicle. This vastly expands the scope and robustness of AI systems.

Another persistent issue is bias in datasets. Real data often reflects historical prejudices or demographic imbalances. Synthetic data enables the creation of balanced datasets that counteract these biases, leading to fairer AI systems.

AI development needs rapid access to vast amount of quality data. Real data collection and processing has overhead. Synthetic data can be generated instantly, supporting fast testing and iteration with diverse datasets. It improves model accuracy while cutting the need for data labeling and collection drastically.

This article takes you on a journey into the fascinating world of synthetic data. We'll uncover the top techniques used to create it. We'll see how it's solving critical problems across sectors. And we'll give you an unfiltered look at the pros and cons you need to know. Or how you can implement it for your specific business need.

If you're in a rush, we have a TLDR at the end.

How To Generate Synthetic Data

Synthetic data is statistically accurate simulations of real-world datasets. Two forms of it exist: fully synthetic data, created without real data, and partially synthetic data, which may include original dataset elements.

Data quality depends on generation technique and type. In this article, we have explored three major useful synthetic data generation technique families Generative, Evolutionary, and Marginal-based. Technique choice depends on data type (tabular, time-series, images) and desired properties (diversity, realism, privacy).

Generative Methods

Generative AI has had a big impact across all industries pushing teams to perform at a higher rate. The goal of generative AI is essentially to generate new data, be it image, audio, video, text, etc. It is, thus, natural to rely on generative AI techniques to create new synthetic data.

Below we discuss three of the most popular generative synthetic data generation methods below, Variational Autoencoders (VAEs), Generative Adversarial Networks (GANs) and Diffusion Models.

1. Variational Autoencoders (VAEs)

A Variational Autoencoder (VAE) is an algorithmic tool that compresses and decompresses data. VAEs are practical tools used in diverse fields such as computational chemistry for generating molecular graphs, anomaly detection in data streams, and even in game design to create diverse and complex environments.

The goal of VAEs is to transform data points from a high-dimensional space, to a lower-dimensional, more meaningful latent space. The transformation process involves two main components: the encoder and the decoder. The encoder's task is to map data to the latent space, while the decoder reconstructs the data back to the original high-dimensional space. The goal is to minimize the error between the original and reconstructed data, denoted as

where x is an input, E and D denote the encoder and decoder functions, and ϵ is an error function.

A VAE differs from a standard autoencoder by introducing a probabilistic twist. The encoder in a VAE doesn't directly output coordinates in the latent space. Instead, it outputs parameters to a probability distribution—typically a Gaussian denoted by

where μ denotes the mean and σ it’s standard deviation where both can depend on the input x. The figure below shows the general architecture of a VAE.

The decoder samples from the distribution generated by the encoder to generate new data points, effectively turning the decoder into a generator of synthetic data. This process is governed by a loss function L, which balances two terms

where KL denotes the Kullback-Leibler divergence, a measure of how one probability distribution diverges from a second, expected probability distribution.

Pros

• Data Type Flexibility. VAEs are adaptable to various data types, including multimodal data.

• Anomaly Detection. It excels in identifying data points that deviate significantly from the norm, as these will have high reconstruction errors.

• Efficient Sampling. Post-training, sampling from the latent space and generating new data via the decoder is straightforward and efficient.

Cons

• Loss Function Balancing. The dual nature of the loss function, combining reconstruction and regularization, requires careful tuning, which is still an active area of research (Asperti et al. 2021).

• Blurriness in Output. Often, the output images or data from VAEs lack sharpness (Asperti 2020).

• Clustering in Latent Space. In scenarios where the dataset inherently contains subcategories, these often manifest as clusters in the latent space, which can be problematic for generative tasks where uniformity might be desired (Asperti et al. 2021).

2. Generative Adversarial Networks (GANs)

In a Generative Adversarial Networks (GANs), two AI algorithms are locked in a game competing against each other. We have two neural networks, the generator (G) and the discriminator (D), engaging in a continuous game of deception and detection. The generator strives to create data so authentic that the discriminator cannot distinguish it from real data. Conversely, the discriminator's goal is to accurately identify whether the data it reviews is genuine or fabricated by the generator. This dynamic competition drives both networks towards perfection, simulating a game where the prize is the ability to replicate reality.

GANs have found their way into various industries showcasing their versatility. NVIDIA, for instance, uses GANs to create lifelike artwork and to transform 2D images into 3D shapes. eBay employs them for image-based search functionalities in their marketplace. Audi leverages GANs for innovative wheel design, while Zalando generates new textures for the fashion industry.

In a GAN, the generator crafts fake data, aiming to pass it off as real. The discriminator tries to separate the true from the false. This process is a high-level game, where each player sharpens their skills in response to the other's moves. The generator uses input 𝑥 from a distribution p​ corresponding to a random variable X to produce output aiming to mimic a target distribution T. The discriminator examines outputs 𝑦=𝐺(𝑥), guessing if they're real or fake. Ideally, the generator's output distribution 𝐺(𝑋) aligns so closely with 𝑇 that the discriminator is left guessing, assigning a 50% probability to both real and fake data.

Pros

• Anomaly Detection. GANs excel in identifying outliers, making them invaluable for tasks like fraud detection.

• High-Quality Outputs. They are capable of producing images so realistic that they can fool the human eye.

• Privacy Protection. Incorporating differential privacy during training ensures sensitive data remains secure.

Cons

• Mode Collapse. A scenario where the generator produces limited varieties of output, making it easier for the discriminator to identify fakes.

• Training Challenges. Balancing the training pace of both networks is crucial, yet difficult, to ensure both improve simultaneously.

• Non-Convergence. The unique loss function of GANs can lead to challenges in achieving convergence through gradient descent.

3. Diffusion Models

Diffusion models represent an alternative generative AI approach, primarily used for creating synthetic images. They appear in the automatic generation of high-resolution images and videos where notable implementations include Stability AI's Stable Diffusion, OpenAI's Dall-E and Sora, and Meta's ventures into 3D image generation. These models operate through a dual-process system: a forward diffusion that adds noise, and a reverse diffusion that reconstructs the image. This method offers a robust framework based on principles of non-equilibrium thermodynamics.

The forward diffusion is a Markov chain process. It incrementally introduces Gaussian noise into an image, transforming it step-by-step into a noisy version. This process is mathematically represented as:

The mean of the normal distribution is

and the variance is

The sequence β1,…,βt, known as a schedule, varies between 0 and 1.

The reverse diffusion employs a deep neural network, specifically a U-Net architecture, to reconstruct a new image from the noisy data. The goal is to learn a probability distribution

where θ represents the model parameters. Training this model involves maximizing the likelihood of matching the learned probability distributions of the equation above with the original distributions q.

Pros

• Theoretical foundations. These models are grounded in the solid theoretical frameworks of thermodynamics and differential equations.

• Efficient inference. They are capable of generating high-resolution images efficiently, even on personal computing devices.

• High-quality outputs. The images produced are detailed and realistic, making them central to technologies like Dall-E and Sora.

Cons

• Training time. Developing a diffusion model from scratch demands significant computational resources and expertise.

• Complexity. A deep understanding of the underlying theoretical principles is necessary to effectively implement and optimize these models (Karras et al. (2024).

Evolutionary Methods

Evolutionary methods comprise a collection of algorithms and techniques that iteratively construct synthetic data from a seed dataset. By applying certain operations of combination and mutation of data points, evolutionary methods look to generate synthetic data that is diverse and deep in the context of the initial dataset.

1. Genetic Algorithms

The algorithm Private-GSD (Private Genetic Synthetic Data) proposed by Liu et al. (2023) is a genetic algorithm designed to generate synthetic data that approximates the statistical properties of an underlying sensitive dataset while ensuring differential privacy. It leverages principles from biological evolution to iteratively optimize a population of synthetic datasets. By applying selection pressure and introducing variation through mutation and recombination, these methods evolve datasets that can produce increasingly diverse and high-quality synthetic examples over generations.

The process of Private-GSD is shown in the figure below.

1. Initialization: Private-GSD starts with an initial population of synthetic datasets randomly generated. Each dataset in this population is a candidate solution to the problem of generating synthetic data that closely matches the statistical properties of the original dataset.

2. Evaluation: Each candidate synthetic dataset is evaluated based on how well it approximates the statistical queries of interest on the original dataset. This evaluation is quantified using a loss function that measures the difference between the statistical properties of the synthetic and original datasets.

3. Selection: The algorithm selects a subset of the best-performing synthetic datasets from the current population. These datasets are considered "elite" and are carried over to the next generation.

4. Crossover and Mutation: Private-GSD generates new synthetic datasets by combining features from the elite datasets (crossover) and introducing random changes (mutations) to some of these datasets. This process is inspired by genetic recombination and mutation in biological evolution.

5. Iteration: Steps 2 through 4 are repeated for a specified number of generations or until a convergence criterion is met. With each iteration, the population of synthetic datasets evolves, ideally becoming better approximations of the original dataset's statistical properties.

6. Output: The algorithm outputs the best-performing synthetic dataset from the final generation.

Pros

• Flexibility. Private-GSD can generate synthetic data that approximates a wide range of statistical queries, including those that are non-differentiable and thus challenging for other algorithms that rely on gradient-based optimization.

• Privacy Preservation. By design, Private-GSD ensures differential privacy, making it suitable for generating synthetic data from sensitive datasets without compromising individual privacy.

• No Requirement for Differentiability. Unlike methods that rely on gradient-based optimization, Private-GSD does not require the objective function to be differentiable. This allows it to work with a broader range of statistical queries.

Cons

• Computational Complexity. The iterative nature of genetic algorithms, combined with the need to evaluate multiple candidate solutions in each generation, can make Private-GSD computationally intensive, especially for large datasets and complex statistical queries (Liu et al. 2023).

• Parameter Sensitivity. The performance of Private-GSD can be sensitive to its hyperparameters, such as the size of the population, the number of generations, and the rates of crossover and mutation. Finding the optimal set of parameters may require extensive experimentation (Liu et al. 2023).

• Convergence Guarantees: While Private-GSD is designed to improve the approximation of statistical properties over generations, there may not be strong theoretical guarantees on the convergence rate or the quality of the final synthetic dataset compared to the original dataset (Liu et al. 2023).

2. Self-instruct

The process of training Large Language Models (LLMs) to follow instructions often involves labor-intensive tasks, particularly in dataset creation. OpenAI, for instance, employed numerous annotators to develop an open-domain instruction dataset for training InstructGPT. Conversely, Vicuna utilized around 70,000 user-shared conversations from ShareGPT.com to fine-tune a LLaMA model. Human annotation introduces notable challenges. According to Xao et al. (2023), these include the tendency of human-created instructions to skew towards easier levels and the issue of annotator fatigue, which limits the production of complex instructions over extended periods. This can lead to LLMs generating hallucinations. To counter these issues, synthetic open-domain instruction data generation is essential.

In computational linguistics, Self-instruct was developed by Wang et al. (2022) as a technique that leverages the inference capabilities of LLMs to generate synthetic, open-domain instruction data. Self-instruct was used in Alpaca, a fine-tuned LLM from Stanford, and it was a key component in the tech stacks of AI companies like Continuum Labs and Lightning AI. This showcases the technique's practical value and versatility.

The process of Self-instruct begins with a seed set of manually-written tasks. This initial step is crucial as it lays the foundation for the model to generate new instructions. GPT-3 is chosen for its robust inference capabilities. It's tasked with a simple job: take an instruction and an input, and produce an output.

1. Generating Task Instructions. Starting from a pool of 175 tasks, the model starts generating new instructions. In each iteration, a mix of human-written and model-generated instructions are used to enrich the task pool, ensuring diversity and complexity.

2. Task Discrimination: Classification or Not? This stage involves discerning whether a task is a classification task. The model is prompted with a mix of classification and non-classification instructions to make this determination, a critical step for tailoring the generation approach.

3. Crafting Input-Output Pairs. Depending on the task's nature, the model adopts either an input-first or output-first strategy. This bifurcation allows for a more nuanced generation of data, aligning closely with the task's requirements.

4. Filtering and Postprocessing. The final step involves filtering to ensure quality. Using metrics like ROGUE-L similarity, the model removes redundant or low-quality instructions, refining the task pool to only include the most unique tasks.

Pros

• Diverse Instruction Sets. By generating instructions from the model itself, Self-instruct facilitates the creation of a more diverse and creative set of tasks.

• Cost-effectiveness and Scalability. Generating instructions using the model itself is more cost-effective compared to the traditional method of collecting large-scale human-annotated datasets.

• Reduction in Dependency on Human-Written Instructions. Self-instruct minimizes this dependency by generating its own instructions.

Cons

• Model Dependency: The technique's reliance on a specific LLM inherits its limitations (Wang et al. 2022).

• Bias Reinforcement: There's a struggle to produce balanced labels, inadvertently amplifying the model's inherent biases (Wang et al. 2022).

• Dataset Growth Plateau: The exponential growth of new instructions eventually levels off, indicating a saturation point (Wang et al. 2022).

3. Evol-instruct

Evol-instruct is a technique revolutionizing the creation of synthetic open-domain instruction data across various difficulty levels. Evol-instruct has been instrumental in projects like Ragas's RAG pipeline and the WizardCoder LLM, surpassing models like Anthropic's Claude and Google's Bard in performance. Its application continues to expand across AI enterprises, like Clarifai, demonstrating its pivotal role in advancing LLM capabilities.

The core concept involves evolving a known set of instruction-answer pairs over a series of epochs, each designed to enhance the dataset's complexity, richness, and diversity.

The initial dataset comprises pairs of instruction-reply of the form

for 𝑘 instructions. Each epoch 𝑡 updates the dataset to D(t+1) by refining each instruction 𝐼(𝑡)​ through a prompt to an LLM, which then generates an improved instruction 𝐼(𝑡+1)​. This new instruction is used to obtain a corresponding answer R(𝑡+1). After 𝑀 epochs, this iterative process results in multiple datasets 𝐷(1),…,𝐷(𝑀).

The evolution process bifurcates into two types of prompts: in-depth and in-breadth evolving. In-depth evolving involves enhancing instructions through adding constraints, deepening content, concretizing details, increasing reasoning steps, and complicating inputs. Each modification aims to incrementally raise the instruction's difficulty, with a word limit of 10 to 20.

In-breadth evolving, on the other hand, focuses on generating new instructions from existing ones to broaden topic and skill coverage and enhance diversity.

An additional step, Elimination evolving, filters out unsuccessful instruction evolutions. An instruction is deemed unsuccessful if it adds no new information, complicates response generation for the LLM, consists only of punctuation and stop words, or merely copies words from the prompt.

Upon completing all epochs, the instruction-answer pairs from 𝐷(1),…,𝐷(𝑀) are shuffled randomly to ensure a uniform distribution of varying difficulty levels.

Pros

• Variety and Challenge. The technique excels in generating tasks of varying difficulty, pushing the boundaries of what models can learn (Xu et al. 2023).

• Quality. Human annotator experiments reveal superior performance on complex instructions, showcasing the high caliber of the synthetic dataset (Xu et al. 2023).

Cons

• Failure Rate in Instruction Evolution. The Evol-instruct method sometimes fails during the instruction evolution process (Xu et al. 2023).

• Complexity Management. While Evol-instruct is designed to increase the complexity of instructions, managing this complexity effectively is challenging. There is a risk of generating instructions that are too complex (Xu et al. 2023).

• Quality Control. Ensuring consistent quality across the evolved instructions is another challenge. The method relies heavily on the initial quality of the seed instructions and the effectiveness of the LLM used in the evolution process (Xu et al. 2023).

Marginal-based Methods

Marginal-based methods focus on modelling the marginal distributions and inter-attribute dependencies in the original data. They typically construct a probabilistic model, such as a graphical model or Bayesian network, to capture these statistical properties. Synthetic data is then generated by sampling from this model. Some marginal-based methods provide rigorous guarantees of differential privacy, making them well-suited for tabular data synthesis when preserving privacy is crucial.

Marginal-based methods is a nascent field in private synthetic data and it is still looking to be adopted in relation to other competing techniques like GANs. Tumult Labs is looking at these techniques and currently researching its potential.

1. Multiplicative Weights Exponential Mechanism (MWEM)

The MWEM (Multiplicative Weights Exponential Mechanism) algorithm proposed by Hardt et al. (2010) is designed for differentially private data release, particularly focusing on producing synthetic datasets that respect differential privacy while answering a set of linear queries. The algorithm is a combination of the multiplicative weights update rule and the Exponential Mechanism, which are used to iteratively refine an approximation of a dataset in a way that balances privacy and accuracy.

Let D be a dataset with d columns. A marginal for a subset r of the d columns is a histogram for r, that is, a table that counts the number of occurrences of tuples with columns in r. A marginal is often referred to as a marginal query. A workload W is a collection of marginal queries. A marginal-based method takes a workload as input, then it adapts intelligently to the queries in the workload, and generates synthetic data that is tailored to these queries.

The MWEM has three main components. A Multiplicative Weights Update Rule that is used to adjust the weights of an approximating dataset to better reflect the true dataset based on the discrepancies in query responses. The Exponential Mechanism selects the most informative queries to improve the dataset approximation. The Laplace Mechanism for adding noise to the query results to ensure differential privacy.

The MWEM algorithm proceeds as shown in the figure below.

It starts with an initial approximation of the dataset, typically a uniform distribution over the data domain. Then it uses the exponential mechanism to select a query that is poorly explained by the current approximation. Then the selected query on the true dataset is measured, adding noise via the Laplace mechanism to ensure privacy. These steps of select-measure are repeated for a number of iterations. Finally, the output is an average of the approximations across all iterations, which forms the synthetic dataset.

An improved algorithm based on MWEM was developed by McKenna et al. (2022), where MWEM is enhanced with the Probabilistic Graphical Model or PGM and the Gaussian mechanism is used instead of the Laplace mechanism.

Pros

• Privacy. MWEM is designed with differential privacy mechanisms (McKenna et al. 2022).

• Effectiveness. MWEM+PGM is currently the best method for generating synthetic data with privacy guarantees (Tao et al. 2023).

Cons

• Limited data types. MWEM have only been tested on discrete data. Much research still needs to be done in categorical and numerical data (McKenna et al. 2022).

• Lack of maturity. The field of marginal-based methods is very recent, and still much research is needed to have a good understanding of its impact and applications (Ponomareva et al. 2023).

2. PrivBayes

PrivBayes is a differentially private method proposed by Zhang et al. (2017) for releasing high-dimensional data, which is effective when dealing with datasets that contain a large number of attributes. The algorithm operates in three main phases: network learning, distribution learning, and data synthesis.

In the first phase, PrivBayes constructs a Bayesian network that approximates the full-dimensional distribution of the dataset. This network is built using a differentially private method, ensuring that the privacy of the data is maintained. The Bayesian network is a graphical model that represents a set of variables and their conditional dependencies via a directed acyclic graph.

Once the Bayesian network is established, PrivBayes computes a set of differentially private conditional distributions for the data in the subspaces defined by the network. This involves adding noise to the distributions to ensure differential privacy, typically using mechanisms like the Laplace mechanism.

In the final phase, PrivBayes uses the noisy conditional distributions and the structure of the Bayesian network to generate a synthetic dataset. This dataset is an approximation of the original data but is constructed in such a way that it maintains the privacy of the individuals in the dataset.

Pros

• Handling High-Dimensional Data. PrivBayes handles the challenge of releasing high-dimensional data, which is always a barrier for traditional differential privacy methods.

• Reduced Noise Addition. By focusing on low-dimensional marginals, PrivBayes can inject less noise resulting in higher utility of the released data.

• Flexibility in Query Evaluation. The synthetic data generated by PrivBayes can support a wide range of queries allowing the estimation of various statistical properties.

• Scalability. The method scales well with the size of the data and the number of attributes, making it suitable for large datasets.

Cons

• Complexity in Network Construction. The process of learning a differentially private Bayesian network is complex and can be computationally intensive.

• Accuracy Dependence on Network Quality. The accuracy of the synthetic data heavily depends on the quality of the Bayesian network constructed.

• Parameter Sensitivity. The performance of PrivBayes can be sensitive to the choice of parameters, such as the privacy budget allocated to different phases of the algorithm.

TLDR

This is a quick summary of the key points covered in the article about different techniques for generating synthetic data for machine learning.

Generative Methods use AI models to generate new synthetic data. Popular techniques include Variational Autoencoders (VAEs), Generative Adversarial Networks (GANs), and Diffusion Models. VAEs compress data into a lower-dimensional space and then reconstruct it, allowing for data generation. GANs have two competing AI models, a generator creating fake data and a discriminator identifying real vs. fake data. Diffusion Models gradually add noise to an image and then learn to reverse the process, generating new images.

Evolutionary Methods iteratively construct synthetic data from a seed dataset by applying operations like combination and mutation. Techniques like Private-GSD use genetic algorithms, while Self-instruct and Evol-instruct leverage large language models to generate and evolve open-domain instruction data.

Marginal-based Methods model the marginal distributions and inter-attribute dependencies in the original data using probabilistic models like graphical models or Bayesian networks. Synthetic data is then generated by sampling from these models. Techniques like MWEM and PrivBayes provide differential privacy guarantees, making them suitable for sensitive tabular data.

The table below summarizes the different aspects of the methods.

Synthetic data is disrupting the AI industry status quo. As techniques like generative models, evolutionary methods, and marginal-based approaches continue to advance, the potential applications of synthetic data are boundless. Embracing this technology can empower businesses to unlock more accessible, reliable, and impactful AI.

Decentralized AI networks amplify this problem, as nodes host various models and compete for users' attention and money. We need a way to verify their claims and protect clients from fraud. For example, paying top dollar for the latest LLAMA outputs but getting inferior model results. The stakes couldn't be higher in a decentralized AI ecosystem.

The solution to this is Verifiable inference, using mechanisms to verify the use of specific models for generating inferences. This enables honest competition and protects clients.

While zero-knowledge proofs as a solution have been explored extensively by the decentralized AI community, and we discussed it in-depth in a previous article, they are currently too slow and expensive. At Bagel, we have investigated more practical and widely used alternatives from the traditional AI world, such as watermarking (Jia et al. 2021) and fingerprinting (Chen et al. 2019). Initially developed to prevent model extraction attacks, these methods also serve as unique model identifiers thus enabling verifiability. Compared to zero-knowledge approaches, they offer enhanced efficiency and ease of use, better aligning with current applications of verifiable deep learning models.

Today, we're open sourcing our research. Our goal is to empower the decentralized AI community and inspire builders to explore diverse, high-performance solutions from traditional AI. Together, we can create a more robust decentralized AI ecosystem that benefits the mainstream AI market.

If you're in a rush, we have a TLDR at the end.

Watermarking

How it works

In the landscape of machine learning, protecting foundation model as an intellectual property (IP) rights is a paramount concern. Enter watermarks – a clever defensive mechanism against model extraction attacks. Just like classical digital watermarking, a form of steganography where a message is concealed within a digital object, watermarks for machine learning models serve as a guardian of ownership, authenticity, and integrity. They help safeguard IP and enforce licenses, ensuring that the hard work of model creators doesn't go unrecognized.

Lederer et al. (2023) outline a plethora of characteristics that a watermark for a machine learning model must possess, but three stand out as the most crucial:

Effectiveness - Effectiveness means that the watermark can be verified at any time by the model's creator.

Fidelity - Fidelity ensures that the model's accuracy remains unaffected by the watermark's presence.

Robustness - That's the watermark's ability to withstand a barrage of attacks, from fine-tuning and model compression to watermark detection, removal, overwriting, or invalidation.

But how does the magic happen? Any watermarking embedding method consists of two algorithms: an extraction algorithm that retrieves the watermark from a model, and a verification algorithm that confirms its presence.

There are two main families of watermark embedding techniques: i) white-box watermarking and, ii) black-box watermarking. The general process of watermarking is illustrated in the figure below.

White-box watermarking requires full access to the model. The idea is to embed a signature s into the model's weights during training by adding an extra term to the loss function. This is done carefully to maintain the model's accuracy. The embedding is achieved by modifying the loss function through a regularizer parameter. Let w be a vector of all weights in a model. The goal is to embed s into w using an embedding matrix M, which acts as a secret key usually held by the model owner. The watermark is extracted by applying M to the weight vector w, followed by a threshold function (see Nagai et al. (2018) for details).

Black-box watermarking, on the other hand, only requires query access to the model. It creates backdoors on data using data poisoning. In this context, a backdoor is a set of input-output pairs known to the model owner that triggers a behavior not predictable by model consumers. For example, deliberately adding wrong labels to data points. The goal is to ensure that the model performs correctly on the main classification task, but the backdoor exhibits a specific behavior defined by the model owner.

Next we show two applications of watermarks in the context of verifying ownership of models and verifying inference from models.

Watermarking for Verifiable Inference

Public Models

Black-box watermarking is a game-changer when it comes to verifiable inference for public models. It's like a secret handshake between the model creator and the user, ensuring that the model is authentic and trustworthy.

Here's how it works: the model creator embeds a watermark and then discloses its presence to the world. It's like a badge of honor, a mark of quality. Any user who interacts with the model can then authenticate it, verifying that it's the real deal.

Zhong et al. (2020) took this concept to the next level. They developed a black-box watermarking technique that adds new labels to inputs that have nothing to do with the original dataset. It's like adding a secret code that only the model creator knows.

When a user wants to verify the watermark, they simply query these special inputs. If the watermark is present in the inference results, it's a clear sign that the model is authentic. It's like a digital signature that can't be forged.

Chen et al. (2019) took a slightly different approach. They generated watermark keys and used fine-tuning to embed a signature in the model. It's like hiding a secret message in plain sight.

With the watermark keys in hand, any user can interact with the model and extract the signature through its predictions. It's like unlocking a hidden layer of authentication, ensuring that the model is genuine and trustworthy.

Private Models

When a machine learning model is privately hosted, the mechanism is almost the same, black-box watermarking is used, except the model owner hosts the model privately. Model consumers gets access to the model via and API or gateway. First the model creator runs the watermark algorithm on his model and loads it to the hosting service. Then, via API access, users can query the model via inputs and obtains inferences as output. The backdoor is still public, so users can later verify the correct use of the model.

With knowledge of the backdoor, then users can run the verification algorithm to check for the watermark. The verification algorithm will tell if the correct model was used or not.

For private models offering API access, Xu and Yuan (2019) showed how to add unique serial numbers to the trigger-set of watermarks. Their serial number technique is independent of labels and can be supported by digital certification authorities.

Fingerprinting

How it works

Fingerprinting is a fascinating approach to model identification that differs from watermarking in a key way. While watermarking embeds a secret message into the model, fingerprinting relies on the model's inherent characteristics to create a unique identifier (Chen et al., 2019; Lukas et al., 2019). This identifier, akin to a digital DNA, can be transferred to any models derived from the original, making it a robust tool for proving model provenance.

A fingerprinting scheme consists of two essential components: a generation algorithm for creating fingerprints and a verification algorithm for confirming their presence.

• Generate(M,D): Given white-box access to a model M and a dataset D, this procedure outputs a fingerprint F and verification keys K={M(x) : x in F}.

The figure below illustrates the process of generating fingerprints.

• Verify(M’(F),K): Given black-box access to a model M’, fingerprint F, and verification keys K, this procedure outputs 1 if M’ is verified by the fingerprint and 0 otherwise.

Generating fingerprints is a complex task that requires a deep understanding of the model and its training data. One approach is to use adversarial examples. By adding carefully crafted noise to a correctly predicted data point, the model can be tricked into predicting a desired label. The data point and noise pair form an adversarial example, which can serve as a fingerprint.

Cao et al. (2019) demonstrated how to construct adversarial examples near the model's decision boundary and leverage their transferability to surrogate models. Peng et al. (2022) employed Universal Adversarial Perturbations (UAPs), which are vectors drawn from a low-dimensional subspace containing most normal vectors of the decision boundary. UAPs can function as a model's fingerprint, allowing the owner to verify if a given UAP vector v lies within the suspect model's UAP space.

In the following section, we will explore how fingerprints can be applied to verifiable inference. Fingerprinting provides a unique method for identifying models based on their inherent characteristics, distinguishing it from watermarking's embedded secret messages. As we delve further into this captivating topic, we will uncover the power of fingerprinting in ensuring model provenance.

Fingerprinting for Verifiable Inference

Public Models

When a model is public, the creator generates a fingerprint and keys, making them accessible to all. Consumers can download the model, fingerprint, and keys as a package.

With full access, consumers can modify the model through fine-tuning or compression. But how can they verify it's the genuine article?

Enter the verification algorithm. By running it on the query results and fingerprint, consumers can confirm the model's authenticity.

For public models, unique identifier fingerprints are crucial. Lukas et al. (2021) proposed a game-changing technique: conferrable examples. These crafted examples act as adversarial inputs not just for the target model, but also for any imitators. They transfer to surrogates but not to independently trained models.

Zhao et al. (2020) introduced adversarial marks, another transferrable fingerprint that can't be removed without sacrificing significant accuracy.

These innovative techniques ensure that public models remain trustworthy and authentic. They give consumers the power to verify, while creators can share their work with confidence.

In a world where AI models are increasingly open and accessible, fingerprinting is a vital tool. It's the key to maintaining trust and integrity in the face of modification and imitation.

So, the next time you download a public model, look for the fingerprint. It's your guarantee of authenticity in an ever-evolving AI landscape.

Private Models

Model creators may keep their models private to protect intellectual property rights. When generating a fingerprint, they can host the model privately and provide consumer access through a public API.

The verification algorithm for fingerprints always works on black-box access to the model. With public knowledge of the fingerprint, any consumer can verify it using the API.

Private models with public APIs are crucial for giving users access to LLMs like ChatGPT and Claude. However, LLMs pose a challenge in generating fingerprints due to their vast number of parameters.

Recent work by Gu et al. (2002), Li et al. (2023), and Xu et al. (2024) demonstrates how to construct fingerprints for LLMs. These methods implant input-output pairs that exploit the model's inherent characteristics.

Xu et al. (2024) presents the fastest fingerprinting method, claiming to fingerprint LLAMA2-13B in less than a minute using a single A100 GPU. Their technical contribution is that their methods can work in either white-box mode using an F-transformer or black-box mode using fine-tuning.

As AI continues to evolve, fingerprinting will play a vital role in ensuring the integrity of private models accessed through public APIs. It's a powerful tool that strikes a balance between accessibility and security.

TLDR

Verified inference is crucial in MLaaS, both centralized and decentralized. Consumers should be able to confirm the model generating the inference is the one they requested and paid for.

ZKML proves a model was executed on provided data, but compromises data privacy. Current ZK proof generation techniques are time-consuming due to cryptographic operations on top of the computation being proved (Garg et al., 2023).

ZKML allows perfect proofs of inference and provenance but struggles with real-time proofs for relevant deep neural networks like LLMs (GPT2, Llama).

Watermarking and fingerprinting fill the gaps. Practical and effective inference verification is a must in decentralized AI where trust cannot be assumed. Xu et al. (2024) generated fingerprints in the largest LLMs in under one minute.

Verifiable inference is achievable through all three techniques. Watermarking and fingerprinting provide extraction and verification algorithms, while ZK offers proofs of knowledge.

Efficiency varies significantly. ZK is the slowest, while watermarking and fingerprinting can be fast, making them more applicable considering current state-of-the-art.

Knowledge distillation, a model compression attack, is a threat watermarking and fingerprinting are designed to protect against by passing the identifying markers to the compressed model. ZK does not address this issue, which is significant in the traditional AI industry and will be covered in a future article.

Output accountability is where ZK shines. While watermarks and fingerprints can establish ownership, they don't inherently ensure model integrity or output correctness. ZKPs prove a model was evaluated correctly on specific inputs, enabling accountability (Lederer et al., 2023; Oliynik et al., 2023).

As AI reshapes our world, we can't build the future on blind trust. Verifiability is essential for decentralized AI networks. We need transparency and accountability. AI's potential is vast, but honesty and integrity must come first. Watermarking and fingerprinting are key alternatives tools for that.

The stakes are high. We have to get this right.

Bagel Labs is a distributed machine learning research lab.

To propel human civilization forward in energy, healthcare, and collaboration, it is crucial to enable AI systems that train and generate inference on data while maintaining full end-to-end privacy. At Bagel, we believe accessing a fundamental resource like knowledge, for both human-driven and autonomous AI, should not entail a compromise on privacy.

We have applied and experimented with almost all the major privacy-preserving machine learning (PPML) mechanisms. Below, we share our insights, our approach, and some research breakthroughs.

And if you're in a rush, we have a TLDR at the end.

Privacy-preserving Machine Learning (PPML)

Recent advances in academia and industry have focused on incorporating privacy mechanisms into machine learning models, highlighting a significant move towards privacy-preserving machine learning (PPML). At Bagel, we have experimented with all the major PPML techniques, particularly those post differential privacy. Our work, positioned at the intersection of AI and cryptography, draws from the cutting edge in both domains.

Our research covered a wide range of PPML techniques suitable for our platform. Among those, Differential Privacy (DP), Federated Learning, Zero-knowledge Machine Learning (ZKML) and Fully Homomorphic Encryption Machine Learning (FHEML) stood out for their potential in PPML.

First, we will delve into each of these, examining their advantages and drawbacks. In subsequent posts, we will describe Bagel's approach to data privacy, which addresses and resolves the challenges associated with the existing solutions.

Differential Privacy (DP)

One of the first and most important techniques with a mathematical guarantee for incorporating privacy into data is differential privacy or DP (Dwork et al. 2006), addressing the challenges faced by earlier methods with a quantifiable privacy definition.

DP ensures that a randomized algorithm, A, maintains privacy across datasets D1 and D2—which differ by a single record—by keeping the probability of A(D1) and A(D2) generating identical outcomes relatively unchanged. This principle implies that minor dataset modifications do not significantly alter outcome probabilities, marking a pivotal advancement in data privacy.

The application of DP in machine learning, particularly in neural network training and inference, demonstrates its versatility and effectiveness. Notable implementations include adapting DP for supervised learning algorithms by integrating random noise at various phases: directly onto the data, within the training process, or during inference, as highlighted by Ponomareva et al. (2023) and further references.

The balance between privacy and accuracy in DP is influenced by the noise level: greater noise enhances privacy at the cost of accuracy, affecting both inference and training stages. This relationship was explored by Abadi et al. in (2016) through the introduction of Gaussian noise to the stochastic gradient descent (DP-SGD) algorithm, observing the noise's impact on accuracy across the MNIST and CIFAR-10 datasets.

An innovative DP application, Private Aggregation of Teacher Ensembles (PATE) by Papernot et al. in (2016), divides a dataset into disjoint subsets, training networks on each without privacy, termed as teachers. These networks' aggregated inferences, subjected to added noise for privacy, inform the training of a student model to emulate the teacher ensemble. This method also underscores the trade-off between privacy enhancement through noise addition and the resultant accuracy reduction.

Further studies affirm that while privacy can be secured with little impact on execution times (Li et a. 2015), stringent privacy measures can obscure discernible patterns essential for learning (Abadi et al. 2016). Consequently, a certain level of privacy must be relinquished in DP to facilitate effective machine learning model training, illustrating the nuanced balance between privacy preservation and learning efficiency.

Pros of Differential Privacy

The advantages of using DP are:

Effortless. Easy to implement into algorithms and code.

Algorithm independence. Schemes can be made independent of the training or inference algorithm.

Fast. Some DP mechanisms have shown to have little impact on the execution times of algorithms.

Tunable privacy. The degree of desired privacy can be chosen by the algorithm designer.

Cons of Differential Privacy

Access to private data is still necessary. Teachers in the PATE scheme must have full access to the private data (Papernot et al. 2016) in order to train a neural network. Also, the stochastic gradient descent algorithm based on DP only adds noise to the weight updates and needs access to private data for training (Abadi et al. 2016).

Privacy-Accuracy-Speed trade-off on data. All implementations must sacrifice some privacy in order to get good results. If there is no discernable pattern in the input, then there is nothing to train (Feyisetan et al. 2020). The implementation of some noise mechanisms can impact execution times, necessitating a balance between speed and the goals of privacy and accuracy.

Zero-Knowledge Machine Learning (ZKML)

A zero-knowledge proof system (ZKP) is a method allowing a prover P to convince a verifier V about the truth of a statement without disclosing any information apart from the statement's veracity. To affirm the statement's truth, P produces a proof π for V to review, enabling V to be convinced of the statement's truthfulness.

Zero-Knowledge Machine Learning (ZKML) is an approach that combines the principles of zero-knowledge proofs (ZKPs) with machine learning. This integration allows machine learning models to be trained and to infer with verifiability.

For an in-depth examination of ZKML, refer to the work by Xin et al. in (2023). Below we provide a brief explanation that focuses on the utilization of ZKPs for neural network training and inference.

ZKML Inference

Consider an unlabeled dataset A and a pretrained neural network N tasked with labeling each record in A. To generate a ZK proof of N's computation during labeling, an arithmetic circuit C representing N is required, including circuits for each neuron's activation function. Assuming such a circuit C exists and is publicly accessible, the network's weights and a dataset record become the private and public inputs, respectively. For any record a of A, N's output is denoted by a pair (l,π), where l is the label and π is a zero-knowledge argument asserting the existence of specific weights that facilitated the labeling.

This model illustrates how ZK proves the accurate execution of a neural network on data, concealing the network's weights within a ZK proof. Consequently, any verifier can be assured that the executing agent possesses the necessary weights.

ZKML Training

ZKPs are applicable during training to validate N's correct execution on a labeled dataset A. Here, A serves as the public input, with an arithmetic circuit C depicting the neural network N. The training process requires an additional arithmetic circuit to implement the optimization function, minimizing the loss function. For each training epoch i, a proof π_i is generated, confirming the algorithm's accurate execution through epochs 1 to i-1, including the validity of the preceding epoch's proof. The training culminates with a compressed proof π, proving the correct training over dataset A.

The explanation above illustrates that during training, the network's weights are concealed to ensure that the training is correctly executed on the given dataset A. Additionally, all internal states of the network remain undisclosed throughout the training process.

Pros of ZKML

The advantages of using ZKPs with neural networks are:

Privacy of model weights. The weights of the neural network are never revealed during training or inference in any way. The weights and the internal states of the network algorithm are private inputs for the ZKP.

Verifiability. The proof certifies the proper execution of training or inference processes and guarantees the accurate computation of weights.

Trustlessness. The proof and its verification properties ensure that the data owner is not required to place trust in the agent operating the neural network. Instead, the data owner can rely on the proof to confirm the accuracy of both the computation and the existence of correct weights.

Cons of ZKML

The disadvantages of using ZKPs with neural networks are:

No data privacy. The agent running the neural network needs access to the data in order to train or do inference. Data is considered a parameter that is publicly known to the data owner and the prover running the neural network (Xing et al. 2023).

No privacy for the model’s algorithm. In order to create a ZK proof, the algorithm of the entire neural network should be publicly known. This includes the activation functions, the loss function, optimization algorithm used, etc (Xing et al. 2023).

Proof generation of an expensive computation. Presently, the process of generating a ZK proof is computationally demanding. Creating a proof for each epoch within a training algorithm can exacerbate the computational burden of an already resource-intensive task.

Federated Learning (FL)

In Federated Learning or FL we look to train a global model using a dataset that is distributed in multiple servers with local data samples but without each server sharing their local data.

In FL there is a global objective function that is being optimized which is defined as

where n is the number of servers, each variables is the set of parameter as viewed by the server i, and each function is a local objective function of server i. FL tries to find the best set of values that optimizes f.

The figure below shows the general process in FL.

1. Initialization. An initial global model is created and distributed by a central server to all other servers.

2. Local training. Each server trains the model using their local data. This ensures data privacy and security.

3. Model update. After training, each server shares with the central server their local updates like gradients and parameters.

4. Aggregation. The central server receives all local updates and aggregates them into the global model, for example, using averaging.

5. Model distribution. The updated model is distributed again with local servers and the previous steps are repeated until a desired level of performance is achieve by the global model.

Since local servers never share their local data, FL guarantees privacy over that data. However, the model being constructed is shared among all parties, and hence, its structure and set of parameters are not hidden.

Pros of FL

The advantages of using FL are:

Data privacy. The local data on the local servers are never shared. All computations are done locally, and there is no need of communication between them.

Distributed computing. The creation of the global model is distributed among local servers, thereby parallelizing a resource-intensive computation. Thus, FL is considered a distributed machine learning framework (Xu et al. 2021).

Cons of FL

The disadvantages of using FL are:

Model is not private. The global model is shared among each local server in order to do their computations locally. This includes the aggregated weights and gradients at each step of the FL process. Thus, each local server is aware of the entire architecture of the global model (Konečný et al. 2016).

Data leakage. Recent research indicates that data leakage remains a persistent issue, notably through mechanisms such as gradient sharing—see for example Jin et al. (2022). Consequently, FL cannot provide complete assurances of data privacy.

Trust. Since no proofs are generated in FL, every party involved in the process need to be trusted that their computation and parameters were computed as expected (Gao et al. 2023).

Fully Homomorphic Encryption (FHE)

At its core, homomorphic encryption permits computations on encrypted data. By "homomorphic," we refer to the capacity of an encryption scheme to allow specific operations on ciphertexts that, when decrypted, yield the same result as operations performed directly on the plaintexts.

Consider a scenario with a secret key k and a plaintext m. In an encryption scheme (E,D), where E and D represent encryption and decryption algorithms respectively, the condition D(k,E(k,m))=m must hold. A scheme (E,D) is deemed fully homomorphic if for any key k and messages m, the properties E(k,m+m’)=E(k,m)+E(k,m’) and E(k,m*m’)=E(k,m)* E(k,m’) are satisfied, with addition and multiplication defined over a finite field. If only one operation is supported, the scheme is partially homomorphic. This definition implies that operations on encrypted data mirror those on plaintext, crucial for maintaining data privacy during processing.

In plain words, if we have a fully homomorphic encryption scheme, then operating over the encrypted data is equivalent to operating over the plaintext. We will write FHE to refer to a fully homomorphic encryption scheme. The figure below shows how an arbitrary homomorphic operation works over a plaintext and ciphertext.

The homomorphic property of FHE makes it invaluable in situations where data must remain secure while still being used for computations. For instance, if we possess sensitive data and require a third party to perform data analysis on it, we can rely on FHE to encrypt the data. This allows the third party to conduct analysis on the encrypted data without the need for decryption. The mathematical properties of FHE guarantee the accuracy of the analysis results.

FHE Inference

Fully Homomorphic Encryption (FHE) can be used to perform inference in neural networks while preserving data privacy. Let's consider a scenario where N is a pretrained neural network, A is a dataset, and (E,D) is an asymmetric FHE scheme. The goal is to perform inference on a record a of A without revealing the sensitive information contained in a to the neural network.

The inference process using FHE begins with encryption. The data owner encrypts the record a using the encryption algorithm E with the public key public_key, obtaining the encrypted record a’ = E(public_key, a).

Next, the data owner sends the encrypted record a’ along with public_key to the neural network N. The neural network N must have knowledge of the encryption scheme (E,D) and its parameters to correctly apply homomorphic operations over the encrypted data a’. Any arithmetic operation performed by N can be safely applied to a’ due to the homomorphic properties of the encryption scheme.

One challenge in using FHE for neural network inference is handling non-linear activation functions, such as sigmoid and ReLU, which involve non-arithmetic computations. To compute these functions homomorphically, they need to be approximated by low-degree polynomials. The approximations allow the activation functions to be computed using homomorphic operations on the encrypted data a’.

After applying the necessary homomorphic operations and approximated activation functions, the neural network N obtains the inference result. It's important to note that the inference result is still in encrypted form, as all computations were performed on encrypted data.

Finally, the encrypted inference result is sent back to the data owner, who uses the private key associated with the FHE scheme to decrypt the result using the decryption algorithm D. The decrypted inference result is obtained, which can be interpreted and utilized by the data owner.

By following this inference process, the neural network N can perform computations on the encrypted data a’ without having access to the original sensitive information. The FHE scheme ensures that the data remains encrypted throughout the inference process, and only the data owner with the private key can decrypt the final result.

It's important to note that the neural network N must be designed and trained to work with the specific FHE scheme and its parameters. Additionally, the approximation of non-linear activation functions by low-degree polynomials may introduce some level of approximation error, which should be considered and evaluated based on the specific application and accuracy requirements.

FHE Training

The process of training a neural network using Fully Homomorphic Encryption (FHE) is conceptually similar to performing inference, but with a few key differences. Let's dive into the details.

Imagine we have an untrained neural network N and an encrypted dataset A’ = E(public_key, A), where E is the encryption function and public_key is the public key of an asymmetric FHE scheme. Our goal is to train N on the encrypted data A’ while preserving the privacy of the original dataset A.

The training process unfolds as follows. Each operation performed by the network and the training algorithm is executed on each encrypted record a’ of A'. This includes both the forward and backward passes of the network. As with inference, any non-arithmetic operations like activation functions need to be approximated using low-degree polynomials to be compatible with the homomorphic properties of FHE.

A fascinating aspect of this approach is that the weights obtained during training are themselves encrypted. They can only be decrypted using the private key of the FHE scheme, which is held exclusively by the data owner. This means that even the agent executing the neural network training never has access to the actual weight values, only their encrypted counterparts.

Think about the implications of this. The data owner can outsource the computational heavy lifting of training to a third party, like a cloud provider with powerful GPUs, without ever revealing their sensitive data. The training process operates on encrypted data and produces encrypted weights, ensuring end-to-end privacy.

Once training is complete, the neural network sends the collection of encrypted weights w’ back to the data owner. The data owner can then decrypt the weights using his private key, obtaining the final trained model. He is the sole party capable of accessing the unencrypted weights and using the model for inference on plaintext data.

There are a few caveats to keep in mind. FHE operations are computationally expensive, so training a neural network with FHE will generally be slower than training on unencrypted data.

Pros of FHE

The advantages of using FHE are:

Data privacy. Third-party access to encrypted private data is effectively prevented, a security guarantee upheld by the assurances of FHE and lattice-based cryptography (Gentry 2009).

Model privacy. Training and inference processes are carried out on encrypted data, eliminating the need to share or publicize the neural network's parameters for accurate data analysis.

Effectiveness. Previous studies have demonstrated that neural networks operating on encrypted data using FHE maintain their accuracy—see for example Nandakumar et al. (2019) and Xu et al. (2019). Therefore, we can be assured that employing FHE for training and inference processes will achieve the anticipated outcomes.

Quantum resistance. The security of FHE, unlike other encryption schemes, is grounded in difficult problems derived from Lattice theory. These problems are considered to be hard even for quantum computers (Regev 2005), thus offering enhanced protection against potential quantum threats in the future.

Cons of FHE

The disadvantages of using FHE are:

Verifiability. FHE does not offer proofs of correct encryption nor correct computation. Hence, we must rely on trust that the data intended for encryption is indeed the correct data (Viand et al. 2023).

Speed. Relative to conventional encryption schemes, FHE is still considered to be slow during parameter setups, encryption and decryption algorithms (Gorantala et al. 2023).

Memory requirements. The number of weights that need to be encrypted are proportional to the size of the network. Even for small networks, the RAM memory requirements are in the order of gigabytes (Chen et al. 2018), (Nandakumar et al. 2019).

Usability. FHE schemes use many parameters that need to be carefully tuned and requires extensive experience from users (Al Badawi et al. 2022), (Halevi & Shoup 2020).

TLDR

We examined the four most widely used privacy-preserving techniques in machine learning, focusing on neural network training and inference. We evaluated these techniques across four dimensions: data privacy, model algorithm privacy, model weights privacy, and verifiability.

Data privacy considers the model owner's access to private data. Differential privacy (DP) and zero-knowledge machine learning (ZKML) require access to private data for training and proof generation, respectively. Federated learning (FL) enables training and inference without revealing data, while fully homomorphic encryption (FHE) allows computations on encrypted data.

Model algorithm privacy refers to the data owner's access to the model's algorithms. DP does not require algorithm disclosure, while ZKML necessitates it for proof generation. FL distributes algorithms among local servers, and FHE operates without accessing the model's algorithms.

Model weights privacy concerns the data owner's access to the model's weights. DP and ZKML keep weights undisclosed or provide proofs of existence without revealing values. FL involves exchanging weights among servers for decentralized learning, contrasting with DP and ZKML's privacy-preserving mechanisms. FHE enables training and inference on encrypted data, eliminating the need for model owners to know the weights.

Verifiability refers to the inherent capabilities for verifiable computation. ZKML inherently provides this capability. DP, FL, and FHE would not provide similar levels of integrity assurance.

The table below summarizes our findings:

Subscribe now