Read Like Humans: Autonomous, Bidirectional and Iterative LanguageModeling for Scene Text Recognition

Posted on In Views:

前言

在文本识别领域，如何在端到端的训练任务中使用语言信息是一个具有挑战性的任务，在本篇论文中，作者认为语言模型的限制在于三个方面：1. 隐式的语言建模 2. 特征表达没有方向性 3. 语言模型的输入有噪声。于是作者提出了ABINet，针对这三个方面进行改善。1. 将视觉和语言模型的梯度反传隔绝开，来显式地加强语言模型 2. 使用双向特征表达 3. 采用迭代的方式来优化语言模型的输入。

此外，作者还提出了可以有效使用无标注数据的自监督训练方法

论文：Read Like Humans: Autonomous, Bidirectional and Iterative LanguageModeling for Scene Text Recognition

代码：ABINet

背景

作者根据人类阅读文本的特性，提出了设计文本识别算法需要参考的三个准则：

  1. 自主：鉴于聋哑人和盲人都可以分别拥有完整的视觉和语言功能，因此视觉和语言可以独立进行学习，同时这也意味着视觉与语言之间有良好的相互作用，独立学习语言知识有助于视觉中的字符识别。
  2. 双向：根据文本上下文推理看不清的字符，类似于完形填空任务，因此，可以同时利用不可读字符的左右两侧的可读字符线索进行预测，这与双向性相对应。
  3. 迭代：在具有挑战性的环境下，人类会采用渐进式策略，通过迭代修正已识别的结果，来提高预测的置信度

将这三个准则应用到文本识别算法，则需要进行以下设计：

  1. 文本模型和视觉模型应该分开单独训练。现有的算法大多将语言模型隐式的嵌入模型中(图a)，因此，语言模型是否学习到了字符之间的关系，以及学习的程度无法衡量
  2. 双向的语言模型较没有方向的语言模型，可以学到2倍的信息量。最显而易见的方法为建立从左到右和从右到左两个模型(图e)，但由于语言特征是没有方向的，所以这样做没有很大提升，而且会增加计算量。直接将Bert方法应用于文本识别也比较难，因为每次只能mask掉一个字符
  3. 使用迭代的方式可以优化语言模型的输出，但现有的模型很少使用该技巧。现有的算法有些采用自回归解码(图d)，该方法会导致误差累计，有些采用并行解码，该方法中视觉模型的误差会降低语言模型的性能，且会面临mis-align的问题

因此，根据上述三个原则，本文提出了ABINet，主要从以下几个方面

  1. 将视觉模型和语言模型的梯度反传进行阻断(图b)，强制语言模型显示地学习语言规则。同时视觉模型和语言模型可以分开单独进行训练
  2. 设计了一个双向填空网络BCN(图c)，它同时以左右上下文为条件，通过指定注意力掩码来控制对两边字符的访问。此外，为了防止信息泄露，不允许跨字符访问
  3. 提出了通过迭代修正语言模型输出结果的方法。通过将ABINet的输出反复输入语言模型，逐步改善预测，并在一定程度上缓解长度不对齐问题。此外，通过将迭代预测视为一种ensemble，本文探索了一种基于自我训练的半监督方法

框架

视觉模型

使用ResNet+Transformer组合作为backbone，分别用于提取特征以及序列化。backbone提取的特征为
$$
\mathbf{F}_b=\mathcal{T}(\mathcal{R}(\boldsymbol{x})) \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times C}
$$
然后通过位置编码模块，即：
$$
\mathbf{F}_v=\operatorname{softmax}\left(\frac{\mathbf{Q K}^{\top}}{\sqrt{C}}\right) \mathbf{V} .
$$
其中$\mathbf{Q} \in \mathbb{R}^{T\times C}$为位置编码，$\mathbf{K}=\mathcal{G}\left(\mathbf{F}_b\right) \in \mathbb{R}^{\frac{H W}{16} \times C}$，为$\mathbf{F}_b$经过U-Net后得到，$\mathbf{V}=\mathcal{H}\left(\mathbf{F}_b\right) \in \mathbb{R}^{\frac{H W}{16} \times C}$为$\mathbf{F}_b$经过Identity Mapping后得到

语言模型

自主性

语言模型需要有以下特性：1）是一个独立的模型，用于纠正预测错误的字符，2）梯度反传在输入处被阻断(无法传播给视觉模型)，3）可以被未标注文本单独训练

由于语言模型的自主性，它可以随时被更加优秀的语言模型替代，且梯度反传的阻断会强制其显式地学习语言知识，另外语言模型可以受益于NLP算法的发展，如大模型的预训练权重等

双向特征表达

假设输入文本$\bold{y}=(y_1,…,y_n)$长度为$n$，字符的种类为$c$，对于$y_i$，传统的双向和无向模型可以表达为$P\left(y_i \mid y_n, \ldots, y_{i+1}, y_{i-1}, \ldots, y_1\right)$以及$P\left( y_{i-1}, \ldots, y_1\right)$，其entropy分别为$H_y=(n-1) \log c$以及$\frac{1}{n} \sum_{i=1}^n(i-1) \log c = \frac{1}{2}H_y$，及双向模型本质上为两个无向模型的ensemble，且无向模型本身的表达能力就有限

直接将Bert中的MLM方法应用到文本识别中，即每次讲文本中的一个字符$y_i$替换为$\text{[MASK]}$，由于每次仅能mask掉一个字符，因此如果要预测文本中的每个字符，需要重复调用$n$次，计算开销较大。因此作者提出了BCN

在该网络中，字符被直接输入multi-head attention block，除此之外，在multi-head attention中还加入了额外的attention mask来防止看到当前字符，除此之外，为了防止信息泄露，BCN中没有使用self-attention。整体的计算公式如下所示：
$$
\begin{aligned}
\mathbf{M}{i j} & = \begin{cases}0, & i \neq j \
-\infty, & i=j\end{cases} \
\mathbf{K}_i & =\mathbf{V}_i=P\left(y_i\right) \mathbf{W}_l \
\mathbf{F}{m h a} & =\operatorname{softmax}\left(\frac{\mathbf{Q K}^{\top}}{\sqrt{C}}+\mathbf{M}\right) \mathbf{V}
\end{aligned}
$$
其中$\mathbf{Q} \in \mathbb{R}^{T\times C}$为位置编码，$P\left(y_i\right)$为字符串的概率分布，$\mathbf{W}_l \in \mathbb{R}^{c \times C}$为线性映射矩阵，$\mathbf{M} \in \mathbb{R}^{T \times T}$为attention mask，防止当前字符参与到计算中。经过多个BCN层的堆叠，可以获得用于计算$y$的$F_l$

通过类似完形填空的方式，BCN可以学到比双向网络更多的表达。得益于transformer的结果，BCN可以进行并行独立的计算。此外，BCN模型比ensemble更有效率，因为只需要花费一半的计算量

循环校正

为了解决输入噪声的问题，作者提出了循环语言模型，他会进行$M$次循环，首次循环的输入为视觉模型的输出，后续的输入为上一轮循环中，视觉模型和语言模型的混合输出。通过这种方式，语言模型可以对输出结果进行循环校正

如果视觉模型输出的文本长度和GT不同，transformer无法进行很好的校正视觉模型的输出，这是由于文本预测中使用的padding mask导致的。本文的语言模型可以较好的缓解该问题，因为该方法中，视觉模型的输出和语言模型的输出会混合多次，所以预测的文本长队会进行多次refine

混合网络

为了将语言模型和视觉模型两种模态的输入映射到同一空间，作者采用了以下方式将两种特征进行混合：
$$
\begin{aligned}
\mathbf{G} & =\sigma\left(\left[\mathbf{F}_v, \mathbf{F}_l\right] \mathbf{W}_f\right) \
\mathbf{F}_f & =\mathbf{G} \odot \mathbf{F}_v+(1-\mathbf{G}) \odot \mathbf{F}_l
\end{aligned}
$$
其中，$\mathbf{W}_f \in \mathbb{R}^{2 C \times C}$，且$ \mathbf{G} \in \mathbb{R}^{T \times C}$

有监督的训练

优化函数为：
$$
\mathcal{L}=\lambda_v \mathcal{L}_v+\frac{\lambda_l}{M} \sum_{i=1}^M \mathcal{L}l^i+\frac{1}{M} \sum{i=1}^M \mathcal{L}_f^i
$$
损失函数均为cross entropy，不同的下标代表对不同网络的监督(视觉模型/语言模型/混合模型)，$i$代表不同的迭代次数

半监督训练

作者采用了self-training的方法，即先使用模型预测伪标签，然后使用这些伪标签进一步训练模型，该策略的关键为如何获得较高质量的伪标签

作者设计了一个衡量文本置信度的指标：
$$
\begin{cases}\mathcal{C} & =\min _{1 \leq t \leq T} e^{\mathbb{E}\left[\log P\left(y_t\right)\right]} \ P\left(y_t\right) & =\max _{1 \leq m \leq M} P_m\left(y_t\right)\end{cases}
$$
对于预测文本伪标签中的每个字母$y_t$，选择在$M$此迭代中最大的概率$P(y_t)$，然后选择整个文本中最小的$P(y_t)$作为文本的置信度。当其小于阈值时，这个样本被删除。策略如下所示

实验

消融实验

视觉模型

  1. 只扩大视觉模型即可提高性能，但会降低速度
  2. 升级视觉模型可以使用更深的transformer以及使用position attention

语言模型

自主策略：

对比实验结果可以发现：

  1. 视觉模型的pretrain能够提高性能 (第一行 vs 第二行)
  2. 在MJ和ST上pretrain语言模型收益不大 (第一行 vs 第三行，第二行 vs 第四行)
  3. 使用额外的语言数据来pretrain语言模型有较高的收益(第二行 vs 第五行)
  4. 允许语言模型和视觉模型间的梯度传播会损害性能(第二行 vs 第六行)。作者认为这是由于视觉模型有助于语言模型的收敛，从而使语言模型过拟合了(Loss下降很快)。因此需要阻断语言模型和视觉模型间的梯度传播

双向特征表达：

作者对比了BCN、SRN(双向特征表达)和SRN-U(无向特征表达)，并使用两种不同大小的backbone来进行比较。与SRN-U相比，BCN的参数和计算速度相似，但性能更好；与SRN相比，BCN性能更好且运算速度更快

之前提到语言模型可以视为一个单独的文本校正网络，为了验证这个猜测，作者将MJ和ST作为测试集，并随机采样了20000个文本，对其中的部分文本做了删除/替换一个字符的操作，并使用语言模型校正，BCN网络的性能较SRN提高了4.5%

为了更好的探究BCN如何工作，作者将top-5的准确率进行了可视化，当只有一个字母出错时，BCN可以以很高的概率预测遮盖掉的字母，这样有助于文本识别。但当2个字母出错时，对于其他字母来说噪音太大，BCN无法进行很好的预测

循环校正：

作者使用不同大小的模型验证循环校正的作用，发现进行3次循环可以提高性能，且在越难的数据集(IC15,SVT,SVTP)上表现的越好。

对比训练和测试时不同的次数可以发现：

  1. 训练中不使用循环校正(iter=1)，直接在测试中使用循环校正(iter>1)效果也可以
  2. 训练中使用循环校正有助于提升性能，因为可以提供更多的sample
  3. 循环次数超过三次就饱和了。

将实验结果可视化后发现，视觉模型预测的字符串存在错误，经过多次的校正，最终预测结果正确

经过以上的消融实验可以发现：

  1. BCN是一种有效的语言模型，可以提升速度和准确率(相比双向模型)
  2. BCN+循环校正可以纠正错误的预测结果

与SOTA对比

与SRN对比，性能要优于SRN，且同时超过了其他SOTA算法。由于语言模型+循环校正的帮助，ABINet在含有较多低质图片(噪音、模糊)的数据集SVT、SVTP、IC15上提升更加明显

半监督学习

作者将ST和MJ用于监督学习，并在Uber-Text数据集上使用半监督方法训练，即$\text{ABINet-LV}{st}$，性能较$\text{ABINet-LV}$ 有所提升，进行ensemble后，$\text{ABINet-LV}{est}$性能进一步提升