togethercomputer/Aurora-Spec-Minimax-M2.1

Aurora-Spec-Minimax-M2.1

Model Description

This is an EAGLE3 draft model trained from scratch (random initialization) using the Aurora inference-time training framework for speculative decoding. Unlike traditional approaches that fine-tune pre-trained models, this model is built entirely through Aurora's online training process. The model is optimized to generate high-quality draft tokens for the MiniMax M2.1 target model, achieving significant speedups across various batch sizes.

Key Features

• Training Approach: Trained from scratch (random initialization) - no pre-training required
• Framework: Trained with Aurora - an advanced inference-time training system
• Architecture: EAGLE3 speculative decoding draft model
• Target Model: MiniMax M2.1
• Performance: Achieves 2.62 average accept length with lookahead 4 (recommended configuration)
• Training: 44,000 inference requests on NVIDIA H200 GPU
• Speedup: Up to 1.58× speedup at batch size 1 (lookahead 3), 1.57× with lookahead 4 (recommended)

Target Model

This draft model is specifically designed to work with:

• Model: MiniMax M2.1
• Type: General-purpose language model
• Domain: Broad language understanding and generation

The draft model learns to predict the target model's token distribution during inference-time training, enabling efficient speculative decoding.

Architecture

EAGLE3 Speculative Decoding

This model implements the EAGLE3 (Extrapolation Algorithm for Greater Language-model Efficiency) architecture:

• Draft Model: Lightweight model that generates candidate tokens
• Tree-based Attention: Enables parallel verification of multiple draft tokens
• Auto-regressive Generation: Produces speculative token sequences
• Dynamic Adaptation: Updates during inference to match target model distribution

Model Structure

• Initialization: Trained from scratch (random initialization, no pre-training)
• Base Architecture: Single-layer Transformer decoder
• Recommended Configuration: Lookahead 4 (speculative_num_steps=4)
• Attention Mechanism: Tree-based for parallel draft verification
• Training Paradigm: Online learning during inference (Aurora framework)

Training Details

Aurora Framework

This model was trained from scratch using Aurora, an inference-time training framework that:

• No Pre-training Required: Starts from random initialization and learns entirely through online training
• Updates the draft model dynamically during inference
• Uses reverse KL divergence for distribution matching (minimizing KL(target || draft))
• Employs online learning with periodic model updates
• Optimizes for both draft quality and speculative acceptance rate
• Demonstrates that effective draft models can be built from scratch without expensive pre-training

Training Configuration

• Hardware: NVIDIA H200 GPU
• Training Requests: 44,000 inference requests
• Synchronization Interval: Every 800 requests
• Recommended Configuration: Lookahead 4
• KL Divergence: Reverse KL divergence (draft → target)

Dataset

Trained on diverse prompts suitable for general-purpose language modeling and speculative decoding.

Benchmarks

End-to-End Throughput Performance

Measured on a holdout evaluation dataset using the final Aurora checkpoint.

MiniMax M2.1: end-to-end throughput under varying batch size and lookahead

We report tokens-per-second (TPS) statistics and speedup relative to the no-speculation baseline.

| BS | Config | Mean TPS | P50 TPS | P05 TPS | P95 TPS | Count | Speedup | Acc Len | |:---:|:---------|:--------:|:-------:|:-------:|:-------:|:-----:|:-------:|:-------:| | 1 | w/o spec | 134.9 | 136.4 | 130.6 | 136.9 | 257 | -- | -- | | | lookahead 3 | 213.0 | 213.7 | 169.8 | 256.3 | 257 | 1.58× | 2.42 | | | lookahead 4 | 211.8 | 210.6 | 163.1 | 270.3 | 257 | 1.57× | 2.62 | | 8 | w/o spec | 79.0 | 78.7 | 73.7 | 85.1 | 257 | -- | -- | | | lookahead 3 | 106.5 | 105.2 | 84.0 | 134.8 | 257 | 1.35× | 2.43 | | | lookahead 4 | 107.1 | 104.5 | 79.9 | 137.1 | 257 | 1.36× | 2.62 | | | lookahead 5 | 106.6 | 104.8 | 79.3 | 140.9 | 257 | 1.35× | 2.70 | | 16 | w/o spec | 64.5 | 63.7 | 58.9 | 72.3 | 257 | -- | -- | | | lookahead 3 | 83.2 | 81.4 | 62.2 | 110.3 | 257 | 1.29× | 2.43 | | | lookahead 4 | 83.1 | 82.9 | 60.9 | 112.0 | 257 | 1.29× | 2.62 | | | lookahead 5 | 82.6 | 81.0 | 58.1 | 116.1 | 257 | 1.28× | 2.69 | | 32 | w/o spec | 53.5 | 52.9 | 47.1 | 67.1 | 257 | -- | -- | | | lookahead 3 | 67.1 | 64.9 | 45.2 | 97.8 | 257 | 1.25× | 2.44 | | | lookahead 4 | 67.1 | 64.7 | 44.0 | 100.5 | 257 | 1.25× | 2.62 | | | lookahead 5 | 67.3 | 64.9 | 45.2 | 99.7 | 257 | 1.26× | 2.71 |

Performance Across Different Batch Sizes

Aurora provides consistent speedups across all batch sizes for MiniMax M2.1, demonstrating the effectiveness of speculative decoding across diverse deployment scenarios:

• Batch Size 1 (Best Case): Up to 1.58× speedup with lookahead 3 configuration. The recommended lookahead 4 achieves 1.57× speedup with 2.62 average accept length. At low batch sizes, the cost of draft generation and verification is well amortized by reduced target model forward passes, providing the largest gains for latency-critical scenarios.

• Batch Size 8 (Strong): 1.36× speedup with lookahead 4 configuration (2.62 average accept length). Speculative decoding continues to provide substantial throughput improvements for moderate batching scenarios.

• Batch Size 16 (Moderate): 1.29× speedup with lookahead 4 configuration (2.62 average accept length). Benefits remain significant as the verification overhead is effectively managed.

• Batch Size 32 (Consistent): 1.25-1.26× speedup with lookahead 4-5 configurations. Unlike some models, MiniMax M2.1 maintains positive speedups even at large batch sizes, demonstrating robust performance across the batching spectrum.

Metrics Explained:

• TPS: Tokens per second (throughput)
• Acc Len: Average accept length (number of draft tokens accepted per verification step)
• Speedup: Relative to the no-speculation baseline
• P05/P95: 5th and 95th percentile throughput values
• Count: Number of evaluation samples

Notably, this performance is…