python - 跨多个维度的 Top-k 评分

Question

例如在机器学习中的自然语言处理中，束搜索通常用于预测下一个要添加到序列中的对象并对它们进行排序。 beam-search 的一个关键部分是 top-k 得分指标，它实际上是:给定长度为 N 的概率得分选择列表，返回最高 >k 个评分项目，共 N 个。这就像对列表进行排序然后取最高值一样简单。
引用视觉示例 https://www.researchgate.net/figure/A-partially-completed-beam-search-procedure-with-a-beam-width-of-5-for-an-example-input_fig2_317377611在 beam-search 中(在上述情况下，k=5，并且“最高”分数是最小值)，在每次迭代中，每个节点选择最高 k 项来自 N 个选择列表，导致 k² 个潜在路径。从这些路径中，总的 k 被过滤掉，它们形成了下一次迭代的节点。在前面的示例中，您只能看到每个时间步的过滤节点。 https://d2l.ai/_images/beam-search.svg全面扩展了 k=2, N=5 的情况。

想象一下，不是为每个分支/节点优化 N 中的一个选择，而是必须选择多个值:从节点探索时，您有一组维度选择( N, q) 您要从中选择 q 个值，每列 q 一个。然后，要找到得分最高的选项集，您需要考虑这些列中值的组合。例如: 对于选择矩阵 N=5，q=4:

+---+--------+--------+--------+--------+
| N |   q0   |   q1   |   q2   |   q3   |
+---+--------+--------+--------+--------+
| 0 | 0.9763 | 0.0791 | 0.1530 | 0.5565 |
| 1 | 0.1560 | 0.1014 | 0.6932 | 0.7551 |
| 2 | 0.8142 | 0.9494 | 0.4582 | 0.4411 |
| 3 | 0.3807 | 0.2403 | 0.6897 | 0.7356 |
| 4 | 0.0156 | 0.9419 | 0.9568 | 0.2266 |
+---+--------+--------+--------+--------+

如果 k=5，此 top-k 函数应返回以下内容:

1. 3.6376 = q₀[0] + q₁[2] + q₂[4] + q₃[1]
2. 3.6301 = q₀[0] + q₁[4] + q₂[4] + q₃[1]
3. 3.6181 = q₀[0] + q₁[2] + q₂[4] + q₃[3]
4. 3.6106 = q₀[0] + q₁[4] + q₂[4] + q₃[3]
5. 3.4755 = q₀[2] + q₁[2] + q₂[4] + q₃[1]

这是最大可能的总和，使用每一列中的一个值。

针对任意 N 和 q 解决这个问题，天真的方法是计算所有 N^q 求和，对它们进行排序，然后取前 k 个结果。优化的第一步是对每一列进行排序，然后只计算每一列中前 k 值的总和组合，将复杂度降低到 k ^q。

但是，考虑到这个查找最高分的函数必须在波束搜索的每个时间步调用 k 次，如果希望扩展到高 k，每一种可能的加速都是至关重要的 或高 q。我提出的最佳解决方案(浓缩为最小示例，假设 ma​​trix 是一个形状为 (N, q) 的 numpy 数组, 取 q 为 4):

import numpy as np
from itertools import combinations


class Beamsearch():
    def __init__(self, klen, q=4):
        self.klen = klen
        self.combis = []
        for lens in range(klen):
            self.combis.extend(list(self.partition(lens, q)))
        self.width = q
        self.wdth = list(range(q))

    def partition(self, N, size):
        n = N + size - 1
        for splits in combinations(range(n), size - 1):
            yield [s1 - s0 - 1 for s0, s1 in zip((-1,) + splits, splits + (n,))]

    def getkmaxscores(self, matrix):
        matrix_argsort = np.argsort(-matrix, axis=0)
        sums = []
        for comb in self.combis:
            midxs = matrix_argsort[comb, self.wdth]
            midxslist = midxs.tolist()
            msum = (sum(matrix[midxs, self.wdth]),
                    midxslist)
            sums.append(msum)
        sums.sort(reverse=True)
        return sums[:self.klen]

此方法为整数 0 ≤ p ≤ k 创建整数 p 到给定宽度 q 的分区，例如对于q=4:

p0: [0, 0, 0, 0]
p1: [0, 0, 0, 1], [0, 0, 1, 0], [0, 1, 0, 0], [1, 0, 0, 0]
p2: [0, 0, 0, 2], [0, 0, 1, 1], [0, 0, 2, 0], [0, 1, 0, 1], [0, 1, 1, 0], [0, 2, 0, 0], [1, 0, 0, 1], [1, 0, 1, 0], [1, 1, 0, 0], [2, 0, 0, 0]

等等

这些然后用于索引 argsorted 输入矩阵，以选择每个组合进行求和。在 q=4 的情况下，p_i 的长度遵循三角金字塔序列 (https://oeis.org/A000292):这将搜索空间缩小到所有 p_0...k 的总和，即二项式系数 (k,4) = k(k-1)(k-2)(k-3)/24(https://oeis.org/A000332)。这是对小 k 的 k⁴ 解决方案的巨大改进(对于 k < 30，这更少比 k³)，但仍以 k⁴ 的数量级增长。是否存在复杂度 k^q) 的任意情况的解决方案？

Answer 1

这个问题在文献中被称为从 X + Y 中选择。规范引用是 Frederickson and Johnson当对 X 和 Y 进行排序时，谁给出了 O(k) 时间最优算法。您的列未排序，并且 F&J 的算法非常复杂，所以让我画出更简单的 O(k log k) 算法。

首先对X和Y都选择前k个元素排序。初始化最大堆，其中元素 (i, j) 的优先级为 X[i] + Y[j]。插入 (0, 0)。重复以下 k 次:弹出最大元素 (i, j) 并记录其优先级。插入 (i, j+1)。如果 j = 0，也插入 (i+1, 0)。这一切都需要时间 O(n + k log k)，其中 n 是列中元素的数量。

最后，让我们将问题简化为两列。如果有两个以上，例如 X, Y, Z，那么我们可以从 X + Y 中选择前 k 个元素，然后从 (X + Y) + Z 中选择前 k 个元素。