Chain-of-Experts (CoE)

Chain-of-Experts (CoE) is a novel approach for sparse Large Language Model (LLM) processing that replaces the standard parallel expert processing found in Mixture-of-Experts (MoE) models with sequential communication.

Instead of processing tokens independently in parallel, CoE implements an iterative mechanism where the output of one expert iteration serves as the input for the next. This allows experts to "communicate" and accumulate information over multiple iterations within a single layer, improving performance and memory efficiency compared to traditional MoE architectures.

Experimental methods in this project are denoted using the notation ${\rm{CoE/MoE}} (C, K, N)$.

• C: Number of iteration cycles.
• K: Number of experts selected per iteration.
• N: Total number of experts per layer.

Example: CoE (C=2, K=4, N=64) means using CoE with 2 iteration cycles, selecting 4 experts each time, from a total pool of 64 experts per layer.

CoE relies on two critical architectural components for performance:

1. Independent Gating

Instead of using a shared gating mechanism, use independent gating $g_{t,i}$ for each iteration $t$. Shared gating ($g_i$) results in worse performance than standard MoE.

2. Inner Residual Connections

Use inner residual connections rather than outer ones. Inner residuals allow the model to treat sequential processing as an increase in effective depth, where each expert learns to process the residual of the previous step.

Correct Implementation (Inner Residual + Independent Gating): $x^{(0)} = x$ $x^{(t)} = \sum_{i=1}^{K/C} g_{t,i} \cdot \text{FFN}_i(x^{(t-1)}) + I_r \cdot x^{(t-1)}, \quad t = 1, 2, ..., C$ $y = x^{(C)}$

CoE utilizes an independent gating mechanism where each iteration $t$ uses its own router parameters $u_{t,i}$ and token embeddings $e_{t,i}$ to determine expert weights. This allows the model to re-evaluate routing decisions based on the refined information from previous iterations.

Gating Logic: The gate value $g_{t,i}$ for expert $i$ at iteration $t$ is defined as: $g_{t,i} = \begin{cases} s_{t,i}, & s_{t,i} \in \text{TopK}(s_{t,j}|1 \leq j \leq K/C, N), \ 0, & \text{otherwise}, \end{cases}$

where the softmax score is: $s_{t,i} = \text{Softmax}(u_{t,i}^\top e_{t,i})$

The experimental method follows a specific notation: ${\rm{CoE/MoE}}-C(K/N)$.

• C: The number of iteration cycles.
• K: The number of experts selected in each iteration.
• N: The total number of experts available per layer.

Example: CoE-2(4/64) means using CoE with 2 iteration cycles, selecting 4 experts each time, from a total of 64 experts per layer.

To replicate the experimental training setup, use the following hyperparameters:

• Optimizer: AdamW
• Learning Rate: 3e-4
• Learning Rate Schedule: 10% Warmup
• Betas: [0.9, 0.95]
• Weight Decay: 0.01
• Warmup Steps Ratio: 0.1
• Gradient Clipping: 1.0
• Batch Size: 64
• Sequence Length: 512
• Training Steps: 1000

Note: The implementation uses a modified version of the veRL architecture.

You can run the basic experiments using the provided shell scripts.

• To run the latest experiments: bash runs/run_latest.sh
• To run the specific performance comparison experiment (CoE vs MoE on Math tasks): bash runs/run.sh

bash runs/run_latest.sh # run the basic experiments

bash runs/run.sh # run the performance comparison experiment

To prepare your environment for running CoE experiments, execute the provided setup script. This script handles the necessary dependencies and environment configurations.

bash scripts/setup.sh # setting up environment

When implementing or configuring a CoE model based on the provided experimental setup, use the following parameter groups:

Basic Model Parameters

• hidden_size: 1024
• num_hidden_layers: 4
• num_attention_heads: 8

MoE Parameters

• total_routed_experts: 63
• shared_experts: 1
• experts_selected_per_token: 8 (routed) + 1 (shared)
• moe_intermediate_size: 704
• moe_layer_frequency: 1 (every layer is an MoE layer)

CoE Specific Parameters

• num_iterations: 2
• residual_connection_approach: Inner (Inner residual connections are more effective than outer ones)
• independent_gating: Enabled (Using independent gating $g_{t,i}$ per iteration is more effective than shared gating $g_i$)

The following configuration represents the baseline model used in the research (based on DeepSeek-V2-Lite architecture).

### Basic model parameters:
- Hidden size: 1024
- Number of hidden layers: 4
- Number of attention heads: 8

### MoE related parameters:
- Total routed experts: 63
- Shared experts: 1
- Experts selected per token: 8 (routed) + 1 (Shared)
- MoE intermediate size: 704
- MoE layer frequency: 1 (every layer is an MoE layer)

### CoE specific parameters:
- Number of iteration: 2
- Residual connection approach: Inner
- Independent gating: Enabled

Standard training hyperparameters used in the experiments:

Batch size: 64
Sequence length: 512
Training steps: 1000-10000
Optimizer: AdamW
Learning rate: 3e-4
Learning rate schedule: 10% Warmup
betas: [0.9, 0.95]
weight_decay: 0.01
warmup_steps_ratio: 0.1
clip_grad: 1.0