此前的 Temporal ensembling Model 中存在一个问题, 即在每个 epoch 只进行一次 EMA, 因此在学习大型数据集时, Temporal Ensembling 成变得比较笨拙, 且无法实现模型的在线训练. 为了克服这个问题, 文中提出了 Mean Teacher, 能在每个 epoch 中的每个 step 进行模型权重的 EMA. 也就是将原来计算输出向量 z ~ \tilde{z} z~ 的过程变成了计算整个网络的参数 θ \theta θ.
论文地址: Mean teachers are better role models
代码地址: https://github.com/CuriousAI/mean-teacher
会议: NeurIPS 2017
任务: 分类
Mean Teacher
Mean Teacher 思想比较简单, 是对 Temporal ensembling 的改进. 与 Temporal Ensembling Model 相比, 在 Mean Teacher 中, 将原始模型看作 Student Model, 然后新添加一个架构相同的 Teacher Model. Teacher Model 更加健壮, 可以指导 Student Model 学习. 其中, Teacher 的参数由 Student 计算指数移动平均值(EMA)得到, 见式(2).
Mean Teacher 中一致性损失
J
J
J 定义如下:
J
(
θ
)
=
E
x
,
μ
′
,
μ
[
∣
∣
f
(
x
,
θ
′
,
μ
′
)
−
f
(
x
,
θ
,
μ
)
)
∣
∣
2
]
(1)
J(\theta)=\mathbb{E}_{x,\mu',\mu}[\vert \vert f(x,\theta',\mu')-f(x,\theta,\mu))\vert\vert^2] \tag{1}
J(θ)=Ex,μ′,μ[∣∣f(x,θ′,μ′)−f(x,θ,μ))∣∣2](1)
θ
t
′
=
α
θ
t
−
1
′
+
(
1
−
α
)
θ
t
(2)
\theta'_t=\alpha\theta'_{t-1}+(1-\alpha)\theta_t \tag{2}
θt′=αθt−1′+(1−α)θt(2)
其中,
θ
t
\theta_t
θt,
θ
t
′
\theta'_t
θt′ 分别为 Student 和 Teacher 的参数,
μ
\mu
μ,
μ
′
\mu'
μ′ 分别为 Student 和 Teacher 的扰动(噪声). 当
α
=
0
\alpha=0
α=0 时, Mean Teacher Model 与
Π
\Pi
Π-model 在形式上等价.
算法流程
Mean Teacher 算法模型如下:
对于有标记样本
{
(
x
i
,
y
i
)
}
\{(x_i,y_i)\}
{(xi,yi)}:
- 1.将样本 x i x_i xi 输入到 Student Model, 然后得到预测标签 y ^ i \hat{y}_i y^i.
- 2.构造损失函数 L 1 ( y i , y ^ i ) L_1(y_i,\hat{y}_i) L1(yi,y^i), 一般使用交叉熵.
对于所有样本 { ( x l ) } \{(x_l)\} {(xl)}:
- 1.将 { ( x l ) } \{(x_l)\} {(xl)} 同时输入到 Student 及 Teacher Model, 然后得到预测标签 y l 1 y_{l1} yl1, y l 2 y_{l2} yl2.
- 2.构造 y l 1 y_{l1} yl1, y l 2 y_{l2} yl2 的损失函数 L 2 L_2 L2, 文中使用 MSE, 见上面式(1).
- 3.通过总损失函数 L 1 + λ L 2 L_1+\lambda L_2 L1+λL2 梯度下降, 更新 Student Model 参数 θ \theta θ, 然后通过 EMA 更新 Teacher Model 参数 θ ′ \theta' θ′.