【大模型】- 多 Token 预测

宋 算法
  2026-08-19 09:18 2026-08-19 09:18 创建   2026-08-21 16:20:11 2026-08-21 16:20:11 更新      1.7k 字  7 分钟

多 Token 预测

一次预测多个 token，加速生成

类型： 学习　｜　语言： Python　｜　🏷 前置：《原生稀疏注意力》（本系列第 16 篇）

学习目标

  • 理解多 token 预测的原理
  • 实现多 token 预测头
  • 训练多 token 预测模型
  • 评估加速效果
  • 分析多 token 预测的权衡

多 token 预测概述

多 token 预测（Multi-Token Prediction）是指模型一次预测多个未来的 token，而不是逐个预测。

标准预测的问题

  • 逐个预测：每次只预测下一个 token
  • 推理慢：需要多次前向传播
  • 串行依赖：无法并行生成

多 token 预测的优势

  • 并行生成：一次预测多个 token
  • 加速推理：减少前向传播次数
  • 更好表示：学习更丰富的上下文表示

实现

基本架构

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiTokenPrediction(nn.Module):
    def __init__(self, vocab_size, d_model, n预测头=4):
        super().__init__()
        self.d_model = d_model
        self.n预测头 = n预测头
        
        # 基础 Transformer
        self.transformer = Transformer(d_model)
        
        # 多个预测头
        self.prediction_heads = nn.ModuleList([
            nn.Linear(d_model, vocab_size) 
            for _ in range(n预测头)
        ])
    
    def forward(self, x):
        # 基础表示
        hidden = self.transformer(x)
        
        # 多个预测
        predictions = []
        for head in self.prediction_heads:
            pred = head(hidden)
            predictions.append(pred)
        
        return torch.stack(predictions, dim=1)  # (batch, n预测头, seq_len, vocab_size)

训练目标

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def multi_token_loss(predictions, targets, n预测头=4):
    """多 token 预测损失"""
    batch_size, seq_len, vocab_size = predictions.shape
    
    total_loss = 0
    
    for i in range(n预测头):
        # 第 i 个预测头预测第 i+1 个未来 token
        pred = predictions[:, i, :-(i+1), :]  # 去掉最后 i+1 个位置
        target = targets[:, i+1:]  # 偏移 i+1 个位置
        
        loss = F.cross_entropy(
            pred.reshape(-1, vocab_size),
            target.reshape(-1)
        )
        
        total_loss += loss
    
    return total_loss / n预测头

自回归生成

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
@torch.no_grad()
def generate_multi_token(model, prompt, max_new_tokens, n预测头=4):
    """多 token 预测生成"""
    model.eval()
    
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    generated = input_ids.clone()
    
    while generated.shape[1] < max_new_tokens:
        # 预测多个 token
        predictions = model(generated)
        
        # 从每个预测头采样
        next_tokens = []
        for i in range(min(n预测头, max_new_tokens - generated.shape[1])):
            pred = predictions[:, i, -1, :]
            probs = F.softmax(pred, dim=-1)
            token = torch.multinomial(probs, 1)
            next_tokens.append(token)
        
        # 追加到序列
        next_tokens = torch.cat(next_tokens, dim=1)
        generated = torch.cat([generated, next_tokens], dim=1)
    
    return generated

高级架构

层级预测头

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class HierarchicalPrediction(nn.Module):
    def __init__(self, vocab_size, d_model, n预测头=4):
        super().__init__()
        self.d_model = d_model
        self.n预测头 = n预测头
        
        # 基础 Transformer
        self.transformer = Transformer(d_model)
        
        # 层级预测头
        self.prediction_heads = nn.ModuleList()
        for i in range(n预测头):
            head = nn.Sequential(
                nn.Linear(d_model, d_model // 2),
                nn.GELU(),
                nn.Linear(d_model // 2, vocab_size)
            )
            self.prediction_heads.append(head)
        
        # 上下文融合
        self.context_fusion = nn.Linear(d_model * n预测头, d_model)
    
    def forward(self, x):
        # 基础表示
        hidden = self.transformer(x)
        
        # 层级预测
        predictions = []
        context = hidden
        
        for i, head in enumerate(self.prediction_heads):
            # 融合上下文
            fused = self.context_fusion(
                torch.cat([hidden] + [p.argmax(dim=-1).unsqueeze(-1).expand(-1, -1, self.d_model) 
                                      for p in predictions], dim=-1)
            ) if predictions else hidden
            
            # 预测
            pred = head(fused)
            predictions.append(pred)
        
        return torch.stack(predictions, dim=1)

自适应预测长度

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
class AdaptiveMultiTokenPrediction(nn.Module):
    def __init__(self, vocab_size, d_model, max预测头=8):
        super().__init__()
        self.d_model = d_model
        self.max预测头 = max预测头
        
        # 基础 Transformer
        self.transformer = Transformer(d_model)
        
        # 预测头
        self.prediction_heads = nn.ModuleList([
            nn.Linear(d_model, vocab_size) 
            for _ in range(max预测头)
        ])
        
        # 自适应长度预测器
        self.length_predictor = nn.Sequential(
            nn.Linear(d_model, d_model // 4),
            nn.GELU(),
            nn.Linear(d_model // 4, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        # 基础表示
        hidden = self.transformer(x)
        
        # 预测长度
        length_pred = self.length_predictor(hidden[:, -1, :])
        predicted_length = (length_pred * self.max预测头).int().item()
        predicted_length = max(1, min(predicted_length, self.max预测头))
        
        # 只使用前 predicted_length 个预测头
        predictions = []
        for i in range(predicted_length):
            pred = self.prediction_heads[i](hidden)
            predictions.append(pred)
        
        return torch.stack(predictions, dim=1), predicted_length

训练策略

渐进式训练

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def progressive_training(model, train_data, config):
    """渐进式训练：从 1 个预测头逐渐增加到 n 个"""
    
    for epoch in range(config["epochs"]):
        # 当前 epoch 使用的预测头数量
        current_n_heads = min(epoch // config["heads_per_epoch"] + 1, config["max_heads"])
        
        print(f"Epoch {epoch + 1}: 使用 {current_n_heads} 个预测头")
        
        # 训练
        for batch in train_data:
            # 前向传播
            predictions = model(batch["input_ids"])
            
            # 只使用前 current_n_heads 个预测
            predictions = predictions[:, :current_n_heads]
            
            # 计算损失
            loss = multi_token_loss(predictions, batch["input_ids"], current_n_heads)
            
            # 反向传播
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

样本权重

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def weighted_multi_token_loss(predictions, targets, n预测头=4):
    """带权重的多 token 预测损失"""
    batch_size, seq_len, vocab_size = predictions.shape
    
    total_loss = 0
    
    for i in range(n预测头):
        pred = predictions[:, i, :-(i+1), :]
        target = targets[:, i+1:]
        
        # 权重：预测越远，权重越低
        weight = 1.0 / (i + 1)
        
        loss = F.cross_entropy(
            pred.reshape(-1, vocab_size),
            target.reshape(-1)
        ) * weight
        
        total_loss += loss
    
    return total_loss / n预测头

评估加速效果

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def evaluate_speedup(model_standard, model_multi, test_prompts, n预测头=4):
    """评估加速效果"""
    
    # 标准模型
    start_time = time.time()
    for prompt in test_prompts:
        inputs = tokenizer.encode(prompt, return_tensors="pt")
        model_standard.generate(inputs, max_new_tokens=100)
    standard_time = time.time() - start_time
    
    # 多 token 预测模型
    start_time = time.time()
    for prompt in test_prompts:
        generate_multi_token(model_multi, prompt, max_new_tokens=100, n预测头=n预测头)
    multi_time = time.time() - start_time
    
    speedup = standard_time / multi_time
    
    print(f"标准模型时间: {standard_time:.2f}s")
    print(f"多 token 模型时间: {multi_time:.2f}s")
    print(f"加速比: {speedup:.2f}x")
    
    return speedup

质量评估

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def evaluate_quality(model_standard, model_multi, eval_dataset, n预测头=4):
    """评估生成质量"""
    
    # 标准模型困惑度
    ppl_standard = compute_perplexity(model_standard, eval_dataset)
    
    # 多 token 模型困惑度
    ppl_multi = compute_perplexity(model_multi, eval_dataset)
    
    # 生成质量
    samples_standard = [generate(model_standard, p) for p in eval_dataset[:100]]
    samples_multi = [generate_multi_token(model_multi, p, 100, n预测头) for p in eval_dataset[:100]]
    
    # 使用 BLEU 或其他指标评估
    bleu_standard = compute_bleu(samples_standard, eval_dataset[:100])
    bleu_multi = compute_bleu(samples_multi, eval_dataset[:100])
    
    print(f"标准模型 - 困惑度: {ppl_standard:.4f}, BLEU: {bleu_standard:.4f}")
    print(f"多 token 模型 - 困惑度: {ppl_multi:.4f}, BLEU: {bleu_multi:.4f}")

权衡分析

精度-速度权衡

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
def analyze_tradeoff(n预测头_list, eval_data):
    """分析不同预测头数量的权衡"""
    
    results = []
    
    for n_heads in n预测头_list:
        model = MultiTokenPrediction(vocab_size, d_model, n_heads)
        
        # 训练模型
        train_model(model, train_data)
        
        # 评估
        perplexity = compute_perplexity(model, eval_data)
        speedup = measure_speedup(model, n_heads)
        
        results.append({
            "n_heads": n_heads,
            "perplexity": perplexity,
            "speedup": speedup,
            "efficiency": speedup / perplexity
        })
        
        print(f"预测头数: {n_heads}")
        print(f"  困惑度: {perplexity:.4f}")
        print(f"  加速比: {speedup:.2f}x")
        print(f"  效率: {speedup / perplexity:.4f}")
    
    return results

内存使用

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def analyze_memory_usage(n预测头_list):
    """分析内存使用"""
    
    results = []
    
    for n_heads in n预测头_list:
        model = MultiTokenPrediction(vocab_size, d_model, n_heads)
        
        # 计算参数量
        n_params = sum(p.numel() for p in model.parameters())
        
        # 估算内存
        memory_mb = n_params * 4 / 1024 / 1024  # float32
        
        results.append({
            "n_heads": n_heads,
            "n_params": n_params,
            "memory_mb": memory_mb
        })
        
        print(f"预测头数: {n_heads}")
        print(f"  参数量: {n_params:,}")
        print(f"  内存: {memory_mb:.2f}MB")
    
    return results

最佳实践

  1. 选择合适的预测头数量：通常 2-4 个预测头效果最好
  2. 使用渐进式训练：从 1 个预测头开始，逐渐增加
  3. 平衡精度和速度：根据应用需求选择
  4. 监控生成质量：确保多 token 预测不降低质量

总结

多 token 预测通过一次预测多个 token，显著加速了自回归生成。关键组件包括多个预测头、训练目标和生成策略。需要权衡预测头数量、生成速度和质量。

下一步

下一课将介绍 DualPipe 并行，优化流水线并行。

📚 本文改编自 AI Engineering from Scratch（MIT License · 作者 Rohit Ghumare），中文内容来自官方中文镜像。原课程共 503 课 · 20 阶段 · 免费开源，教程网站见 aiengineeringfromscratch.com。

  • 标题: 【大模型】- 多 Token 预测
  • 作者: 宋
  • 创建于 : 2026-08-19 09:18:00
  • 更新于 : 2026-08-21 16:20:11
  • 链接: https://sxl-space.tk/2026/08/19/010_LLM/010_LLM-18-MultiTokenPrediction/
  • 版权声明: 版权所有 © 宋，禁止转载。
目录
【大模型】- 多 Token 预测
  1. 1. 多 Token 预测
  2. 2. 学习目标
  3. 3. 多 token 预测概述
    1. 3.1. 标准预测的问题
    2. 3.2. 多 token 预测的优势
  4. 4. 实现
    1. 4.1. 基本架构
    2. 4.2. 训练目标
    3. 4.3. 自回归生成
  5. 5. 高级架构
    1. 5.1. 层级预测头
    2. 5.2. 自适应预测长度
  6. 6. 训练策略
    1. 6.1. 渐进式训练
    2. 6.2. 样本权重
  7. 7. 评估加速效果
    1. 7.1. 质量评估
  8. 8. 权衡分析
    1. 8.1. 精度-速度权衡
    2. 8.2. 内存使用
  9. 9. 最佳实践
  10. 10. 总结
  11. 11. 下一步
  12. 12. 下一课将介绍 DualPipe 并行，优化流水线并行。