Visual-Semantic Transformer for Scene Text Recognition

Text recognition has been an active research area for decades. See (Long, He, and Yao 2021; Chen et al. 2021) for comprehensive reviews. Due to page limitation, we can only list a portion of recent work here.

Recently, (Nguyen et al. 2021) incorporate a dictionary in training and inference stage to help selecting the most compatible outcome for STR. (Feng et al. 2021) introduce character center segmentation branch to extract semantic features which encode the category and position of characters for improving video text detection performance. (Patel et al. 2016) propose to generate contextualized lexicons for scene images with only visual information. (Sabir, Moreno-Noguer, and Padró 2018) use language model to build the semantic correlation between scene and text in order to re-rank the recognition results. (Zheng, Wang, and Betke 2019) also propose to use pretrained language models for correcting predictions. Very recently, (Wang et al. 2021b) propose to learn the linguistic rules in the visual space by randomly mask out some characters from the image and predict them back.

Similar to speech recognition, scene text recognition can be treated as a sequence-to-sequence (seq2seq) mapping problem (Jaderberg, Vedaldi, and Zisserman 2014; Jaderberg et al. 2015; Qiao et al. 2020; Yue et al. 2020). (Li et al. 2019) combine convolution and LSTM as an encoder, then use another LSTM as decoder to predict text attentively, quite similar to the show, attend and tell work on image captioning (Xu et al. 2015). ASTER (Shi et al. 2018) takes a two-stage approach to first rectify curved text images and then perform recognition using seq2seq model with attention. CRNN (Shi, Bai, and Yao 2016) adapts CNN to obtain visual features, which are then fed into LSTM module with CTC loss (Graves et al. 2006) for text prediction. STAR-Net (Liu et al. 2016) uses spatial transformer (Jaderberg et al. 2015) to tackle challenges brought by image distortion. Attention can be added to the seq2seq models (Li et al. 2019; Bhunia et al. 2021) in a straightfoward way to alleviate bottleneck effects brought by seq2seq models. (Litman et al. 2020) propose stacked block architecture with intermediate supervision to train a deep BiLSTM encoder, while attention is used in decoding stage to exploit contextualized visual features. (Aberdam et al. 2021) propose seq2seq contrastive learning of visual representations which can be applied to text recognition. DAN (Wang et al. 2020) decouples alignment from the decoding stage into the early conventional encoder network. (Wang et al. 2021a) propose an alignment module enabling text recognizer to recognize document level image. (Wang et al. 2019b) use adversarial loss to handle low-resolution image.

Transformers have also been successfully applied to STR. ABINet (Fang et al. 2021) enforces the bidirectional language-model (LM) to only learn linguistic rules by gradient-stopping in training. The decoding is in an iterative way allowing the predictions to be refined progressively. Their conv+transformer visual module and transformer-LM can be separately pretrained to improve performance. HRGAT (Yang et al. 2020) connects CNN feature maps to a transformer-based autoregressive decoder, where the cross-attention is guided by holistic representation obtained by average-pooling of 2d feature maps.

Perhaps more related, SRN (Yu et al. 2020) incorperates visual-to-semantic embedding block and cross-entropy loss to align with ground-truth text, but they use argmax embedding, which is different from our direct use of probability vector that enables smooth gradient flow in training. Our work is also innovative in many ways such as visual-semantic alignment, multiple-stage semantic processing and transformer-based visual-semantic interaction. Instead of using argmax, (Bhunia et al. 2021) use Gumbel-softmax (Jang, Gu, and Poole 2016) for extracting semantic information, which is then fed into succeeding transformer-based visual-semantic reasoning module. The decoding involves complex multiple-stage attentional LSTM that couples with feature pyramid networks. Our approach is not only conceptually much simpler and computationally more efficient, but also more effective in solving STR.

We use resnet-like architecture in extracting local features. The input images are first resized to have the same height, with aspect ratio kept unchanged. During training and testing, replication-padding (padding using values from the image border) is used for batch processing.

Assuming that input image is of size $W\times H\times 3$, the convnet module will produce feature maps of size $w\times h\times d$. The features will be fed into the visual module.

Convnets are good at learning local features but hard to learn global correlation between features that locate far apart. Inspired from the success of transformers when applied to vision task, we insert a transformer module here in order to enhance the feature maps by allowing them to interact with each other. The visual modules takes the convolution feature maps as input, producing primary visual features. The primary visual features together encode completely the appearance information of the input image. They are then converted to primary semantic features by the vs-align module.

The visual module consists of MHSA, layer-norm, and feed-forward network, as described in (Vaswani et al. 2017), but with minor modification that puts layer-norm before MHSA (Dosovitskiy et al. 2020; Wang et al. 2019a). Note that Fig. 1 does NOT reflect this modification. The feature maps are viewed/reshaped as a sequence of dimension $d$ and length $n=w\times h$. the module V, takes the feature sequence as input and generate a sequence of the same dimension and length. The interaction between visual features themselves mainly takes place in the MHSA layers.

Parallel attentions has been used to map visual features into semantics(Wang et al. 2021b; Lyu et al. 2019). Compared with previous work, the vs-align module in this work is much simpler but still effective. Note that we do not use the name semantic decoder here because the output of the module is not actually the character index, but rather the probability distribution of the characters. In other words, the semantic feature encodes the probability distribution at each candidate location, which can be readily converted into character index with $argmax$ operation in decoding stage (not in training). This allows the gradients to flow smoothly inside the VST.

The architecture of vs-align is depicted in Fig. 2. As shown in the figure, the visual features (viewed as feature maps) are projected using a linear layer and normalize using softmax operator, obtaining $t$ attention maps each of which has the same spatial dimension as the origin feature map. The alignment module can also be formulated mathematically as follows,

$$S=\text{softmax}(QV^{T})V$$

(1)

where $S\in R^{t\times d}$ is the semantic sequence, $V\in R^{n\times d}$ is the visual sequence and $Q\in R^{t\times d}$ is the trainable projection matrix.

Our VS-align module is able to learn the relation between the semantic and visual features. This is different from CTC (Graves et al. 2006), which is good at aligning two sequences with beam search, but not capable of handling information along the height direction.

As shown in Fig. 1, there are two vs-align modules, which share weights during training and inference. This weight-sharing scheme is the key to successfully learning 1st semantic information. It enables the weights learned in aligning 2nd visual features to transfer to aligning 1st visual features, making the early learning of 1st semantics possible. The 1st semantic in return enhances the 1st visual features through interaction module, further improving the training of the 2nd vs-align module. The early learning of semantics provides more chances to correct the semantic features through interaction in later stages. If semantic had been learned in the last stage, we would have not change to correct it, unless incorporating additional lexicon or dictionary information (Nguyen et al. 2021)

The interaction module plays the key role of mixing primary semantic features and visual features. The modules takes two streams , namely the primary semantic features and primary visual features as input, producing the secondary visual and semantic features. Fixed positional encoding are added into the visual stream. For the semantic stream, learnable position embedding is used. Other positional encoding schemes probably help as long as it breaks permutation invariant property of transformers.

The design of the interaction module follows from module V, comprising layers of transformer blocks. The module enables the feature interaction between a) the visual stream and semantic will interact with each other, and b) the features themselves inside the same stream.

The interpretation of module can be seen from eq. 2. Following from the notations in (Vaswani et al. 2017), the interaction in the MHSA layer is obvious by judging form the following equation,

$$S=\text{softmax}(\frac{[Q_{s};Q_{v}][K_{s};K_{v}]^{T}}{\sqrt{d_{k}}})[V_{s};V_{v}],$$

(2)

where the subscripts s and v denote semantic and visual part respectively, ’;’ is column representation.

In this way, the primary semantic features will look for support from the spatial visual features. The interaction is useful for enhancing semantic features because they now have attentive access to all spatial location of the visual features. On the other hand, visual features are also enhanced not only because then can interact with all other spatial locations regardless of distance, but they can also learn from semantic features, which we hypothesize is useful for dealing with appearance degradation.

The semantic features can be seen as pseudo-linguistic features, distinguished from visual domain. Hence we add domain embedding before feeding two combined streams into the interaction module.

One of the outputs of interaction module is secondary visual sequence in $R^{n\times d}$, which will be converted to tertiary semantic features using vs-align module following by cross-entropy loss with softmax, against the ground-truth text. The other is semantic features in $R^{n\times d}$, which will be directly compared against ground-truth text with the same activation ans loss type. Note that at inference time, the two loss branches are unneeded.

The semantic module is inserted optionally to play the role of further fusing the two semantic streams. When inserted, our model is named VST-F (full); otherwise, it is named VST-B (basic). The two streams are the secondary semantic features generated from the intersection module and the tertiary (third) semantic features generated from the second vs-align module respectively.

The layer configureation of this module also follows from the design of module V and I. Fixed positional encoding is added, but there is no domain or segment embedding because the two streams are both semantic. Since the transformer blocks keep the input shape, we concatenate the two streams along channel direction to $t\times 2d$ and feed it to linear and softmax layer, obtaining the final text predictions.

.

As shown in Fig. 1, there are three losses indicated by the ladders in different colors. The green one measures the discrepancy between the the true text labels and semantic features output from second vs-align module using softmax activation and cross-entropy loss. The blue one measures how different are the secondary semantic features from the true text. The yellow one is the main loss for decoding final text for VST-F, while the other two are auxiliary for improving convergence performance. For VST-B, since there is no module S, only blue and green loss are used. For all loss branches, the discrepancy are measured by cross-entropy loss with softmax activation.

From the optimization perspective, the VST-F solves the following optimization problem,

	$\displaystyle\min_{\theta_{v},\theta_{a},\theta_{i},\theta_{s}}$	$\displaystyle L(s_{2},y)+L(s_{3},y)+L(f_{\theta_{s}}(s_{2},s_{3}),y)$		(3)
	subject to	$\displaystyle\quad v_{1}=f_{\theta_{v}}(x),s_{1}=f_{\theta_{a}}(v_{1})$
		$\displaystyle\quad s_{2},v_{2}=f_{\theta_{i}}(s_{1},v_{1}),s_{3}=f_{\theta_{a}}(v_{2})$

where $\theta_{v},\theta_{a},\theta_{i},\theta_{s}$ denote the parameters of module C+V, A, I, and S respectively, $L$ is cross-entropy loss, $s_{1},s_{2},s_{3},v_{1},v_{2}$ are the primary/secondary/tertiary semantic/visual features resp., $x,y$ are the input mage and text label resp. For VST-B, the third term and the corresponding constraint in the above objective function is discarded.

Our model is trained on three datasets: SynthText (ST) (Gupta, Vedaldi, and Zisserman 2016), MJSynth (MJ) (Jaderberg et al. 2014, 2016) and SynthAdd (SA) (Li et al. 2019). The training dataset contains totally 15.7 million synthetic images. For evaluation, we use four regular text datasets: IIIT 5K-words (IIIT) (Mishra, Alahari, and Jawahar 2012) , Street View Text (SVT) (Wang, Babenko, and Belongie 2011), ICDAR 2003 (IC03) (Lucas et al. 2003), ICDAR 2013 (IC13) (Karatzas et al. 2013) , and three irregular text datasets: ICDAR 2015 (IC15) (Karatzas et al. 2015), Street View Text Perspective (SVTP) (Phan et al. 2013) and CUTE (Risnumawan et al. 2014) .

IIIT contains $3000$ cropped images collected from Google image search. SVT is collected from Google Street View containing $647$ testing images. Following from (Wang, Babenko, and Belongie 2011), we select $867$ cropped images from IC03 for testing. IC13 contains $1095$ testing images and we discard images containing non-alphanumeric characters or contains fewer than three characters, resulting in $1015$ images for testing, following from previous work (Yu et al. 2020).

For irregular datasets, IC15 contains $2077$ cropped images by Google Glasses. We use $1811$ images after discarding extremely distorted images. SVTP and CUTE contains $639$ and $288$ images respectively.

Module C is implemented as a four-layers resnet, with each layer having 1,2,5 and 3 blocks respectively. Module V, I and S each consists of 3 transformer layers. The number of parameters are about 63.99M for VST-F, and 63.65M for VST-B.

The image patches are resized with aspect-ratio unchanged so that $H=48$, and the width is trimed or padded to $W=160$ pixels. The feature map size is $6\times 40\times 512$, $w=40,h=6,n=240$ The number of class is set to 38, including 0-9,a-z,[unk],[eos]. We assume that the maximum character length is 25, i.e., $t=25$. We use online data augmentations including distortion, stretch, perspective transform, blurring, colour jitter etc. In training, we set the batch size to 256 and sample from ST, MJ, SA datasets with weight 0.4,0.4 and 0.2 to balance the datasets (Baek et al. 2019). We use Adam optimizer with the initial learning rate 1e-4 and decreased lr to 1-e5 when the loss plateaus. We use for and the training takes roughly 3 days to converge on four Tesla V100 GPUs.

We compare VST-B and VST-F with previous work over seven public benchmarks in table 1. The proposed VST-B and VST-F both achieve or on par with state-of-the-art approaches. The large version of the very recent work by (Fang et al. 2021) performs slightly better than ours on three of the seven datasets, but they used extra curbersome pretrained language models. Moreover, our VST-B and VST-F are about three times faster than theirs in inference. On the rest four datasets, we improve SOTA by large margins (0.2% IIIT, 0.3% SVT, 0.8% IC03, 2.8% CUTE).

To verify whether each module is critical for the final performance, ablation study is conducted by subsequently adding one module starting from the most basic C+V configuration. The results are shown in table 2, from which we observe a consistent performance gain when a new module is added. Also note that C+V+A also performs better than C+V, which verifies the effectiveness of vs-align module. For the C+V+A+I configuration, there are three ways of decoding: decoding from secondary semantics ($s_{2}$), tertiary semantics ($s_{3}$) or by probability voting from both semantics.

Inside the module I, each semantic vector will attend to visual features corresponding to some spatial locations. We can visualize which locations are attended to using heapmaps. Fig. 3 shows three examples. The heatmaps are computed by averaging all 8 attention heads and overlapping onto images after upsampling. Note that inside the module each head focuses on different aspects of the images, but on average the attention is mainly at the right spatial location for each decoded character. For some characters, the focus is slightly shifted, probably due to the translation-invariant property of convnet, which means that the visual features cannot be mapped back uniformly to the corresponding spatial locations. Overall, the visualization proves that the interaction module works as expected.

Likewise, we also visualize the attention heatmaps of the primary and secondary vs-align modules. Even though the two modules share the same weights, their attention maps should be different since the visual and semantic features are different. We expect that the attention maps of 2nd vs-align be more precise than the first one. This is verified in Fig. 4.

Beside visualization, we select some successful and failure cases for illustration purpose, shown in Fig. 5 and Fig. 6 respectively.