ADu2021/arbitrage-advantage-speculation

Overview

Arbitrage improves speculative decoding by introducing a lightweight router trained to identify when the target model will produce meaningfully superior reasoning steps. Rather than fixed acceptance thresholds, the framework dynamically routes generation, achieving near-optimal efficiency-accuracy tradeoffs.

When to Use

• Reasoning tasks with lengthy chain-of-thought processes
• Inference cost reduction without quality loss
• Scenarios with heterogeneous model capabilities
• Mathematical and logical reasoning problems
• Need for dynamic quality-efficiency tradeoff

When NOT to Use

• Simple single-step generation
• Strictly real-time latency-critical applications
• Cases where draft models are unavailable
• Scenarios where all steps need identical quality

Core Technique

Dynamic routing via learned advantage estimation:

# Arbitrage: Dynamic routing based on model advantage
class ArbitrageRouter:
    def __init__(self, draft_model, target_model):
        self.draft = draft_model
        self.target = target_model
        # Lightweight router predicting when target exceeds draft
        self.router = nn.Sequential(
            nn.Linear(hidden_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )

    def predict_target_advantage(self, state):
        """Router predicts probability target model is superior."""
        # Extract features from current state
        features = self.extract_state_features(state)
        # Predict target advantage
        prob_target_better = self.router(features)
        return prob_target_better

    def generate_with_dynamic_routing(self, prompt, num_steps):
        """Route generation: draft vs target based on predicted advantage."""
        state = prompt
        total_steps = 0

        for step in range(num_steps):
            # Predict if target model worth using
            target_advantage = self.predict_target_advantage(state)

            if target_advantage > 0.5:
                # Use target model for this step
                next_step = self.target.generate_step(state)
            else:
                # Use draft model (faster)
                next_step = self.draft.generate_step(state)

            state = state + next_step
            total_steps += 1

        return state

    def train_router(self, trajectories):
        """Train router on step-level comparison data."""
        for trajectory in trajectories:
            for step_idx, (state, action) in enumerate(trajectory):
                # Compute target advantage for this state
                draft_output = self.draft.generate_step(state)
                target_output = self.target.generate_step(state)

                # Semantic verification: which is better?
                advantage = self.compute_semantic_advantage(
                    draft_output,
                    target_output,
                    trajectory[step_idx+1:]  # future trajectory
                )

                # Train router to predict advantage
                features = self.extract_state_features(state)
                pred_advantage = self.router(features)

                loss = torch.nn.functional.mse_loss(
                    pred_advantage,
                    torch.tensor([advantage])
                )
                loss.backward()

        self.optimizer.step()

Key Results

• 2× latency reduction at matched accuracy
• Consistent improvements across mathematical reasoning
• Intelligent model selection per step

References

• Original paper: https://arxiv.org/abs/2512.05033
• Focus: Efficient reasoning through dynamic routing
• Domain: Inference optimization, speculative decoding