🤖 machine learning

Mean-Field Model for Two-Layer Neural Networks Trained with Consensus-Based Optimization

This paper proposes a hybrid Consensus-Based Optimization (CBO) and Adam approach for two-layer neural networks that reduces memory overhead in multi-task learning, while theoretically establishing a mean-field model within the Wasserstein-over-Wasserstein framework that guarantees monotonic variance decrease and convergence.

Original authors: William De Deyn, Michael Herty, Giovanni Samaey

Published 2026-07-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: William De Deyn, Michael Herty, Giovanni Samaey

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to find the absolute lowest point in a vast, foggy mountain range. This mountain range represents the "error" of a computer program (a neural network). Your goal is to get the error as close to zero as possible.

This paper explores a new way to navigate this foggy landscape, comparing it to the standard methods used today. Here is the breakdown in simple terms:

1. The Problem: Getting Stuck in Small Valleys

Usually, computers train these networks by looking at the slope right under their feet and taking a step downhill. This is like a hiker who only looks at the ground immediately in front of them.

The Issue: If the hiker starts in a small valley (a "local minimum"), they might think they've reached the bottom, even though a much deeper valley (the "global minimum") exists just over the next ridge. They get stuck.

2. The New Method: The "Swarm" Approach (CBO)

Instead of one hiker, the authors propose using a swarm of explorers (called "particles").

How it works: Imagine 200 explorers scattered across the mountain. They all talk to each other. Every few minutes, they calculate a "consensus point"—a weighted average of where everyone is.
The Magic: If an explorer is in a high, bad spot, they are pulled strongly toward the group's average. If the group is mostly in a good spot, the whole swarm drifts that way.
The Benefit: Because they are a group, they are less likely to get stuck in a tiny, shallow valley. They can "feel" the landscape better and find the deepest valley together.

3. The Experiments: Testing the Swarm

The authors tested this "Swarm" method (called Consensus-Based Optimization or CBO) against the standard "Hiker" method (called Adam) on two tasks:

Task A: Drawing a Sine Wave (Regression)
- Result: The Swarm found a slightly better, smoother drawing of the wave than the Hiker. It was also more stable, meaning it didn't wobble as much.
- Catch: The Swarm was slower because every single explorer had to check the map at every step.
Task B: Recognizing Handwritten Numbers (MNIST)
- Result: The standard Hiker (Adam) was actually faster and found a very good solution. The Swarm alone was a bit slow.
- The Hybrid Solution: The authors created a Hybrid Team. They took the best of both worlds: they let the Hiker take the lead for speed but kept the Swarm nearby to stabilize the group and prevent them from falling off cliffs.
- Outcome: This Hybrid team was the fastest and most stable of all.

4. The "Recycling" Trick (Multi-Task Learning)

Usually, if you want a computer to learn two different things (like recognizing cats and dogs), you need two separate teams of explorers. This uses up a lot of memory.

The Innovation: The authors realized that if the two tasks are similar, the "best spot" for cats is probably close to the "best spot" for dogs.
The Analogy: Instead of hiring two new teams, you just tell the same 200 explorers to split up. Half of them focus on the cat mountain, and the other half focus on the dog mountain. They share the same starting gear.
Result: You can train on many tasks at once without needing extra memory, because you are "recycling" the same explorers for different jobs.

5. The Big Picture: The "Infinite" View (Mean-Field Models)

The authors didn't just run simulations; they did some heavy math to understand what happens if you have infinite explorers and infinite neurons in the network.

The Math Metaphor: Instead of tracking 200 individual dots, they looked at the "cloud" of dots as a whole fluid.
The Finding: They proved mathematically that as the swarm moves, the "spread" (variance) of the group shrinks steadily. The cloud gets tighter and tighter around the best solution, like a net slowly closing in on a fish.
Verification: They ran computer experiments to show that as they increased the number of neurons and the number of explorers, the error consistently went down, confirming their math was correct.

Summary

The Goal: Train AI better by avoiding getting stuck in bad solutions.
The Tool: A "Swarm" method (CBO) that uses many agents to explore together.
The Win: A Hybrid version (Swarm + Standard) works faster and more reliably than the standard method alone.
The Efficiency: You can reuse the same "swarm" to learn multiple tasks at once, saving memory.
The Theory: They proved mathematically that this swarm method naturally tightens up and converges to a solution, even when viewed as an infinite cloud of data.

The paper concludes that while this method is powerful, it is currently best used for simple, two-layer networks, and adding "noise" (randomness) to the math is still a work in progress.

Technical Summary: Mean-Field Model for Two-Layer Neural Networks Trained with Consensus-Based Optimization

Problem Statement

The training of neural networks is fundamentally an optimization problem aimed at minimizing an empirical risk function. While gradient-based methods like Stochastic Gradient Descent (SGD) and adaptive variants (e.g., Adam) are the standard, they are prone to getting trapped in local minima due to the non-convex nature of the objective function. Furthermore, multi-task learning often requires balancing conflicting gradients and managing significant memory overhead when training separate models or ensembles for different tasks.

This paper investigates the feasibility of using Consensus-Based Optimization (CBO), a gradient-free, global optimization method based on interacting particle systems, for training two-layer neural networks. The authors aim to:

Compare CBO performance against Adam.
Develop a hybrid approach combining CBO and Adam to leverage the strengths of both.
Adapt CBO for multi-task learning to reduce memory overhead.
Establish rigorous mean-field models for the training dynamics in two regimes: infinite network width ( $M \to \infty$ ) and infinite particle count ( $N \to \infty$ ).

Methodology

1. Optimization Frameworks

The study formulates the training of a two-layer neural network $\hat{g}(x; \theta)$ as minimizing the empirical risk $\hat{R}(\theta)$ .

Adam: Used as the baseline gradient-based optimizer, employing adaptive learning rates based on first and second moment estimates of gradients.
Consensus-Based Optimization (CBO): A particle-based method where an ensemble of $N$ particles $\theta_n$ evolves via a stochastic differential equation. The dynamics consist of a drift term towards a consensus point $V^k$ (a weighted average of particles favoring lower risk) and a diffusion term for exploration.
Hybrid Method: A novel update rule combining an Adam step and a CBO step:
$\theta_n^{k+1} = \theta_n^k - \gamma \Delta t (\text{Adam update}) + (1-\gamma) (\text{CBO update})$
where $\gamma \in [0,1]$ controls the balance.
Multi-Task CBO: A formulation where a single particle ensemble is recycled across multiple tasks. Instead of maintaining separate ensembles, particles are assigned to specific tasks, reducing memory overhead. The consensus point is computed per task using the shared particle set.

2. Mean-Field Formulation

The authors reformulate CBO within the Optimal Transport (OT) framework to handle the infinite-width limit of neural networks.

Infinite Width ( $M \to \infty$ ): Instead of representing a network as a point in $\mathbb{R}^{M(d+2)}$ , the network is represented by a probability measure $\mu$ over the parameter space. The CBO dynamics are lifted to the Wasserstein space $P_2(\mathbb{R}^{d+2})$ . The consensus point becomes a barycenter in the Wasserstein space, defined as the minimizer of the weighted sum of squared Wasserstein distances.
Infinite Particles ( $N \to \infty$ ): The dynamics of the particle ensemble are lifted to the space of measures over the Wasserstein space (Wasserstein-over-Wasserstein). The authors derive a time-discrete mean-field model and prove that the variance of the particle distribution decreases monotonically.

3. Numerical Experiments

The paper validates the methods on three tasks:

Sine Approximation: A 1D regression problem comparing CBO and Adam.
MNIST Classification: A multi-class classification task comparing CBO, Adam, and the Hybrid method.
Multi-Task Learning: Approximating 100 shifted sine functions simultaneously using the Multi-Task CBO strategy.
Square Approximation (Mean-Field Validation): Training networks with varying widths ( $M$ ) and particle counts ( $N$ ) to verify convergence to the mean-field limits.

Key Results

Performance Comparison

Sine Regression: CBO achieved a lower final empirical risk and exhibited greater stability than Adam, though it incurred a higher computational cost per iteration due to the need for forward passes for all particles.
MNIST Classification: Adam converged faster and achieved lower risk than standard CBO. However, the Hybrid Method outperformed both, converging faster than Adam alone. The hybrid approach allowed for a larger learning rate (which would be unstable for pure Adam) because the CBO component stabilized the dynamics.
Multi-Task CBO: The method successfully minimized risk across 100 distinct tasks using a single particle ensemble. Both the median and minimum empirical risks decreased monotonically, confirming that particle recycling effectively handles multi-objective optimization without the memory cost of separate ensembles.

Mean-Field Convergence

Variance Decay: Theoretical analysis (Proposition 3) proves that in the infinite-particle limit, the variance of the particle ensemble decays geometrically as $(1-\Delta t)^{2k}$ , ensuring consensus.
Empirical Validation: Numerical experiments on square approximation showed that empirical risk decreases monotonically as both the network width $M$ and the number of particles $N$ increase, saturating near the mean-field regimes.

Significance and Claims

The paper claims the following contributions and significance:

Hybrid Optimization: The study demonstrates that combining CBO with Adam yields a robust optimizer that is faster and more stable than either method individually, particularly in non-convex landscapes like MNIST.
Memory-Efficient Multi-Task Learning: The Multi-Task CBO formulation offers a practical way to train models on multiple related tasks by recycling particles, significantly reducing memory overhead compared to training separate models.
Theoretical Unification: The work bridges particle-based optimization and neural network theory by:
- Reformulating CBO in the Optimal Transport framework, allowing for the treatment of infinitely wide neural networks as continuous measures.
- Deriving a time-discrete mean-field model for CBO dynamics on these infinite-width networks.
- Proving that the variance of the particle population decreases monotonically in the mean-field limit.
Limitations and Future Directions: The authors modestly note that while variance decay ensures consensus, it does not strictly guarantee convergence to a global minimizer in the OT setting (a link that remains an open problem). Additionally, the current framework is restricted to two-layer networks, and the OT formulation currently lacks a rigorous diffusion term (relying on added noise in practice). Future work is suggested for extending the analysis to deeper architectures and establishing rigorous convergence proofs.

In summary, the paper posits that CBO provides a viable, globally convergent alternative to gradient descent for neural network training, with specific advantages in stability and multi-task efficiency, and that its dynamics can be rigorously analyzed through mean-field limits in the Wasserstein space.