AdaptiveRiverLM (also known as ROE - River of Experts) is an experimental hybrid architecture combining State Space Models (Mamba) with Mixture-of-Experts (MoE) layers for efficient and adaptive language modeling.

Model Status

⚠️ Experimental Release ⚠️

This model is a test model with a single pass through a collection of synthetic mathematics datasets at a budget of 1. Further training runs are to be performed soon should the architecture prove resourceful.

Extensive tests have not been performed yet.

Architecture Overview

ROE (River of Experts) / AdaptiveLM Architecture

The main theory behind the ROE architecture is made to target a range of setups while preserving the performance of a model of the same size or higher.

By preserving complexity with multiple layers of experts, we allow for higher performance in both token generation and model prediction.

Three-Zone Design

┌─────────────────────────────────────────┐
│  Zone 1: Early Mamba (Layers 0-1)       │
│  • Fast sequence processing             │
│  • O(n) complexity                      │
│  • d_state=16, expand=2                 │
├─────────────────────────────────────────┤
│  Zone 2: MoE Layers (Layers 2-21)       │
│  • Conditional computation              │
│  • MoE Attention: 6 experts             │
│  • MoE FFN: 4 experts                   │
│  • Dynamic expert routing               │
├─────────────────────────────────────────┤
│  Zone 3: Enhanced Mamba (Layers 22-23)  │
│  • High-capacity refinement             │
│  • d_state=16, expand=4 (2× capacity)   │
└─────────────────────────────────────────┘

Key Features

Hybrid Architecture: Combines Mamba State Space Models (SSM) with Mixture-of-Experts
24 Layers: 2 early Mamba + 20 MoE + 2 enhanced Mamba
~1B Parameters: Efficient parameter usage through sparse activation
Budget-Aware Inference: Runtime control over active experts (speed vs. quality tradeoff)
Stateless Generation: Simplified deployment without KV cache requirements

Model Specifications

Parameter	Value
Total Layers	24
Hidden Size (d_model)	1024
FFN Hidden Size	4096
Vocabulary Size	50,257
Attention Experts	6 per MoE layer
FFN Experts	4 per MoE layer
Mamba d_state	16
Mamba d_conv	4
Position Encoding	RoPE (Rotary Position Embedding)

Requirements

bash

pip install torch transformers
pip install mamba-ssm  # Required for Mamba layers

Note: The mamba-ssm package is required for the model to function. Without it, Mamba layers will be non-functional.

Usage

python inference_tester.py --model_dir /path/to/adaptiveriverlm --interactive

Architecture Details

Mamba Blocks

Early and late Mamba layers use selective state space models for efficient O(n) sequence processing:

Early Layers (0-1): Standard capacity (expand=2)
Late Layers (22-23): Enhanced capacity (expand=4)

MoE Layers

Middle layers use conditional computation with dynamic expert routing:

MoE Attention:

6 expert attention heads per layer
Dynamic top-k selection (typically 4-6 active)
Per-expert Q/K/V projections with RoPE
Router network for content-based selection

MoE FFN:

4 expert feed-forward networks per layer
Token-level routing with top-k gating
Straight-through estimator for differentiable routing
Load balancing to ensure expert utilization

Training Objectives

The model is trained with multiple auxiliary losses:

Load Balancing Loss: Ensures even expert utilization
Router Z-Loss: Prevents logit magnitude explosion
Entropy Regularization: Encourages diverse expert selection

Adaptive Expert Selection

One of the key innovations of AdaptiveRiverLM is its ability to maintain strong performance while dynamically adjusting the number of active attention experts. This allows the model to adapt to different computational budgets and deployment scenarios without requiring separate model checkpoints.

How Expert Scaling Works

The model's MoE attention layers contain 6 expert heads each, but not all experts need to be active for every input. The router network intelligently selects which experts to activate based on:

Input Content: Content-based routing determines which experts are most relevant
Budget Ratio: User-defined parameter controlling the expert activation range (0.0 to 1.0)

Expert Activation Formulas

The model uses different scaling strategies for attention and FFN experts:

python

# Attention expert selection (more aggressive scaling)
k_attention = max(1, int(round(top_k * (0.25 + 0.75 * budget_ratio))))
# Example: top_k=6, budget_ratio=0.5 → k_target = 3 experts (50% active)
# Example: top_k=6, budget_ratio=0.8 → k_target = 5 experts (83% active)

# FFN expert selection (conservative scaling)
k_ffn = max(1, int(round(base_top_k * (0.5 + budget_ratio / 2.0))))
# Example: base_top_k=2, budget_ratio=0.5 → k_target = 1-2 experts
# Example: base_top_k=2, budget_ratio=1.0 → k_target = 2 experts

Key Insight: Attention experts scale more aggressively (25-100%) while FFN experts scale conservatively (50-100%), as attention routing has been found to be more critical for maintaining quality.

Performance Characteristics by Budget

Budget Ratio	Active Attn Experts	Active FFN Experts	Relative Speed	Quality Retention	Recommended Use Case
1.0 (Full)	6/6 (100%)	1/4 (25%)	1.0×	100%	Maximum quality, complex reasoning
0.9	5-6/6 (83-100%)	1/4 (25%)	~1.1×	95-98%	High-quality production
0.75	4-5/6 (67-83%)	1/4 (25%)	~1.4×	90-95%	Balanced performance
0.6	4/6 (67%)	1/4 (25%)	~1.7×	85-90%	Efficient inference
0.5	3/6 (50%)	1/4 (25%)	~2.0×	80-85%	Fast generation, good quality
0.35	2-3/6 (33-50%)	1/4 (25%)	~2.3×	70-80%	Speed-optimized
0.25	2/6 (33%)	1/4 (25%)	~2.5×	60-75%	Minimal mode, basic tasks

Important Notes:

Quality retention percentages are task-dependent (simple tasks degrade less)
Speed improvements are approximate and vary by hardware
The model uses sparse activation, so even at full budget, not all parameters are active

Why This Matters

Graceful Degradation: Unlike traditional transformers that operate at fixed capacity, AdaptiveRiverLM provides smooth quality/speed tradeoffs:

Budget = 1.0: Full model capacity, all 6 attention experts available per layer
Budget = 0.5: 50% fewer active attention parameters while maintaining 80-85% performance
Budget = 0.25: Minimal mode (33% experts) suitable for simple queries or edge deployment

Preserved Complexity: Even at lower budgets, the model maintains architectural richness through:

Expert Specialization: Different experts learn complementary skills during training
Intelligent Routing: Most relevant experts are activated first (content-aware selection)
Hybrid Design: Mamba layers provide stable base performance regardless of budget
Residual Connections: Information bypassing when experts aren't needed

Real-World Benefits:

Deploy one checkpoint across different hardware (high-end server to mobile edge)
Dynamically adjust compute based on query complexity or available resources
Batch processing with mixed budgets for different priority levels
Cost optimization: use lower budgets for simple queries, full budget for complex reasoning
A/B testing: compare quality degradation vs. speed gains for your specific use case

Technical Implementation Details

The router architecture includes:

Temperature-scaled Gating: Controlled via gate_temperature=0.7 for smooth probability distributions
Straight-Through Estimator (STE): Enables differentiable top-k selection during training
Auxiliary Losses:
- Load Balancing Loss: ((usage - uniform)² ).sum() prevents expert collapse
- Router Z-Loss: (logits²).mean() prevents magnitude explosion
- Entropy Regularization: Encourages diverse expert utilization
Top-k Masking: Hard selection with soft backpropagation via STE

Performance Validation

Expected Behavior:

Tasks requiring broad knowledge benefit more from high budgets (budget ≥ 0.7)
Narrow, specialized tasks show minimal degradation even at budget = 0.5
Simple pattern matching (arithmetic, templates) works well at budget ≥ 0.3
The Mamba layers (zones 1 and 3) provide stable performance regardless of MoE budget

Recommended Testing:

python

# Benchmark across budgets for your specific use case
for budget in [1.0, 0.75, 0.5, 0.25]:
    results = evaluate_model(test_set, budget_ratio=budget)
    print(f"Budget {budget}: Accuracy={results.accuracy}, Speed={results.tokens_per_sec}")

This adaptive mechanism is what allows AdaptiveRiverLM to maintain strong performance even when constrained to fewer active experts, making it particularly suitable for production deployments with varying resource availability, cost constraints, or latency requirements.

Training Data

This initial release was trained on a collection of synthetic mathematics datasets, focusing on:

Arithmetic reasoning
Algebraic problem solving
Mathematical expressions
Step-by-step solutions

Performance Notes

Current Status: This is an experimental architecture test. Comprehensive benchmarking is ongoing.

Limitations

Experimental: Architecture and training are still being refined
Limited Training: Single-pass on synthetic math data only
No KV Cache: Current implementation uses stateless generation
Untested at Scale: Performance characteristics under investigation

Roadmap

Complete comprehensive benchmark suite
Multi-pass training with diverse datasets
Implement efficient KV caching for MoE attention
Optimize expert routing strategies
Scale to larger model sizes (3B, 7B variants)

License and Use

This model is licensed under Apache 2.0. It is intended for research and educational use.

Note: This is an experimental model release. Results may vary significantly from production-ready models. Use with appropriate caution and validation.

Acknowledgments

Mamba: State Space Model architecture from Mamba paper
Mixture-of-Experts: Inspired by modern MoE architectures
Built with PyTorch and Hugging Face Transformers

Alienanthony/ROE_EDU_BASE_Undercooked

AdaptiveRiverLM-1B