TSP问题从DP算法到深度学习4：概率最大状态序列算法

本篇是TSP问题从DP算法到深度学习系列第四篇，这一篇我们会详细举例并比较在 seq-to-seq 或者Markov Chain中的一些常见的搜索概率最大的状态序列的算法。这些方法在深度学习的seq-to-seq 中被用作decoding。在第五篇中，我们使用强化学习时也会使用了本篇中讲到的方法。

• 第一篇: 递归DP方法 AC AIZU TSP问题

• 第二篇: 二维空间TSP数据集及其DP解法

• 第三篇: 深度学习 Pointer Networks 的 Pytorch实现

• 第四篇: 搜寻最有可能路径：Viterbi算法和其他

• 第五篇: 深度强化学习无监督算法的 Pytorch实现

马尔科夫链问题

在 seq-to-seq 问题中，我们经常会遇到需要从现有模型中找概率最大的可能状态序列。这类问题在机器学习算法和控制领域广泛存在，抽象出来可以表达成马尔可夫链模型：给定初始状态的分布和系统的状态转移方程（称为系统动力，dynamics），找寻最有可能的状态序列。

举个例子，假设系统有 \(n\) 个状态，初始状态由 $s_0 = [0.35, 0.25, 0.4] $ 指定，表示初始时三种状态的分布为 0.35，0.25和0.4。

状态转移矩阵由 \(T\) 表达，其中 $ T[i][j]$ 表示从状态 \(i\) 到状态 \(j\) 的概率。注意下面的矩阵 \(T\) 每行的和为 1.0，对应了从任意状态出发，下一时刻的所有可能转移概率和为1。 \[ T= \begin{matrix} & \begin{matrix}0&1&2\end{matrix} \\\\ \begin{matrix}0\\\\1\\\\2\end{matrix} & \begin{bmatrix}0.3&0.6&0.1\\\\0.4&0.2&0.4\\\\0.3&0.3&0.4\end{bmatrix}\\\\ \end{matrix} \]

至此，系统的所有参数都定下来了。接下去的各个时刻的状态分布可以通过矩阵乘法来算得。比如，记\(s_1\) 为 \(t_1\) 时刻状态分布，计算方法为 \(s_0\) 乘以 \(T\)，动画如下：

\(s_1\) 数值计算结果如下。

\[ s_1 = \begin{bmatrix}0.35& 0.25& 0.4\end{bmatrix} \begin{matrix} \begin{bmatrix}0.3&0.6&0.1\\\\0.4&0.2&0.4\\\\0.3&0.3&0.4\end{bmatrix}\\\\ \end{matrix} = \begin{bmatrix}0.325& 0.35& 0.255\end{bmatrix} \] 矩阵左乘行向量可以理解为矩阵每一行的线性组合，直觉上理解为下一时刻的状态分布是上一时刻初始状态分布乘以转移关系的线性组合。 \[ \begin{bmatrix}0.35& 0.25& 0.4\end{bmatrix} \begin{matrix} \begin{bmatrix}0.3&0.6&0.1\\\\0.4&0.2&0.4\\\\0.3&0.3&0.4\end{bmatrix}\\\\ \end{matrix} = 0.35 \times \begin{bmatrix}0.35& 0.6& 0.1\end{bmatrix} + 0.25 \times \begin{bmatrix}0.4& 0.2& 0.4\end{bmatrix} + 0.4 \times \begin{bmatrix}0.3& 0.3& 0.4\end{bmatrix} \] 同样的，后面每一个时刻都可以由上一个状态分布向量乘以 \(T\)，当然这里我们假设每个时刻的转移矩阵是不变的。当然，问题也可以是每个时刻都有不同的转移矩阵来定义，例如深度学习 seq-to-seq 模型。当然，这个设定的变化不会影响搜索最可能状态序列的算法。出于简单考虑，本篇中我们假定所有时刻的状态转移矩阵都是 \(T\)。

下面我们通过多种算法来找出由上述参数定义的系统中前三个时刻的最有可能序列，即概率最大的 \(s_0 \rightarrow s_1 \rightarrow s_2\)。

令 \(L\) 是阶段数，\(N\) 是每个阶段的状态数，则我们的例子中 \(L=N=3\) 。并且，总共有 \(N^L\) 种不同的路径。

穷竭搜索

若给定一条路径，计算特定路径的概率是很直接的，例如，若给定路径为 \(2(s_0) \rightarrow 1(s_1) \rightarrow 2(s_2)\)，则这条路径的概率为

\[ p(2 \rightarrow 1 \rightarrow 2) = s_0[2] \times T[2][1] \times T[1][2] = 0.4 \times 0.3 \times 0.4 = 0.048 \]

因此，我们可以通过枚举所有 \(N^L\) 条路径并计算每条路径的概率来找到最有可能的状态序列。

下面是Python 3的穷竭搜索代码，输出为最有可能的概率及其路径。样例问题的输出为 0.084 和状态序列 \(0 \rightarrow 1 \rightarrow 2\)。

{linenos

def search_brute_force(initial: List, transition: List, L: int) -> Tuple[float, Tuple]:
    from itertools import combinations_with_replacement
    v = [0, 1, 2]
    path_all = combinations_with_replacement(v, L)

    max_prop = 0.0
    max_route = None
    prob = 0.0
    for path in list(path_all):
        for idx, v in enumerate(path):
            if idx == 0:
                prob = initial[v]  # reset to initial state
            else:
                prev_v = path[idx-1]
                prob *= transition[prev_v][v]
        if prob > max_prop:
            max_prop = max(max_prop, prob)
            max_route = path
    return max_prop, max_route

贪心搜索

穷竭搜索一定会找到最有可能的状态序列，但是算法复杂度是指数级的 \(O(N^L)\)。一种最简化的策略是，每一时刻都只选取下一时刻最可能的状态，显然这种策略没有考虑全局最优，只考虑下一步最优，因此称为贪心策略。当然，贪心策略虽然牺牲全局最优解但是换取了很快的时间复杂度。贪心搜索算法动画如下。

Python 3 实现中我们利用了 numpy 类库，主要是 np.argmax() 可以让代码简洁。代码本质上是两重循环，（一层循环是np.argmax中），对应时间算法复杂度是 \(O(N\times L)\)。

{linenos

def search_greedy(initial: List, transition: List, L: int) -> Tuple[float, Tuple]:
    import numpy as np
    max_route = []
    max_prop = 0.0
    states = np.array(initial)

    prev_max_v = None
    for l in range(0, L):
        max_v = np.argmax(states)
        max_route.append(max_v)
        if l == 0:
            max_prop = initial[max_v]
        else:
            max_prop = max_prop * transition[prev_max_v][max_v]
        states = max_prop * states
        prev_max_v = max_v

    return max_prop, max_route

Beam 搜索

贪心策略只考虑了下一步的最大概率状态，若我们改进一下贪心策略，将下一步的最大 \(k\) 个状态保留下来就是beam 搜索了。具体来说， \(k\) beam search表示每个阶段保留 \(k\) 个最大概率路径，下一阶段扩展这 \(k\) 条路径至 \(k \times N\) 条路径再选取最大的top k。以上例来说，选取\(k=2\)，则初始 \(s_0\)时选取最大概率的两种状态 0和 2，下一阶段 \(s_1\)，计算以0和2开始的共 \(2 \times 3\) 条路径，并保留其中最大概率的两条，如此往复。显然，beam search也无法找到全局最优解，但是它能以线性时间复杂度探索更多的路径空间。

以下是Python 3 的代码实现，利用了 PriorityQueue 选取 \(k\) 路径。由于PriorityQueue 无法自定义比较关系，我们定义了 @total_ordering 标注的类来实现比较关心。时间算法复杂度是 \(O(k\times N \times L)\) 。

{linenos

def search_beam(initial: List, transition: List, L: int, K: int) -> Tuple[float, Tuple]:
    N = len(initial)
    from queue import PriorityQueue
    current_q = PriorityQueue()
    next_q = PriorityQueue()

    from functools import total_ordering
    @total_ordering
    class PQItem(object):
        def __init__(self, prob, route):
            self.prob = prob
            self.route = route
            self.last_v = int(route[-1])

        def __eq__(self, other):
            return self.prob == other.prob

        def __lt__(self, other):
            return self.prob > other.prob

    for v in range(N):
        next_q.put(PQItem(initial[v], str(v)))

    for l in range(1, L):
        current_q = next_q
        next_q = PriorityQueue()
        k = K
        while not current_q.empty() and k > 0:
            item = current_q.get()
            prob, route, prev_v = item.prob, item.route, item.last_v
            k -= 1
            for v in range(N):
                nextItem = PQItem(prob * transition[prev_v][v], route + str(v))
                next_q.put(nextItem)

    max_item = next_q.get()

    return max_item.prob, list(map(lambda x: int(x), max_item.route))

Viterbi 动态规划

和之前TSP 动态规划算法的思想一样，最有可能状态路径问题解法有可以将指数时间复杂度 \(O(N^L)\) 降到多项式时间复杂度 \(O(L \times N \times N)\) 的算法，就是大名鼎鼎的 Viterbi 算法（维特比算法）。核心思想是在每个阶段，用数组保存每个状态结尾路径的阶段最大概率（及其对应路径）。在不考虑优化空间的情况下，我们开一个二维数组 \(dp[][]\)，第一维表示阶段序号，第二维表示状态序号。例如，\(dp[1][0]\) 是 \(s_1\) 阶段时以状态0结尾的所有路径中的最大概率，即 \[ dp[1][0] = \max \\{s_0[0] \rightarrow s_1[0], s_0[1] \rightarrow s_1[0], s_0[2] \rightarrow s_1[0]\\} \]

实现代码中没有返回路径本身而只是其概率值，目的是通过简洁的三层循环来表达算法精髓。

{linenos

def search_dp(initial: List, transition: List, L: int) -> float:
    N = len(initial)
    dp = [[0.0 for c in range(N)] for r in range(L)]
    dp[0] = initial[:]

    for l in range(1, L):
        for v in range(N):
            for prev_v in range(N):
                dp[l][v] = max(dp[l][v], dp[l - 1][prev_v] * transition[prev_v][v])

    return max(dp[L-1])

概率采用

以上所有的算法都是确定性的。在NLP 深度学习decoding 时候会带来一个问题：确定性容易导致生成重复的短语或者句子。比如，确定性算法很容易生成如下句子。

This is the best of best of best of ...

一种简单的方法是采用概率采用的方式回避这个问题。也就是我们不寻找确定的局部最优或者全局最优的解，而是通过局部路径或者全局路径的概率信息进行采样生成序列。例如，对于穷竭搜索的 \(N^L\) 条路径计算得到对应概率，转变成 \(N^L\) 个点的 categorical 分布，采样生成某条路径。也可以如下改造贪心或者beam 这类阶段性生成算法一个时刻一个时刻的输出采样的状态序列。

{linenos

def search_prob_greedy(initial: List, transition: List, L: int) -> Tuple[float, Tuple]:
    import random
    N = len(initial)
    max_route = []
    max_prop = 0.0
    vertices = [i for i in range(N)]
    prob = initial[:]

    for l in range(0, L):
        v_lst = random.choices(vertices, prob)
        v = v_lst[0]
        max_route.append(v)
        max_prop = prob[v]
        prob = [prob[v] * transition[v][v_target] for v_target in range(N)]

    return max_prop, max_route