DINOv2中的Sinkhorn-Knopp，收敛原理以及EMD距离

DINOv2的训练方法

DINOv2 采用了一种判别式自监督方法（Discriminative Self-supervised Pre-training）来学习特征，这种方法可以看作是以SwAV为中心的DINO和iBOT 损失的组合。具体而言，作者们实现了以下几种方法：

图像级损失(Image-level objective) 利用一种交叉熵损失函数来比较从学生和教师网络中提取出的特征。这些特征都来自于ViT的cls token，通过对同一图像的不同裁剪图像进行提取得到。作者们使用指数移动平均法（EMA）来构建教师模型，学生模型的参数则通过训练得到（DINO的方法）。

Patch级损失(Patch-level objective) 随机遮盖一些patch输入给student，teacher不遮盖，计算在遮盖的 patch上特征的交叉熵损失函数。

解耦上述的两个优化目标(Untying head weights between both objectives)： 作者们发现，将两种目标的权重绑定在一起会导致模型在 Patch-level 欠拟合，在 Image-level 过拟合。于是不复用权重，使用不同的head就可以避免这个问题。

其中Sinkhorn-Knopp centering 使用SwAV中提到的Sinkhorn-Knopp(SK) Batch Normalization来代替DINO和IBot中的teacher网络中的softmax-centering步骤；对student进行3次Sinkhorn-Knopp算法迭代，再使用softmax进行normalization。

到这里就看不懂了，由此往下研究

Sinkhorn-Knopp算法

Sinkhorn-Knopp 算法解释一

Sinkhorn-Knopp (SK) 算法是一种用于将矩阵归一化为双重随机矩阵的迭代算法。该算法通过交替的行和列归一化操作，使矩阵的行和列之和分别满足给定的目标向量。其广泛应用于最优传输问题和矩阵标度问题。

算法步骤

1. 输入：矩阵 $M$，目标行和向量 $u$，目标列和向量 $v$，迭代次数 $K$。
2. 初始化：矩阵 $P$ 为 $M$ 的归一化版本。
3. 迭代：
   • 对行进行归一化，使每行的和接近 $u$。
   • 对列进行归一化，使每列的和接近 $v$。
4. 输出：迭代 $K$ 次后的矩阵 $P$。

数学定义

对于矩阵 $M$，目标是找到一个矩阵 $P$，使得：

• 矩阵 $P$ 的行和等于向量 $u$。
• 矩阵 $P$ 的列和等于向量 $v$。

具体迭代公式为：

$$P=\frac{P}{\text{np.outer}(u,\text{sum}(P,\text{axis}=1))}$$
$$P=\frac{P}{\text{np.outer}(\text{sum}(P,\text{axis}=0),v)}$$

示例代码

以下是 Sinkhorn-Knopp 算法的实现和一个有效的实例：

import numpy as np

def sinkhorn_knopp(M, u, v, K):
    M = M / np.max(M)  # 对矩阵 M 进行归一化
    P = np.diag(u) @ M @ np.diag(v)
    epsilon = 1e-10  # 防止除零的极小值
    for i in range(K):
        row_sums = np.sum(P, axis=1) + epsilon  # 防止除零
        P = P / row_sums[:, np.newaxis]
        P = P * u[:, np.newaxis]
        
        col_sums = np.sum(P, axis=0) + epsilon  # 防止除零
        P = P / col_sums[np.newaxis, :]
        P = P * v[np.newaxis, :]
    return P

#### 示例使用
M = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
u = np.array([2, 1, 3])  # 改变 u 的值
v = np.array([1, 3, 2])  # 改变 v 的值
K = 10

P = sinkhorn_knopp(M, u, v, K)
print("Normalized matrix P:")
print(P)
print("Row sums:", np.sum(P, axis=1))
print("Column sums:", np.sum(P, axis=0))

输出结果

运行上述代码后，可以得到以下结果：

归一化后的矩阵 $P$:

$$\left[\begin{array}{ccc}0.22358876& 0.99087123& 0.78554002\\ 0.18094781& 0.50118793& 0.31786425\\ 0.59546343& 1.50794084& 0.89659573\end{array}\right]$$

$$\left[\begin{array}{c}0.22358876\\ 0.18094781\\ 0.59546343\end{array}\right]\left[\begin{array}{c}0.99087123\\ 0.50118793\\ 1.50794084\end{array}\right]\left[\begin{array}{c}0.78554002\\ 0.31786425\\ 0.89659573\end{array}\right]$$

行和：

$$[2.0,1.0,3.0]$$

列和：

$$[1.0,3.0,2.0]$$

可以看到满足了需求，这是迭代10次后的结果，然而实际上第一次迭代后的结果就已经很好了

那么下一个问题就是

Sinkhorn-Knopp 算法能够收敛的原理是什么

Sinkhorn-Knopp 算法能够收敛的原理基于几个数学和算法方面的性质。以下是一些关键的原理和性质，它们解释了为什么这个算法能够收敛：

1. 正则化最优传输问题

Sinkhorn-Knopp 算法解决的是一个正则化的最优传输问题。给定两个概率向量 $u$ 和 $v$，目标是找到一个矩阵 $P$，使得：

• $P$ 的每一行的和等于向量 $u$
• $P$ 的每一列的和等于向量 $v$
• 矩阵 $P$ 的所有元素都是非负的

通过加入正则化项（通常是熵正则化），问题变得更加可解，并且解是唯一的。

2. 双重缩放算法

Sinkhorn-Knopp 算法本质上是一种双重缩放算法（Iterative Proportional Fitting Procedure, IPFP）。双重缩放算法的基本步骤包括交替地对矩阵的行和列进行缩放，以逼近目标行和和列和。

迭代步骤：

• 行归一化：对矩阵 $P$ 的每一行进行缩放，使得行和接近目标向量 $u$。
• 列归一化：对矩阵 $P$ 的每一列进行缩放，使得列和接近目标向量 $v$。

这个过程不断重复，逐步调整矩阵 $P$ 的元素，使其行和和列和分别逼近 $u$ 和 $v$。

3. 收敛性分析

3.1 Birkhoff-von Neumann 定理
Birkhoff-von Neumann 定理表明，任意一个双随机矩阵（行和列和都为1的非负矩阵）可以表示为若干置换矩阵的凸组合。虽然 Sinkhorn-Knopp 算法求解的矩阵未必是双随机矩阵，但这个定理为矩阵的归一化操作提供了理论支持。

3.2 均匀缩放的性质
每次迭代中，矩阵 $P$ 的元素都被缩放，以使行和和列和逐渐逼近目标向量 $u$ 和 $v$。这个过程中的缩放操作是线性的，并且因为 $u$ 和 $v$ 是非负向量，所以这个缩放操作总是有效的。

3.3 单调性
在每次迭代中，矩阵 $P$ 的元素变化是单调的，逐渐逼近目标行和和列和。算法的收敛性可以通过观察迭代过程中行和和列和的变化来证明。这些变化逐渐减小，最终稳定在目标向量 $u$ 和 $v$ 附近。

4. 数值稳定性

为了确保算法的数值稳定性，加入了极小值 $ϵ$ 来防止除零错误。此外，通过对矩阵 $M$ 进行归一化处理，进一步增强了算法的数值稳定性，确保了算法在更大规模的矩阵上也能稳定收敛。

5. 正则化的好处

正则化项（如熵正则化）不仅可以使问题更加可解，还可以加速算法的收敛。正则化项通过惩罚矩阵 $P$ 中的极端值，使得矩阵 $P$ 更加平滑，从而提高了迭代过程的数值稳定性和收敛速度。

熵正则化是一种在优化问题中加入熵项的正则化方法，其目的是增加解的稳定性和可解性。熵正则化的数学表达式为：

$$\text{H}(P)=-\sum _{i,j}{P}_{ij}\log P_{ij}$$

其中，$P$ 是待优化的矩阵，$P_{ij}$ 是矩阵 $P$ 的元素。

结论

Sinkhorn-Knopp 算法通过交替的行和列归一化操作，使得矩阵 $P$ 逐渐逼近目标行和和列和。算法的收敛性基于迭代过程的单调性、均匀缩放的性质和正则化项的引入。这些性质确保了算法在合理的迭代次数内收敛到一个满足约束条件的非负矩阵。

---

Sinkhorn-Knopp(SK) Batch Normalization

Sinkhorn-Knopp (SK) Batch Normalization 是一种结合了 Sinkhorn-Knopp 算法和批归一化（Batch Normalization, BN）的方法，用于在深度学习模型中实现归一化操作。其主要目标是对输入的矩阵进行归一化处理，使得其行和与列和符合预期的分布，从而提高模型的训练效率和性能。

Sinkhorn 算法解释二

最优运输问题的目标就是以最小的成本将一个概率分布转换为另一个概率分布。即将概率分布 $c$ 以最小的成本转换到概率分布 $r$，此时就要获得一个分配方案 $P\in {\mathbb{R}}^{n\times m}$，其中需满足以下条件：

• $P$ 的行和服从分布 $r$
• $P$ 的列和服从分布 $c$

因此在分布 $r$、$c$ 约束下，$P$ 的解空间可以做如下定义：

U ( r , c ) = { P ∈ R ≥ 0 n × m ∣ P 1 m = r , P T 1 n = c } U(r, c) = \{ P \in \mathbb{R}^{n \times m}_{\ge 0} \mid P1_m = r, P^T 1_n = c \} U(r,c)={P∈R≥0n×m​∣P1m​=r,PT1n​=c}

同时希望最小化转换成本，即需要一个成本矩阵（cost matrix）$M$。于是就有了最优传输问题的公式化表示：

$$d_M(r,c)=\underset{P\in U(r,c)}{\min }\sum _{i,j}{P}_{ij}{M}_{ij}$$

此时为 Wasserstein metric 或 earth mover distance（EMD 推土机距离）代价函数。Sinkhorn 距离是对推土机距离的一种改进，在其基础上引入了熵正则化项，则代价函数表示为：

$$d_M^{\lambda }(r,c)=\underset{P\in U(r,c)}{\min }⟨P,M⟩-\frac{1}{\lambda }h(P)$$

其中 $h(P)$ 为添加的正则项，即 $P$ 的信息熵：

$$h(P)=-\sum _{i,j}{P}_{ij}\log P_{ij}$$

上式两侧对 $P_{ij}$ 求导：

$$\frac{\partial d_M^{\lambda }(r,c)}{\partial P_{ij}}=M_{ij}+\frac{1}{\lambda }(\log P_{ij}+1)$$

令其为 0 可得：

$$P_{ij}=e^{-\lambda M_{ij}-1}$$

这是在无约束条件下求得的关联矩阵，若考虑约束条件，则上式变为：

$$P_{ij}={\alpha }_i{\beta }_j{e}^{-\lambda M_{ij}}$$

其中 ${\alpha }_i$ 和 ${\beta }_j$ 分别是使行和列满足约束的因子。

伪代码

以下是 Sinkhorn 算法的伪代码：

Input: Cost matrix M, vectors r and c, max iterations K, regularization parameter λ
Output: Optimal transport matrix P

Initialize α = ones(n), β = ones(m)
for k = 1 to K do
    for i = 1 to n do
        for j = 1 to m do
            P_ij = α_i * β_j * exp(-λ * M_ij)
        end for
    end for
    α = r / (P * ones(m))
    β = c / (P^T * ones(n))
end for

Return P

解释二中又引入了推土机距离，常常在点云分析中使用

点云分析中的 EMD（Earth Mover’s Distance）距离

EMD 距离介绍

EMD 距离，又叫做推土机距离，也叫作 Wasserstein 距离。个人理解，EMD 距离是离散化的 Wasserstein 距离，而 Wasserstein 距离是描述两个连续随机变量的 EMD 距离。二者数学思想是相同的，但是所描述的对象和应用场景稍有区分。由于个人正在做关于点云数据的一些研究，因此这篇文章记录的仅仅是 EMD 距离相关的数学描述，不讨论 Wasserstein 距离。

EMD 距离的出处是 2000 年发表在 IJCV 上的 “The Earth Mover’s Distance as a Metric for Image Retrieval” 一文。最初是作为一种度量用来判断两张图像之间的相似度，也就是用来做图像检索工作的。这里，我们从文章中对于 EMD 的定义出发，最后引出在许多点云分析文章中使用的 EMD 做出了哪些假设和简化。

Signature（可以翻译为特征签名）

Signature 的数学定义为：${\mathbf{s}}_{\mathbf{j}}=({\mathbf{m}}_{\mathbf{j}},w_j)$，代表着一个 features 组的聚类，其中 ${\mathbf{m}}_{\mathbf{j}}$ 代表这个聚类类别的平均值 (mean) 或者模式 (mode)，$w_j$ 代表着图像中属于这个类别的像素的占比（在图像处理中），也就是对应类别的权重 (weight)。Histogram 也是统计类别以及占比的统计学工具，但是相比之下，Histogram 的类别分割是等比的，而 Signature 是相对灵活的。

比如，统计数组 {1, 2, 3, 4, 2, 1, 3, 4, 5, 1, 1, 2, 3, 4} 的 Histogram，则会得到 1 有多少个数字，2 有多少个数字等。如果用 Signature 来统计，则可以划分成属于 {1, 2, 3} 这个集合的数字有多少，属于 {4, 5} 这个集合的数字有多少。Signature 比 Histogram 更加灵活，这也提出 Signature 这个数学概念的意义。

Earth Mover’s Distance

假设有两组 Signatures，$P=\{(p_1,w_{p1}),\dots ,(p_m,w_{pm})\}$ 和 $Q=\{(q_1,w_{q1}),\dots ,(q_n,w_{qn})\}$。$P$ 中有 $m$ 个类别，$Q$ 中有 $n$ 个类别。我们可以将两个集合中的 $p_i$ 看作砂矿，$q_j$ 则是砂石仓库，$w_{pi}$ 为每一个砂矿包含的砂石数量，$w_{qj}$ 是每一个仓库能容纳砂石数量。再引入距离矩阵 $\mathbf{D}$ （$m\times n$ 维），其中 $d_{ij}$ 代表从 $p_i$ 到 $q_j$ 之间的距离，一般为欧氏距离。再定义工作流 Flow，记为矩阵 $\mathbf{F}$ （$m\times n$ 维），其中 $f_{ij}$ 代表从 $p_i$ 到 $q_j$ 之间搬运砂石的数量，所以随后的总工作量为：

$$WORK(P,Q,\mathbf{F})=\sum _{i=1}^m\sum _{j=1}^n{d}_{ij}{f}_{ij}$$

另外，对于 $f_{ij}$ 是有条件限制的：

• $f_{ij}\ge 0$，其中 $1\le i\le m$，$1\le j\le n$，这条约束说明砂石只能从 $p_i$ 运向 $q_j$，不能反向。
• ${\sum }_{j=1}^n{f}_{ij}\le w_{pi}$，其中 $1\le i\le m$，这条约束说明从 $p_i$ 砂矿运出的砂石不能超过该矿蕴含的砂石总量。
• ${\sum }_{i=1}^m{f}_{ij}\le w_{qj}$，其中 $1\le j\le n$，这条约束说明运入 $q_j$ 仓库的砂石数量不能超过该仓库的最大容纳量。
• ${\sum }_{i=1}^m{\sum }_{j=1}^n{f}_{ij}=\min \left({\sum }_{i=1}^m{w}_{pi},{\sum }_{j=1}^n{w}_{qj}\right)$，这条约束说明，整个工作完成时，搬运的总砂石数量要么是所有砂矿的储量总和，要么是所有仓库的容纳量总和。

最终的 EMD 距离定义就是归一化之后的工作量：

$$EMD(P,Q)=\frac{{\sum }_{i=1}^m{\sum }_{j=1}^n{d}_{ij}{f}_{ij}}{{\sum }_{i=1}^m{\sum }_{j=1}^n{f}_{ij}}$$

点云分析中的 EMD 距离

假设 $P$ 和 $Q$ 为两个点集，假设：两个点集所包含的点的数量相等，数量记为 $N$。这个假设决定了 EMD 距离中的 $w_{pi}$ 和 $w_{qj}$ 始终保持一致，为 $\frac{1}{N}$。换句话说，这个假设保证了两个点集中的所有点的地位是平等的，这也符合点云分析中的前提，即点云特征与点的顺序置换无关。由于所有的权重均为 $\frac{1}{N}$，所以：

$$\sum _{i=1}^N\sum _{j=1}^N{f}_{ij}=\min \left(\sum _{i=1}^N{w}_{pi},\sum _{j=1}^N{w}_{qj}\right)=\min \left(N\cdot \frac{1}{N},N\cdot \frac{1}{N}\right)=1$$

$$\sum _{i=1}^N\sum _{j=1}^N{d}_{ij}{f}_{ij}=N\cdot \frac{1}{N}\sum _{i=1}^N\sum _{j=1}^N{d}_{ij}=\sum _{i=1}^N\sum _{j=1}^N{d}_{ij}$$

因此，EMD 距离改写为：

$$EMD(P,Q)=\sum _{i=1}^N\sum _{j=1}^N{d}_{ij}$$

也就是说，在神经网络中选择 EMD 作为损失函数时，就是在点集 $P$ 和 $Q$ 中寻找一个一一对应的关系使得 EMD 最小，即：

$$\min EMD(P,Q)=\min \sum _{i=1}^N\sum _{j=1}^N{d}_{ij}$$

其实，也就是一般在论文中看到的那样：

$$Los{s}_{EMD}(P,Q)=\underset{\varphi :P\to Q}{\min }\sum _{x\in P}\Vert \Vert x-\varphi (x)\Vert {\Vert }_2$$

就是在点集 $P$ 和 $Q$ 中间找到一个双射 $\varphi$，将两个点集一一对应起来，使得二者计算欧式距离的和最小。这就是一般我们在点云补全等论文中看到的 EMD 作为损失函数形式的由来了。