Poss Speculative Decoding

---
name: poss-speculative-decoding
title: "PosS: Position Specialist Layers for Improved Speculative Decoding"
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: "https://arxiv.org/abs/2506.03566"
keywords: [speculative-decoding, inference-acceleration, draft-models]
description: "Improve speculative decoding throughput by employing position-specialized draft layers that handle position-specific error accumulation patterns."
---

# PosS: Position Specialist Generates Better Draft for Speculative Decoding

## Core Concept

Speculative decoding accelerates LLM inference by using a smaller draft model to predict multiple tokens, then verifying them with the larger target model. PosS identifies a critical limitation: draft models suffer from degraded token quality at later positions due to error accumulation. By deploying multiple position-specialized layers—each handling a narrow, predictable range of feature deviation—PosS improves acceptance length and speed-up ratios by up to 5.7%.

## Architecture Overview

- **Position-Specific Feature Deviation**: Each position in a draft has characteristic feature deviation from the target model. Later positions accumulate errors multiplicatively.
- **Position Specialists**: Separate transformer layers trained to handle specific positions with their expected deviation levels
- **Position-Wise Acceptance Rate (Pos-Acc)**: New metric analyzing draft quality across different positions in the prediction sequence
- **Chain Rule Decomposition**: Overall acceptance depends on multiplying individual position acceptance rates; fixing weak positions improves end-to-end throughput
- **Training with Simulated Error**: Specialists learn from features generated by predecessors, simulating inference conditions

## Implementation

### Step 1: Analyze Position-Specific Acceptance Degradation

```python
import torch
import numpy as np
from typing import List, Dict
from collections import defaultdict

class PositionAcceptanceAnalyzer:
    """Diagnose where draft models fail in speculative decoding"""

    def __init__(self, draft_model, target_model):
        self.draft_model = draft_model
        self.target_model = target_model

    def compute_position_wise_acceptance(self,
                                        prompts: List[torch.Tensor],
                                        num_draft_tokens: int = 4) -> Dict:
        """
        For each position k in the draft, measure acceptance rate.
        Key finding: pos-acc rapidly deteriorates beyond k=1.
        """

        position_stats = defaultdict(list)

        for prompt in prompts:
            # Generate multiple draft tokens
            draft_tokens = self.draft_model.generate_tokens(
                prompt, num_tokens=num_draft_tokens
            )

            # For each position, check if target model accepts
            for pos in range(num_draft_tokens):
                draft_token_at_pos = draft_tokens[pos]

                # Target model's prediction at this position
                target_logits = self.target_model.forward(
                    torch.cat([prompt, draft_tokens[:pos]], dim=0)
                )
                target_top_token = torch.argmax(target_logits[-1, :])

                # Check acceptance
                accepted = (draft_token_at_pos == target_top_token)
                position_stats[pos].append(float(accepted))

        # Aggregate statistics
        position_acceptance_rates = {}

        for pos in range(num_draft_tokens):
            acc_rate = np.mean(position_stats[pos])
            position_acceptance_rates[pos] = acc_rate

            print(f"Position {pos}: {acc_rate:.1%} acceptance rate")

        # Compute cumulative impact via chain rule
        overall_acceptance = np.prod(list(position_acceptance_rates.values()))

        print(f"\nOverall acceptance (pos[0] × pos[1] × ...): {overall_acceptance:.1%}")

        return position_acceptance_rates

    def analyze_feature_deviation(self, prompts: List[torch.Tensor]) -> Dict:
        """
        Measure feature deviation at each position.
        Later positions have higher feature deviation due to error accumulation.
        """

        position_deviations = defaultdict(list)

        for prompt in prompts:
            # Run draft and target forward passes
            draft_hidden = self.draft_model.extract_hidden_states(prompt)
            target_hidden = self.target_model.extract_hidden_states(prompt)

            # Measure feature deviation at each generation step
            for pos in range(len(draft_hidden) - 1):
                deviation = torch.norm(
                    draft_hidden[pos] - target_hidden[pos]
                )
                position_deviations[pos].append(deviation.item())

        # Average deviation per position
        avg_deviations = {}

        for pos in range(len(draft_hidden) - 1):
            avg_dev = np.mean(position_deviations[pos])
            avg_deviations[pos] = avg_dev

            print(f"Position {pos}: avg feature deviation = {avg_dev:.4f}")

        return avg_deviations
```

### Step 2: Design Position-Specialist Architecture

```python
class PositionSpecialist(torch.nn.Module):
    """Transformer layer specialized for a specific position"""

    def __init__(self, hidden_dim: int, num_layers: int,
                 position_id: int, expected_deviation: float):
        super().__init__()

        self.position_id = position_id
        self.expected_deviation = expected_deviation

        # Adapt model size based on expected deviation
        # Higher deviation = need more capacity
        adapted_layers = max(2, int(num_layers * expected_deviation / 0.5))

        self.specialist_layers = torch.nn.ModuleList([
            self.TransformerBlock(hidden_dim)
            for _ in range(adapted_layers)
        ])

        # Loss components for training
        self.token_loss_weight = 0.7
        self.feature_loss_weight = 0.2
        self.topk_loss_weight = 0.1

    class TransformerBlock(torch.nn.Module):
        def __init__(self, hidden_dim):
            super().__init__()
            self.self_attn = torch.nn.MultiheadAttention(
                embed_dim=hidden_dim, num_heads=8
            )
            self.ffn = torch.nn.Sequential(
                torch.nn.Linear(hidden_dim, 4 * hidden_dim),
                torch.nn.ReLU(),
                torch.nn.Linear(4 * hidden_dim, hidden_dim),
            )
            self.norm1 = torch.nn.LayerNorm(hidden_dim)
            self.norm2 = torch.nn.LayerNorm(hidden_dim)

        def forward(self, x):
            x = x + self.self_attn(x, x, x)[0]
            x = self.norm1(x)
            x = x + self.ffn(x)
            x = self.norm2(x)
            return x

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        """Process hidden state through specialist layers"""

        x = hidden_state

        for layer in self.specialist_layers:
            x = layer(x)

        return x

class PositionSpecialistEnsemble(torch.nn.Module):
    """Multiple position specialists covering all draft positions"""

    def __init__(self, base_model, num_specialists: int = 4):
        super().__init__()

        self.num_specialists = num_specialists
        self.specialists = torch.nn.ModuleList()

        # Create specialist for each position
        position_deviations = self.estimate_position_deviations(base_model)

        for pos in range(num_specialists):
            specialist = PositionSpecialist(
                hidden_dim=base_model.hidden_dim,
                num_layers=base_model.num_layers,
                position_id=pos,
                expected_deviation=position_deviations.get(pos, 0.5)
            )
            self.specialists.append(specialist)

    def estimate_position_deviations(self, base_model) -> Dict:
        """Estimate expected feature deviation per position"""

        deviations = {}

        # Empirical: deviation grows with position
        for pos in range(self.num_specialists):
            # Approximation: exponential growth
            deviations[pos] = 0.1 * (1.5 ** pos)

        return deviations

    def forward(self, hidden_states: torch.Tensor,
               position: int) -> torch.Tensor:
        """
        Route hidden state through appropriate specialist.
        If position >= num_specialists, use last specialist.
        """

        specialist_idx = min(position, self.num_specialists - 1)
        specialist = self.specialists[specialist_idx]

        output = specialist(hidden_states)

        return output
```

### Step 3: Training Position Specialists

```python
class PositionSpecialistTrainer:
    """Train specialists using simulated inference conditions"""

    def __init__(self, target_model, specialist_ensemble: PositionSpecialistEnsemble):
        self.target_model = target_model
        self.specialists = specialist_ensemble
        self.optimizer = torch.optim.AdamW(
            self.specialists.parameters(), lr=2e-4
        )

    def train_on_batch(self, prompts: torch.Tensor,
                      completions: torch.Tensor) -> Dict[str, float]:
        """
        Train specialists on data with three loss components:
        1. Token-level: predict correct next token
        2. Feature-level: match target model features
        3. Top-K: preserve top-K token distribution
        """

        losses = {}

        # Extract target model's hidden states for supervision
        target_hidden = self.target_model.extract_hidden_states(
            torch.cat([prompts, completions], dim=1)
        )
        target_logits = self.target_model.forward(
            torch.cat([prompts, completions], dim=1)
        )

        total_loss = 0

        # Train each position specialist
        for pos in range(self.specialists.num_specialists):
            # Simulate inference at this position:
            # Use specialist for previous positions, then current position
            if pos == 0:
                current_hidden = target_hidden[-1, :].unsqueeze(0)
            else:
                # Use specialist output from previous position
                prev_output = self.specialists.forward(
                    target_hidden[-2, :].unsqueeze(0),
                    position=pos - 1
                )
                current_hidden = prev_output

            # Run through specialist
            specialist_output = self.specialists.forward(
                current_hidden,
                position=pos
            )

            # Compute three-part loss
            token_loss = self.compute_token_loss(
                specialist_output,
                target_logits[pos],
            )

            feature_loss = self.compute_feature_loss(
                specialist_output,
                target_hidden[pos],
            )

            topk_loss = self.compute_topk_loss(
                specialist_output,
                target_logits[pos],
                k=5
            )

            position_loss = (
                0.7 * token_loss +
                0.2 * feature_loss +
                0.1 * topk_loss
            )

            total_loss += position_loss
            losses[f'pos_{pos}'] = position_loss.item()

        # Backward pass
        self.optimizer.zero_grad()
        total_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.specialists.parameters(), 1.0)
        self.optimizer.step()

        losses['total'] = total_loss.item()

        return losses

    def compute_token_loss(self, specialist_logits: torch.Tensor,
                          target_logits: torch.Tensor) -> torch.Tensor:
        """Cross-entropy loss on predicted token"""

        return torch.nn.functional.cross_entropy(
            specialist_logits.view(-1, specialist_logits.size(-1)),
            target_logits.argmax(-1)
        )

    def compute_feature_loss(self, specialist_hidden: torch.Tensor,
                            target_hidden: torch.Tensor) -> torch.Tensor:
        """L2 loss on hidden state features"""

        return torch.nn.functional.mse_loss(
            specialist_hidden,
            target_hidden.detach()
        )

    def compute_topk_loss(self, specialist_logits: torch.Tensor,
                         target_logits: torch.Tensor,
                         k: int = 5) -> torch.Tensor:
        """KL divergence on top-K token probabilities"""

        specialist_probs = torch.nn.functional.softmax(specialist_logits, dim=-1)
        target_probs = torch.nn.functional.softmax(target_logits, dim=-1)

        # Keep only top-K for divergence
        _, topk_indices = torch.topk(target_probs, k)

        specialist_topk = specialist_probs[topk_indices]
        target_topk = target_probs[topk_indices]

        return torch.nn.functional.kl_div(
            torch.log(specialist_topk + 1e-8),
            target_topk.detach(),
            reduction='batchmean'
        )
```

### Step 4: Integration with Speculative Decoding

```python
class SpeculativeDecodingWithPosS:
    """Enhanced speculative decoding using position specialists"""

    def __init__(self, target_model, draft_model,
                 position_specialists: PositionSpecialistEnsemble):
        self.target_model = target_model
        self.draft_model = draft_model
        self.specialists = position_specialists

    def generate_with_verification(self, prompt: torch.Tensor,
                                   max_length: int = 100,
                                   num_draft_tokens: int = 4) -> torch.Tensor:
        """
        Speculative decoding with position-specialized draft:
        1. Use specialist-enhanced draft to predict k tokens
        2. Verify with target model
        3. Accept accepted tokens, resample rejected ones
        """

        generated = prompt.clone()
        total_draft_tokens = 0
        total_verified_tokens = 0

        while generated.shape[0] < max_length:
            # Generate draft tokens with specialist enhancement
            draft_tokens = []

            for pos in range(num_draft_tokens):
                # Get base draft prediction
                draft_logits = self.draft_model.forward(generated)

                # Enhance with position specialist
                draft_hidden = self.draft_model.extract_hidden_state(generated)
                enhanced_hidden = self.specialists.forward(
                    draft_hidden, position=pos
                )

                # Blend enhanced and base predictions
                enhanced_logits = self.draft_model.decode(enhanced_hidden)
                blended_logits = 0.6 * enhanced_logits + 0.4 * draft_logits

                draft_token = torch.argmax(blended_logits)
                draft_tokens.append(draft_token)

                total_draft_tokens += 1

            # Verify with target model
            draft_sequence = torch.cat([generated, torch.stack(draft_tokens)])

            target_logits = self.target_model.forward(draft_sequence)

            # Check each drafted token
            accepted_count = 0

            for pos, draft_token in enumerate(draft_tokens):
                target_token = torch.argmax(target_logits[-(len(draft_tokens)-pos)])

                if draft_token == target_token:
                    generated = torch.cat([generated, draft_token.unsqueeze(0)])
                    accepted_count += 1
                    total_verified_tokens += 1
                else:
                    # Rejection: sample from target distribution
                    target_probs = torch.softmax(target_logits[-(len(draft_tokens)-pos)], dim=-1)
                    new_token = torch.multinomial(target_probs, 1)
                    generated = torch.cat([generated, new_token])
                    total_verified_tokens += 1
                    break

        # Compute metrics
        acceptance_length = total_verified_tokens / total_draft_tokens
        speedup_ratio = total_verified_tokens / (total_verified_tokens / num_draft_tokens + total_verified_tokens)

        print(f"Acceptance length: {acceptance_length:.2f}")
        print(f"Speedup ratio: {speedup_ratio:.2f}×")

        return generated
```

## Practical Guidance

1. **Measure Pos-Acc First**: Profile your draft model with position-wise acceptance analysis. Identify which positions have lowest acceptance rates—these are bottlenecks.

2. **Specialist Count**: Start with 4-8 specialists covering 4-8 draft positions. More specialists provide finer granularity but increase computational overhead.

3. **Feature Deviation Estimation**: Position deviation grows roughly exponentially. Allocate more parameters to later positions which have higher deviation.

4. **Training Strategy**: Train specialists with simulated inference conditions using previous specialist outputs, not clean ground truth.

5. **Three-Part Loss**: Use token-level (70%), feature-level (20%), and top-K (10%) losses. This balances correctness with feature alignment.

6. **Integration**: Enhance draft model predictions by routing through specialists, then blend enhanced logits with base draft logits (60-40 weighting works well).

## Reference

- Paper: PosS (2506.03566)
- Key Metric: Position-wise acceptance rate (pos-acc) analysis
- Improvements: 4.5% on acceptance length, 5.7% on speed-up ratio
- Architecture: Multiple position-specialized transformer layers trained with simulated error accumulation