ABINet 阅读笔记
CVPR 2021 Oral
Paper : https://arxiv.org/abs/2103.06495
Code : https://github.com/FangShancheng/ABINet (即将开源)
论文总览
基于CTC和基于分割的方法，由于缺乏语义信息，所以不能解决低质量图像的识别问题。而基于语言模型的方法有基于 Attention （RNN/Transformer）的方法，而这些方法是隐性的、单向的，是对视觉特征的上下文建模，而非语言特征的上下文建模。
所以论文围绕如何进行有效地建模语言这个问题展开了一些探索及讨论，总体思想就是从 NLP 字符级别级校正的语言建模角度来思考场景文字识别问题，尤其是针对图像质量退化的这种场景文字识别。文中的语言建模主要是 character 级别的建模，而非 NLP 中更常用的 word 级别或者 token 级别的语言建模。
以自治，双向，迭代为指导的ABINet：
- 首先，自治（Autonomous）显性定义视觉模型与语言模型，视觉模型只负责图像分类的功能，而语言模型只负责语言建模的功能。
- 其次，基于双向（Bidirectional）特征表示的 BCN（Bidirectional cloze network）作为语言模型，消除了组合两个单向模型的问题。通过指定注意掩码来控制双方字符的访问，从而在左右上下文中共同受到限制。同样，不允许跨时间步访问以防止泄漏信息。
- 第三，语言模型迭代（Iterative）校正的执行方式，可以有效减轻噪声输入的影响。通过将 ABINet 的输出反复输入到 LM，可以逐步完善预测，并可以在一定程度上缓解长度不齐的问题。 另外，基于迭代预测的集合，还提出了一种可以有效地从未标记图像中学习的自训练方法。
Method
Vision Model
视觉模型包括一个主干网络和一个位置注意力模块。使用 ResNet 和 Transformer 作为特征提取网络和序列建模网络。
F b = T ( R ( x ) ) ∈ R H 4 × W 4 × C (1) F_b = \mathcal{T}(\mathcal{R}(x)) \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times C} \tag{1} Fb=T(R(x))∈R4H×4W×C(1)
其中 x x x 是输入图像， H , W H, W H,W 是输入图像的大小， C C C 是特征维度
位置注意模块将视觉特征并行转换为字符概率：
F v = softmax ( Q K T C ) V (2) F_v = \text{softmax}(\frac{QK^T}{\sqrt{C}})V \tag{2} Fv=softmax(CQKT)V(2)
Q ∈ R T × C Q \in \mathbb{R}^{T \times C} Q∈RT×C 是字符顺序的位置编码， T T T 是字符序列的长度， K = G ( F b ) ∈ R H W 16 × C K = \mathcal{G}(F_b) \in \mathbb{R}^{\frac{HW}{16} \times C} K=G(Fb)∈R16HW×C 其中 G ( ⋅ ) \mathcal{G}(\cdot) G(⋅) 是通过小型的 U-Net 网络实现的。 V = H ( F b ) ∈ R H W 16 × C V = \mathcal{H}(F_b) \in \mathbb{R}^{\frac{HW}{16} \times C} V=H(Fb)∈R16HW×C 其中 H ( ⋅ ) \mathcal{H}(\cdot) H(⋅) 是特征映射
Language Model
自治策略
自治性在实现上，就是显性定义视觉模型与语言模型，视觉模型只负责图像分类的功能，而语言模型只负责语言建模的功能。视觉模型与语言模型的交互方式为视觉模型的预测概率分布直接作为语言模型的输入。语言模型的输入为概率向量，并使用线性函数实现概率映射（Probability Mapping）。二者分别使用损失函数进行训练，最终结果进行融合。其中，一个核心的点为视觉模型的输出概率向量到语言模型的输入概率向量，通过梯度阻塞的方式使得视觉模型跟语言模型进行分离，进而实现两个模型作为独立的功能个体进行学习，即实现显性地语言建模型。
这样做的好处其实有很多：
- 解藕之后，视觉模型跟语言模型就互相独立了。语言模型在做语言推理的适合也不再依赖于视觉特征，视觉模型跟语言模型都是独立的个体，比如说语言模型的语言修正功能可以直接使用和单独训练，这个是隐性的方法不具备的。因此，视觉模型可以单独训练（有监督/无监督均可），而语言模型更可以轻易从海量文本中无监督预训练学习。这就是为啥是叫自治的模型，而非解藕的模型。
- 强大的可解释性，我们甚至可以直接评估语言模型的性能，比如在实验章节单独针对语言模型展开定量以及可视化定性的评估，在这样的模式下，今后可以针对语言模型更加具体的问题提出解决方案。
- 足够的灵活性，一方面今后如果有更优秀的语言模型能代替 ABINet 中的 BCN 语言模型，那么可以在不调整其他结构比如视觉模型的前提下直接取代 BCN；以及在端到端识别中，只需要替换视觉模型，而语言模型并不需要做改变（甚至已经学习到的模型可以直接拿来做 fine-tune），这样成本就低很多了。
- 更好的拓展性，由于输入输出都是定义成概率的形式，拓展性就很强了，比如说在整体框架下可以很简单的实现第三点迭代性的建模。
- 当然，最大的优势是强迫语言模型真正地学习语言规则方面的建模。不像隐性语言模型那样，可能存在视觉特征作为混杂因子的某种关系导致了模型的偏倚，在反向传播过程中，某条作弊的路径让语言模型的学习是次优的。
双向表示
给定长度为 n ，类别为 c 的文本串 y = ( y 1 , . . . , y n ) y = (y_1, ..., y_n) y=(y1,...,yn) ，对于双向和单向的 y i y_i yi 的条件概率分别为 P ( y i ∣ y n , … , y i + 1 , y i − 1 , … , y 1 ) P(y_i | y_n,\dots, y_{i+1}, y_{i-1}, \dots, y_1) P(yi∣yn,…,yi+1,yi−1,…,y1) 和 P ( y i ∣ y i − 1 , … , y 1 ) P(y_i | y_{i-1}, \dots, y_1) P(yi∣yi−1,…,y1) 。从信息论的角度来看，双向表示的可用熵可以量化为 H y = ( n − 1 ) log c H_y = (n-1)\log c Hy=(n−1)logc。 但是，对于单向表示，信息为 1 n ∑ i = 1 n ( i − 1 ) log c = 1 2 H y \frac{1}{n} \sum_{i=1}^n (i-1)\log c = \frac{1}{2} H_y n1∑i=1n(i−1)logc=21Hy。 我们的见解是，以前的方法通常使用两个单向模型的集合模型，这些模型本质上是单向表示。 单向表示基本上捕获了 1 2 H y \frac{1}{2} H_y 21Hy 信息，与双向对应相比，特征抽象的能力有限。
得益于自治策略，可以使用具有拼写校正功能的现成 NLP 模型。 一种可行的方法是通过利用 BERT 中的掩模语言模型（MLM）使用 [MASK] 替换 y i y_i yi 。 但是这是不可行的，因为对于每个文本实例，应分别将 MLM 调用 n 次，这会导致极低的效率。所以提出 BCN 来指定注意掩码而不是掩模输入字符。
总体上，BCN（Bidirectional Cloze Network）是一个基于完型填空思想，作拼写矫正的语言模型。实现上是一个 L 层的 Transformer decoder 的变种，注意是decoder变种，而非 encoder，也不是一些方法使用的encoder+decoder。为了实现双向的建模，BCN在很大程度上跟Transformer又有区别：
-
该网络以字符位置序号编码作为输入，为非字符概率向量。而字符概率向量直接传入multi-head attention模块。
-
网络通过multi-head attention模块中的对角注意力掩码mask控制字符信息的访问。对于第 t 个时间步的字符来说，注意力机制通过对角mask避免看到当前字符，且实现同时访问该字符左边及右边的信息，并综合左边及右边的信息同时做出预测。
-
为了避免信息泄露，传统Transformer decoder中的自注意模块self-attention并没有在BCN中使用，这样避免了跨时间步上的信息访问。因此BCN的每个时间步的计算均为独立且并行的，也具有高效的特点。
3 这一点是很关键的一点，设想下如果第 t 时间步尽管在网络输入的时候没有看到前后时间步的信息，但却在self-attention跨时间步的信息交互上看到了前后时间步的信息，那么语言模型就不需要费劲去预测了，只需要直接将‘偷看’的信息恒等映射即可，如果有self-attention，效果立竿见影地降低，这也是一些基于Transformer的方法效果不好的原因。此外，1中也不能将字符概率向量作为BCN的直接输入，仔细揣摩下会发现这种情况下Transformer的残差连接也会信息泄漏。因此，BCN的一切，都是为了精准实现 P ( y i ∣ y n , … , y i + 1 , y i − 1 , … , y 1 ) P(y_i | y_n,\dots, y_{i+1}, y_{i-1}, \dots, y_1) P(yi∣yn,…,yi+1,yi−1,…,y1) 这个公式。
在多头注意力的操作：
M i j = { 0 , i ≠ j − ∞ , i = j (3) M_{ij} = \begin{cases} 0, & i \neq j \\ -\infty, & i = j \end{cases} \tag{3} Mij={
0,−∞,i=ji=j(3)
K i = V i = P ( y i ) W l (4) K_i = V_i = P(y_i) W_l \tag{4} Ki=Vi=P(yi)Wl(4)
F m h a = softmax ( Q K T C + M ) V (5) F_{mha} = \text{softmax}(\frac{QK^T}{\sqrt{C}} + M) V \tag{5} Fmha=softmax(CQKT+M)V(5)
其中 Q ∈ R T × C Q \in \mathbb{R}^{T \times C} Q∈RT×C 是第一层或最后一层的输出字符序列的位置编码， K , V ∈ R T × C K, V \in \mathbb{R}^{T \times C} K,V∈RT×C 是从字符概率 P ( y i ) ∈ R c P(y_i) \in \mathbb{R}^c P(yi)∈Rc 中获得的， W l ∈ R c × C W_l \in \mathbb{R}^{c \times C} Wl∈Rc×C 是线性映射权重。 M ∈ R T × T M \in \mathbb{R}^{T \times T} M∈RT×T 是阻止添加当前字符的 attention masks 的矩阵。在将 BCN 层堆叠到深度架构中之后，确定文本 y y y 的双向表示 F l F_l Fl。
迭代校正
Transformer 的并行预测采用噪声输入，这些噪声输入通常是视觉预测或视觉特征的近似值。具体地，如图 ABINet 所示的双向表示的例子，P(“O”) 的期望条件是“ SH-WING”。但是，由于环境模糊和遮挡，从 VM 获得的实际条件是 “SH-VING”，其中 “V”成为噪声并损害了预测的置信度。 随着 VM 中错误预测的增加，它对 LM 影响更大。
为了解决噪声输入的问题，我们提出了迭代 LM（图 ABINet 所示）。 LM 以 y y y 的不同分配重复执行 M 次。 对于第一次迭代， y i = 1 y_{i=1} yi=1 是来自 VM 的概率预测。 对于后续迭代， y i ≥ 2 y_{i \geq 2} yi≥2是根据上次迭代中融合模型的概率预测。 通过这种方式，LM可以迭代地校正视觉预测。
另一个观察结果是，基于 Transformer 的方法通常会遇到长度不对齐的问题，这表示如果字符数与 ground truth 不对齐，则Transformer很难校正视觉预测。 长度不对齐的问题是由于不可避免地采用填充掩码而引起的，该掩码对于过滤超出文本长度的上下文是固定的。 迭代 LM 可以缓解此问题，因为视觉特征和语言特征会多次融合，因此预测的文本长度也会逐渐细化。
融合
在图像上训练的视觉模型和在文本上训练的语言模型来自不同的模式。 为了使视觉特征和语言特征保持一致，仅使用门控机制进行最终决策：
G = σ ( [ F v , F l ] W f ) (6) G = \sigma ([F_v, F_l]W_f) \tag{6} G=σ([Fv,Fl]Wf)(6)
F f = G ⊙ F v + ( 1 − G ) ⊙ F l (7) F_f = G \odot F_v + (1-G) \odot F_l \tag{7} Ff=G⊙Fv+(1−G)⊙Fl(7)
其中 W f ∈ R 2 C × C W_f \in \mathbb{R}^{2C \times C} Wf∈R2C×C， G ∈ R T × C G \in \mathbb{R}^{T \times C} G∈RT×C
监督训练
ABINet 使用一下的多任务目标函数端到端训练
L = λ v L v + λ l M ∑ i = 1 M L l i + 1 M ∑ i = 1 M L f i (8) \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} \tag{8} L=λvLv+Mλli=1∑MLli+M1i=1∑MLfi(8)
其中 L v , L l , L f \mathcal{L}_v, \mathcal{L}_l, \mathcal{L}_f Lv,Ll,Lf 分别是 F v , F l , F f F_v, F_l, F_f Fv,Fl,Ff 的交叉熵损失。 L l i , L f i \mathcal{L}_l^i, \mathcal{L}_f^i Lli,Lfi 是第 i 个迭代的损失， λ v , λ l \lambda_v, \lambda_l λv,λl 是平衡因子
半监督整体训练
为了进一步探索迭代模型的优越性，提出了一种基于自我训练的半监督学习方法。 自我训练的基本思想是，首先通过模型本身生成伪标签，然后使用额外的伪标签对模型进行重新训练。 因此，关键问题在于构建高质量的伪标签。
为了过滤噪声伪标签，我们提出以下方法：1）选择文本实例内字符的最小置信度作为文本确定性。 2）每个字符的迭代预测被看作是使噪声标签的影响平滑的整体。 因此，我们定义过滤功能如下：
{ C = m i n 1 ≤ t ≤ T e E [ log P ( y t ) ] P ( y t ) = m a x 1 ≤ m ≤ M P m ( y t ) (9) \begin{cases} \mathcal{C} & = \underset{1 \leq t \leq T}{min} e^{\mathbb{E}[\log P(y_t)]} \\ P(y_t) & = \underset{1 \leq m \leq M}{max} P_m (y_t) \end{cases} \tag{9} ⎩⎨⎧CP(yt)=1≤t≤TmineE[logP(yt)]=1≤m≤MmaxPm(yt)(9)
其中 C \mathcal{C} C 是文本实例的最小确定性， P m ( y t ) P_{m}(y_t) Pm(yt) 是第 t 个字符在第 m 次迭代的概率分布。训练流程见下图， Q Q Q 是阈值， B l , B u B_l, B_u Bl,Bu 是标注和未标注数据的训练批次， N m a x N_{max} Nmax 是训练的最大步数， N u p l N_{upl} Nupl 是更新伪标签的步数。
Experiment
模型的维度 C = 512 C = 512 C=512，BCN 有 4 层，每一层有 8 个 注意力头。 λ v , λ l \lambda_v, \lambda_l λv,λl 都为1。图像使用数据增强缩放至 32 × 128 32 \times 128 32×128。
Conclusion
在本文中，提出了ABINet，它探索了在场景文本识别中利用语言知识的有效方法。 ABINet是
- 自主，通过明确地强制学习来提高语言模型的能力；
- 双向，通过共同限制双方的字符上下文来学习文本表示；
- 迭代，逐步校正预测以减轻噪声输入的影响。
基于ABINet，进一步提出了一种用于半监督学习的整体自训练方法。
参考文章