DINO结构中的exponential moving average (ema)和stop-gradient (sg)

news2024/10/6 10:35:22

DINO思路介绍

在这里插入图片描述

在 DINO 中，教师和学生网络分别预测一个一维的嵌入。为了训练学生模型，我们需要选取一个损失函数，不断地让学生的输出向教师的输出靠近。softmax 结合交叉熵损失函数是一种常用的做法，来让学生模型的输出与教师模型的输出匹配。具体地，通过 softmax 函数把教师和学生的嵌入向量尺度压缩到 0 到 1 之间，并计算两个向量的交叉熵损失。这样，在训练过程中，学生模型可以通过模仿教师模型的输出来学习更好的特征表示，从而提高模型的性能和泛化能力。当然，这也可以看作是一个分类问题，以便网络可以从局部视图中学习更有意义的全局表示。

论文对图片中DINO结构的解释如下

我们用一个简单的视角对 (x1, x2) 来说明DINO。在该模型中，输入图像经过两种不同的随机变换后分别传递给学生网络和教师网络。两个网络具有相同的架构但参数不同。教师网络的输出以批次的均值进行中心化。每个网络输出一个K 维的特征，并通过特征维度上的温度 softmax进行归一化。然后通过交叉熵损失函数来测量它们的相似性。我们在教师网络上应用一个停止梯度（stop-gradient,sg）操作，以仅通过学生网络传播梯度。教师网络的参数通过学生网络参数的指数移动平均（ema）来更新。

核心特点

无标签自蒸馏： 不需要人工标注的数据，通过模型自身的知识传递来进行训练。
双视角 (x1, x2)： 输入图像经过两种不同的随机变换，生成两种视角，以增加数据的多样性和模型的鲁棒性。
学生网络和教师网络： 两个网络结构相同，但参数不同。学生网络用于训练，教师网络提供稳定的指导信号。
教师网络输出中心化： 教师网络的输出在批次内进行中心化处理，以消除偏差并稳定训练过程。
温度 Softmax 归一化： 网络输出通过温度 Softmax 进行归一化，控制特征向量的平滑度，防止梯度消失或爆炸。
相似性度量： 通过交叉熵损失函数测量学生网络和教师网络输出之间的相似性，鼓励学生网络学习到与教师网络相似的表示。
停止梯度操作（SG）： 在教师网络上应用停止梯度操作，确保梯度只通过学生网络进行反向传播，避免教师网络影响梯度更新。
指数移动平均更新（EMA）： 教师网络的参数通过学生网络参数的指数移动平均进行更新，确保教师网络的参数更加平滑和稳定，从而提供可靠的指导信号。

再总结一下：在两个完全一样的教师和学生网络（ViT/CNN 均可）中

教师网络通过 centering 和 sharpening 正则化避免训练崩塌。
- - Centering：教师模型的输出也经过 EMA 操作，从原始激活值中减去一个平均值。
- - Sharpening：在 Softmax 中加入一个 temperature 参数，强制让概率分布更加尖锐。
- 两个网络的输出都通过 Softmax 层归一化处理。
- 通过交叉熵损失计算损失。
学生网络通过 SGD 更新参数，并通过 EMA 更新教师网络参数，教师网络的参数因为sg不会自己动**。

问题一：centering中使用的EMA（Exponential Moving Average）是什么东西,指数是怎么被体现的

在 DINO 中，教师网络的参数通过学生网络参数的指数移动平均（Exponential Moving Average, EMA）进行更新。以下是 EMA 操作的详细解释：

1. 定义

EMA 是一种用于平滑时间序列数据的技术，它通过对新数据赋予较高权重，同时对旧数据赋予较低权重，从而平滑数据变化。具体来说，EMA 的计算公式为：

$\theta_t^{\text{teacher}} = \alpha \theta_t^{\text{student}} + (1 - \alpha) \theta_{t-1}^{\text{teacher}}$

其中：

$\theta_t^{\text{teacher}}$ 是第 $t$ 次更新后的教师网络参数。
$\theta_t^{\text{student}}$ 是第 $t$ 次学生网络的参数。
$\alpha$ 是平滑系数，介于 0 和 1 之间，通常取一个较小的值（例如 0.99 或 0.999）。
$\theta_{t-1}^{\text{teacher}}$ 是第 $t - 1$ 次更新后的教师网络参数。

2. 指数的体现

EMA 操作中的“指数”体现在计算过程中旧数据的权重以指数形式递减。具体地，如果我们展开几次 EMA 的更新，可以看到：

$\theta_t^{\text{teacher}} = \alpha \theta_t^{\text{student}} + \alpha(1 - \alpha) \theta_{t-1}^{\text{student}} + \alpha(1 - \alpha)^2 \theta_{t-2}^{\text{student}} + \cdots$