每日算法 — 使用java实现束搜索：序列生成与启发式剪枝策略

引言

束搜索（Beam Search）是序列生成任务中广泛使用的启发式搜索算法，在机器翻译、语音识别、文本摘要等领域发挥着关键作用。与贪心搜索每次只保留一个最优候选不同，束搜索在每一步保留 k 个最优候选（称为束宽，beam width），在探索更多可能性的同时控制计算开销。本文将用Java实现标准束搜索及其长度归一化变体，并结合机器翻译场景展示其序列解码能力。

核心原理

贪心搜索的局限

贪心搜索（Greedy Decoding）每一步选择概率最高的词，但局部最优不等于全局最优。一个早期的高概率选择可能导致后续生成质量急剧下降。

束搜索策略

束搜索维护一个大小为 k 的候选集合（束）。在每一步：

1. 对当前束中每个候选，扩展所有可能的下一个词
2. 计算每个扩展路径的累积得分
3. 从所有扩展结果中保留得分最高的 k 个作为新束
4. 重复直到生成结束符或达到最大长度

长度归一化

由于长序列的联合概率是多个小于1的数相乘，束搜索天然偏向短序列。长度归一化通过将得分除以序列长度来消除这种偏差：

score(y) = (1 / |y|^α) · Σ_{t=1}^{|y|} log P(y_t | y_{<t}, x)

其中 α 为长度惩罚系数，通常取0.6-1.0。

Java实现

项目结构

src/
├── BeamSearch.java       // 核心束搜索实现
├── SequenceNode.java     // 搜索树节点
└── BeamSearchDemo.java   // 演示主类

搜索树节点

import java.util.ArrayList;
import java.util.List;

/**
 * 束搜索中的序列节点
 * 代表一个部分解码的序列及其累积得分
 */
public class SequenceNode implements Comparable<SequenceNode> {
    public final List<Integer> tokens;      // 已生成的词ID序列
    public final double logProb;            // 累积对数概率
    public final boolean isComplete;        // 是否已生成结束符

    public SequenceNode(List<Integer> tokens, double logProb, boolean isComplete) {
        this.tokens = new ArrayList<>(tokens);
        this.logProb = logProb;
        this.isComplete = isComplete;
    }

    /**
     * 扩展当前序列，添加一个新词
     */
    public SequenceNode extend(int token, double tokenLogProb) {
        List<Integer> newTokens = new ArrayList<>(tokens);
        newTokens.add(token);
        return new SequenceNode(newTokens, logProb + tokenLogProb, token == 0); // 0表示结束符
    }

    public int length() {
        return tokens.size();
    }

    /**
     * 长度归一化得分
     */
    public double normalizedScore(double alpha) {
        return logProb / Math.pow(length(), alpha);
    }

    @Override
    public int compareTo(SequenceNode other) {
        // 按累积对数概率降序排列
        return Double.compare(other.logProb, this.logProb);
    }

    @Override
    public String toString() {
        return "SequenceNode{tokens=" + tokens + ", logProb=" + String.format("%.4f", logProb) + "}";
    }
}

束搜索核心实现

import java.util.*;

/**
 * 束搜索解码器
 * 
 * @param <T> 模型接口类型，提供nextTokenLogProbs方法
 */
public class BeamSearch<T extends LanguageModel> {
    private final T model;
    private final int beamWidth;
    private final int maxLength;
    private final double lengthPenalty;
    private final int eosToken; // 结束符ID

    public BeamSearch(T model, int beamWidth, int maxLength, 
                      double lengthPenalty, int eosToken) {
        this.model = model;
        this.beamWidth = beamWidth;
        this.maxLength = maxLength;
        this.lengthPenalty = lengthPenalty;
        this.eosToken = eosToken;
    }

    /**
     * 执行束搜索解码
     * @param encoderOutput 编码器输出上下文
     * @return 最优的k个完整序列，按得分降序排列
     */
    public List<SequenceNode> decode(double[] encoderOutput) {
        // 初始化: 从空序列开始
        PriorityQueue<SequenceNode> beam = new PriorityQueue<>();
        beam.offer(new SequenceNode(Collections.emptyList(), 0.0, false));

        // 存储已完成的序列
        List<SequenceNode> completed = new ArrayList<>();

        for (int step = 0; step < maxLength; step++) {
            PriorityQueue<SequenceNode> candidates = new PriorityQueue<>();

            // 对当前束中每个候选进行扩展
            for (SequenceNode node : beam) {
                if (node.isComplete) {
                    // 已完成序列不再扩展，直接保留
                    candidates.offer(node);
                    continue;
                }

                // 获取模型预测的下一个词的对数概率分布
                double[] logProbs = model.nextTokenLogProbs(encoderOutput, node.tokens);

                // 只保留概率最高的topK个扩展，减少计算量
                int topK = Math.min(logProbs.length, beamWidth * 2);
                int[] topIndices = getTopKIndices(logProbs, topK);

                for (int idx : topIndices) {
                    SequenceNode newNode = node.extend(idx, logProbs[idx]);
                    candidates.offer(newNode);
                }
            }

            // 从所有候选中选择最优的beamWidth个作为新束
            beam = new PriorityQueue<>();
            while (!candidates.isEmpty() && beam.size() < beamWidth) {
                SequenceNode best = candidates.poll();
                if (best.isComplete) {
                    completed.add(best);
                } else {
                    beam.offer(best);
                }
            }

            // 如果所有候选都已完成，提前终止
            if (beam.isEmpty()) break;
        }

        // 将未完成的序列也加入结果
        completed.addAll(beam);

        // 按长度归一化得分排序
        completed.sort((a, b) -> Double.compare(
            b.normalizedScore(lengthPenalty),
            a.normalizedScore(lengthPenalty)
        ));

        return completed.subList(0, Math.min(beamWidth, completed.size()));
    }

    /**
     * 获取对数概率最高的k个索引
     */
    private int[] getTopKIndices(double[] logProbs, int k) {
        Integer[] indices = new Integer[logProbs.length];
        for (int i = 0; i < indices.length; i++) indices[i] = i;
        Arrays.sort(indices, (a, b) -> Double.compare(logProbs[b], logProbs[a]));
        int[] result = new int[k];
        for (int i = 0; i < k; i++) result[i] = indices[i];
        return result;
    }
}

/**
 * 语言模型接口
 */
interface LanguageModel {
    /**
     * 预测下一个词的对数概率分布
     * @param context 编码器上下文
     * @param prefix 已生成的词序列
     * @return 词表上每个词的对数概率
     */
    double[] nextTokenLogProbs(double[] context, List<Integer> prefix);
}

演示：简单翻译模型

import java.util.*;

/**
 * 模拟一个简单的统计机器翻译模型
 */
class SimpleTranslationModel implements LanguageModel {
    private final int vocabSize;
    private final Random rand;

    public SimpleTranslationModel(int vocabSize) {
        this.vocabSize = vocabSize;
        this.rand = new Random(42);
    }

    @Override
    public double[] nextTokenLogProbs(double[] context, List<Integer> prefix) {
        // 模拟模型输出: 基于前缀长度和上下文生成概率分布
        double[] probs = new double[vocabSize];
        double sum = 0;

        for (int i = 0; i < vocabSize; i++) {
            // 模拟一些词概率更高（如高频词）
            probs[i] = Math.exp(-0.1 * i) + rand.nextDouble() * 0.5;
            if (prefix.size() > 0 && i == prefix.get(prefix.size() - 1)) {
                // 降低重复词的概率
                probs[i] *= 0.3;
            }
            sum += probs[i];
        }

        // 转换为对数概率
        double[] logProbs = new double[vocabSize];
        for (int i = 0; i < vocabSize; i++) {
            logProbs[i] = Math.log(probs[i] / sum);
        }
        return logProbs;
    }
}

/**
 * 束搜索演示主类
 */
public class BeamSearchDemo {
    public static void main(String[] args) {
        int vocabSize = 100;
        int beamWidth = 5;
        int maxLength = 10;
        double lengthPenalty = 0.8;
        int eosToken = 0;

        SimpleTranslationModel model = new SimpleTranslationModel(vocabSize);
        BeamSearch<SimpleTranslationModel> decoder = new BeamSearch<>(
            model, beamWidth, maxLength, lengthPenalty, eosToken
        );

        // 模拟编码器输出
        double[] encoderOutput = new double[]{1.0, 0.5, 0.3};

        System.out.println("=== 束搜索解码 (beamWidth=" + beamWidth + ") ===\n");
        List<SequenceNode> results = decoder.decode(encoderOutput);

        System.out.println("Top-" + results.size() + " 解码结果:");
        for (int i = 0; i < results.size(); i++) {
            SequenceNode node = results.get(i);
            System.out.printf("%d. 序列: %s%n", i + 1, node.tokens);
            System.out.printf("   累积对数概率: %.4f%n", node.logProb);
            System.out.printf("   长度归一化得分: %.4f%n", node.normalizedScore(lengthPenalty));
            System.out.printf("   长度: %d, 完成: %s%n%n", node.length(), node.isComplete);
        }

        // 对比: 贪心搜索
        System.out.println("=== 贪心搜索对比 ===");
        BeamSearch<SimpleTranslationModel> greedyDecoder = new BeamSearch<>(
            model, 1, maxLength, lengthPenalty, eosToken
        );
        List<SequenceNode> greedyResults = greedyDecoder.decode(encoderOutput);
        if (!greedyResults.isEmpty()) {
            SequenceNode greedy = greedyResults.get(0);
            System.out.println("贪心结果: " + greedy.tokens);
            System.out.println("贪心得分: " + String.format("%.4f", greedy.normalizedScore(lengthPenalty)));
        }
    }
}

复杂度分析

其中 T 为目标序列最大长度，V 为词表大小，k 为束宽。通过topK剪枝（每步只扩展概率最高的几个词），可将每步扩展复杂度从 O(V) 降至 O(k)，整体接近 O(T · k²)。

优化策略

1. 多样化束搜索（Diverse Beam Search）

通过添加多样性惩罚项，鼓励不同束组生成差异更大的序列，避免多个结果过于相似。

2. 约束束搜索（Constrained Beam Search）

在解码过程中强制包含特定词或短语，适用于术语翻译、关键词保留等场景。

3. 非自回归解码

使用迭代 refine 策略，一次性预测整个序列然后逐步修正，突破自回归的 O(T) 串行瓶颈。

总结

本文用Java实现了束搜索算法及其长度归一化变体，通过维护多个候选序列在解码空间中进行启发式搜索，有效平衡了生成质量与计算效率。读者可将此实现集成到机器翻译系统、对话模型或代码生成工具中，并通过调整束宽和长度惩罚系数来适配不同任务需求。