mixture-of-flows 0.2.0
Flow Matching Mixture of Experts (FM-MoE)
License: MIT
Requires Python: <4.0,>=3.10
Keywords: artificial intelligence, deep learning, machine learning, mixture of experts, moe, flow matching, transformers, llm, large language models, neural networks, sparse models, expert routing, gating network, conditional computation, efficient inference, model scaling, attention mechanism, pytorch, nlp, natural language processing, generative ai, foundation models

# Flow Matching Mixture of Experts (FM-MoE)

[![License](https://img.shields.io/badge/License-MIT-blue.svg)](LICENSE)
[![Python](https://img.shields.io/badge/Python-3.8+-green.svg)](https://www.python.org/)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.0+-red.svg)](https://pytorch.org/)

A PyTorch implementation of Flow Matching Mixture of Experts, a novel neural network architecture that integrates continuous normalizing flows with sparse mixture of experts routing for enhanced representational capacity in deep learning models.

---

## Abstract

Mixture of Experts (MoE) architectures have demonstrated remarkable success in scaling neural networks by activating only a subset of parameters per input. However, traditional MoE implementations rely on standard multi-layer perceptrons (MLPs) as expert networks, which may limit the expressiveness of learned transformations. This work introduces **Flow Matching Mixture of Experts (FM-MoE)**, a framework that replaces conventional MLP experts with flow matching networks. Each expert learns a continuous transformation through an ordinary differential equation (ODE), enabling more expressive feature mappings while maintaining the computational efficiency of sparse expert selection.


## Installation

```bash
pip install mixture-of-flows
```


## Usage


```python
import torch
from mixture_of_flows import FlowMatchingMoE, FlowMatchingMoEConfig

# Define configuration
config = FlowMatchingMoEConfig(
    input_dim=512,
    hidden_dim=2048,
    num_experts=8,
    num_selected=2,
    flow_steps=10,
)

# Create model
model = FlowMatchingMoE(config)

# Forward pass
x = torch.randn(4, 128, 512)  # (batch_size, seq_len, input_dim)
output, aux_loss = model(x)

print(f"Output shape: {output.shape}")  # (4, 128, 512)
print(f"Auxiliary loss: {aux_loss.item():.6f}")
```

---

## Introduction


**Mixture of Experts (MoE)** models partition the input space among specialized sub-networks (experts), with a gating mechanism that routes each input to a subset of experts. This approach enables conditional computation, where only a fraction of model parameters are activated for any given input, allowing for significant scaling of model capacity without proportional increases in computational cost (Shazeer et al., 2017; Fedus et al., 2022).

**Flow Matching** is a simulation-free approach to training continuous normalizing flows (CNFs). Unlike traditional normalizing flows that require explicit invertibility constraints or simulation during training, flow matching learns a velocity field that transports samples from a source distribution to a target distribution along optimal transport paths (Lipman et al., 2023).

### Motivation

Traditional MoE architectures employ feed-forward networks as experts, which apply fixed, discrete transformations to inputs. By replacing these with flow matching networks, FM-MoE enables each expert to learn a continuous, time-dependent transformation governed by:

```
dx/dt = v_theta(x, t), where t in [0, 1]
```

This formulation provides several theoretical and practical advantages:

1. **Enhanced Expressiveness**: Continuous transformations can represent more complex mappings than discrete layer compositions
2. **Smooth Interpolation**: The ODE formulation ensures smooth transitions in the learned representation space
3. **Flexible Computation**: Integration steps can be adjusted at inference time for accuracy-efficiency trade-offs

---

## Architecture

```
                                    +------------------+
                                    |   Expert 0       |
                                    | (Flow Matching)  |
                                    +--------+---------+
                                             |
+----------+    +----------+    +------------------+    +----------+
|  Input   | -> |  Router  | -> |   Expert 1       | -> | Weighted |
|  Tokens  |    | (Top-k)  |    | (Flow Matching)  |    |   Sum    | -> Output
+----------+    +----------+    +--------+---------+    +----------+
                    |                    |
                    |           +------------------+
                    +---------> |   Expert N-1     |
                                | (Flow Matching)  |
                                +------------------+
```

### Component Architecture

#### 1. Router Network

The router computes expert assignment probabilities and selects the top-k experts for each token:

```
                    +-------------------------+
                    |     Linear(d -> E)      |
      Input  -----> |          |              | -----> Top-k Selection
     (B,L,d)        |     Softmax             |        Probabilities
                    +-------------------------+
```

#### 2. Flow Matching Expert

Each expert implements a velocity field network integrated via Euler method:

```
    +-------+     +---------------+     +-------+
    | x(t)  | --> | Concat with   | --> |  MLP  | --> v(x,t)
    +-------+     | time embed    |     +-------+
                  +---------------+
         |                                   |
         |        x(t+dt) = x(t) + v*dt      |
         +-----------------------------------+
                  (Euler Integration)
```

#### 3. Flow Matching Expert Internal Structure

```
+------------------------------------------------------------------+
|                    Flow Matching Expert                           |
+------------------------------------------------------------------+
|                                                                   |
|   Input (d)                                                       |
|      |                                                            |
|      v                                                            |
|   +------------------+    +----------------------+                |
|   | Sinusoidal Time  |    |                      |                |
|   | Embedding        |--->|  Concatenate         |                |
|   | (t -> d_time)    |    |  [x; time_embed]     |                |
|   +------------------+    +----------+-----------+                |
|                                      |                            |
|                                      v                            |
|                          +-----------------------+                |
|                          | Linear(d+d_time, h)   |                |
|                          +-----------------------+                |
|                                      |                            |
|                                      v                            |
|                          +-----------------------+                |
|                          | SiLU + LayerNorm      |                |
|                          +-----------------------+                |
|                                      |                            |
|                                      v                            |
|                          +-----------------------+                |
|                          | Linear(h, h)          |                |
|                          +-----------------------+                |
|                                      |                            |
|                                      v                            |
|                          +-----------------------+                |
|                          | SiLU + LayerNorm      |                |
|                          +-----------------------+                |
|                                      |                            |
|                                      v                            |
|                          +-----------------------+                |
|                          | Linear(h, d)          |                |
|                          | (zero-initialized)    |                |
|                          +-----------------------+                |
|                                      |                            |
|                                      v                            |
|                              Velocity v(x,t)                      |
|                                                                   |
+------------------------------------------------------------------+
```

---

## Mathematical Formulation

### Flow Matching Fundamentals

Flow matching learns a time-dependent velocity field `v_theta(x, t)` that defines an ordinary differential equation:

```
dx/dt = v_theta(x, t),    t in [0, 1]
```

The transformation from input `x_0` to output `x_1` is obtained by integrating this ODE:

```
x_1 = x_0 + integral from 0 to 1 of v_theta(x(t), t) dt
```

### Numerical Integration

We employ the Euler method for numerical integration with `N` steps:

```
x_{n+1} = x_n + v_theta(x_n, t_n) * dt

where:
  - dt = 1/N (step size)
  - t_n = n * dt (time at step n)
  - n in {0, 1, ..., N-1}
```

### Expert Routing

Given input `x in R^d`, the router computes expert probabilities:

```
p(e|x) = softmax(W_r * x)_e,    e in {1, ..., E}
```

The top-k experts are selected based on these probabilities:

```
S(x) = argmax_k { p(e|x) : e in {1, ..., E} }
```

### Mixture Computation

The final output combines selected expert outputs:

```
y = sum over e in S(x) of [ p_norm(e|x) * f_e(x) ]

where:
  - f_e(x) is the flow transformation of expert e
  - p_norm(e|x) = p(e|x) / sum over e' in S(x) of p(e'|x)
```

### Auxiliary Losses

#### Load Balancing Loss

Encourages uniform expert utilization:

```
L_balance = E * sum from e=1 to E of (p_bar_e - 1/E)^2

where p_bar_e = (1/BL) * sum over b,l of p(e|x_{b,l})
```

#### Router Z-Loss

Prevents large logit values for training stability:

```
L_z = (1/BL) * sum over b,l of (log(sum over e of exp(z_{b,l,e})))^2
```

### Time Embedding

Sinusoidal positional encoding for continuous time:

```
PE(t, 2i) = sin(t / 10000^(2i/d))
PE(t, 2i+1) = cos(t / 10000^(2i/d))
```

---


# More Examples

### Integration with Transformer

```python
import torch
import torch.nn as nn
from mixture_of_flows import FlowMatchingMoE, FlowMatchingMoEConfig

class TransformerBlockWithFMMoE(nn.Module):
    """Transformer block using FM-MoE instead of standard FFN."""
    
    def __init__(self, d_model: int, n_heads: int, num_experts: int = 8):
        super().__init__()
        
        # Multi-head attention
        self.attention = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm1 = nn.LayerNorm(d_model)
        
        # FM-MoE replaces the feed-forward network
        self.moe = FlowMatchingMoE(FlowMatchingMoEConfig(
            input_dim=d_model,
            hidden_dim=d_model * 4,
            num_experts=num_experts,
            num_selected=2,
            flow_steps=10,
        ))
        self.norm2 = nn.LayerNorm(d_model)
        
    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        # Self-attention with residual
        attn_out, _ = self.attention(x, x, x)
        x = self.norm1(x + attn_out)
        
        # FM-MoE with residual
        moe_out, aux_loss = self.moe(x)
        x = self.norm2(x + moe_out)
        
        return x, aux_loss
```

### Expert Usage Analysis

```python
# Get expert usage statistics
stats = model.get_expert_usage_stats(x)

print(f"Expert probabilities: {stats['expert_probs']}")
print(f"Expert selections: {stats['expert_selections']}")
print(f"Balance score: {stats['balance_score']:.4f}")  # 1.0 = perfect balance
```

### Visualization

```python
from mixture_of_flows.main import visualize_flow_matching_moe

# Generate comprehensive visualization
visualize_flow_matching_moe(
    model=model,
    x=x,
    save_path="fm_moe_analysis.png",
    show_flow_trajectory=True,
)
```

---

## API Reference

### FlowMatchingMoEConfig

Configuration dataclass for FM-MoE hyperparameters.

| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input_dim` | int | required | Dimensionality of input features |
| `hidden_dim` | int | required | Hidden dimension for flow transformations |
| `num_experts` | int | 8 | Number of flow matching experts |
| `num_selected` | int | 2 | Number of top-k experts per token |
| `flow_steps` | int | 10 | Euler integration steps |
| `time_embed_dim` | int | 64 | Time embedding dimension |
| `use_router_aux_loss` | bool | True | Enable auxiliary losses |
| `router_z_loss_coef` | float | 0.001 | Z-loss coefficient |
| `load_balance_loss_coef` | float | 0.01 | Load balancing coefficient |
| `expert_capacity_factor` | float | 1.25 | Expert capacity multiplier |
| `dropout` | float | 0.0 | Dropout probability |

### FlowMatchingMoE

Main module implementing the FM-MoE layer.

#### Methods

**`forward(x, return_aux_loss=True)`**
- **Input**: `x` - Tensor of shape `(batch_size, seq_len, input_dim)`
- **Output**: Tuple of `(output, aux_loss)` where output has same shape as input

**`get_expert_usage_stats(x)`**
- **Input**: `x` - Tensor of shape `(batch_size, seq_len, input_dim)`
- **Output**: Dictionary containing `expert_probs`, `expert_selections`, `balance_score`

### FlowMatchingExpert

Individual flow matching expert network.

#### Methods

**`forward(x, t)`**
- Computes velocity field at position `x` and time `t`

**`flow_transform(x, steps=10)`**
- Applies complete flow transformation via Euler integration

---

## Configuration

### Recommended Configurations

#### Small Model (Research/Prototyping)
```python
config = FlowMatchingMoEConfig(
    input_dim=256,
    hidden_dim=512,
    num_experts=4,
    num_selected=1,
    flow_steps=5,
)
```

#### Base Model
```python
config = FlowMatchingMoEConfig(
    input_dim=512,
    hidden_dim=2048,
    num_experts=8,
    num_selected=2,
    flow_steps=10,
)
```

#### Large Model
```python
config = FlowMatchingMoEConfig(
    input_dim=1024,
    hidden_dim=4096,
    num_experts=16,
    num_selected=2,
    flow_steps=15,
)
```

### Hyperparameter Guidelines

| Parameter | Low | Medium | High | Notes |
|-----------|-----|--------|------|-------|
| `flow_steps` | 5 | 10 | 20 | More steps = higher accuracy, more compute |
| `num_experts` | 4 | 8 | 16 | Scale with model capacity |
| `num_selected` | 1 | 2 | 4 | Trade-off: efficiency vs. expressiveness |
| `load_balance_loss_coef` | 0.001 | 0.01 | 0.1 | Increase if experts collapse |

---

## Experimental Results

### Computational Complexity

| Component | Time Complexity | Space Complexity |
|-----------|-----------------|------------------|
| Router | O(B * L * E) | O(E * d) |
| Per Expert | O(B * L * N * d * h) | O(d * h) |
| Total | O(B * L * k * N * d * h) | O(E * d * h) |

Where: B = batch size, L = sequence length, E = num experts, k = num selected, N = flow steps, d = input dim, h = hidden dim

### Memory Efficiency

Compared to dense models with equivalent expressiveness, FM-MoE achieves:

- **Active Parameters**: k/E fraction of total parameters per forward pass

- **Memory Scaling**: Sub-linear with number of experts due to sparse activation

---

## Citation

If you use this implementation in your research, please cite:

```bibtex
@software{gomez2026fmmoe,
  author       = {Gomez, Kye},
  title        = {Flow Matching Mixture of Experts: Continuous Normalizing Flows for Sparse Expert Networks},
  year         = {2026},
  publisher    = {The Swarm Corporation},
  url          = {https://github.com/The-Swarm-Corporation/MoF}
}
```

### Related Works

```bibtex
@article{lipman2023flow,
  title={Flow Matching for Generative Modeling},
  author={Lipman, Yaron and Chen, Ricky TQ and Ben-Hamu, Heli and Nickel, Maximilian and Le, Matt},
  journal={International Conference on Learning Representations},
  year={2023}
}

@article{fedus2022switch,
  title={Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity},
  author={Fedus, William and Zoph, Barret and Shazeer, Noam},
  journal={Journal of Machine Learning Research},
  volume={23},
  number={120},
  pages={1--39},
  year={2022}
}

@article{shazeer2017outrageously,
  title={Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer},
  author={Shazeer, Noam and Mirhoseini, Azalia and Maziarz, Krzysztof and Davis, Andy and Le, Quoc and Hinton, Geoffrey and Dean, Jeff},
  journal={International Conference on Learning Representations},
  year={2017}
}

@article{chen2018neural,
  title={Neural Ordinary Differential Equations},
  author={Chen, Ricky TQ and Rubanova, Yulia and Bettencourt, Jesse and Duvenaud, David K},
  journal={Advances in Neural Information Processing Systems},
  volume={31},
  year={2018}
}
```

---

## License

This project is licensed under the MIT License. See [LICENSE](LICENSE) for details.

---

## Acknowledgments

This implementation builds upon foundational work in mixture of experts architectures and continuous normalizing flows. We acknowledge the contributions of the broader research community in advancing these methodologies.

---

## Author

**Kye Gomez**  

The Swarm Corporation  

[GitHub](https://github.com/The-Swarm-Corporation)