SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yingjie Liu $^{1}$  Xuan Liu $^{2}$  Hui Yu $^{2}$  Xuan Tang $^{1}$  Xian Wei $^{1*}$
$^{1}$ East China Normal University  $^{2}$ FJIRSM, Chinese Academy of Sciences
Recently, the attention mechanism has shown outstanding competence in capturing global structure information and long-range relationships within data, thus enhancing the performance of deep vision models on various computer vision tasks. In this work, we propose a novel alternative dictionary learning-based attention (Dic-Attn) module, which models this issue as a decomposition and reconstruction problem with the sparsity prior, inspired by sparse coding in the human visual perception system. The proposed  $\text{Dic-Attn}$  module decomposes the input into a dictionary and corresponding sparse representations, allowing for the disentanglement of underlying nonlinear structural information in visual data and the reconstruction of an attention embedding. By applying transformation operations in the spatial and channel domains, the module dynamically selects the dictionary's atoms and sparse representations. Finally, the updated dictionary and sparse representations capture the global contextual information and reconstruct the attention maps. The proposed  $\text{Dic-Attn}$  module is designed with plug-and-play compatibility and facilitates integration into deep attention encoders. Our approach offers an intuitive and elegant means to exploit the discriminative information from data, promoting visual attention construction. Extensive experimental results on various computer vision tasks, e.g., image and point cloud classification, validate that our method achieves promising performance, and shows a strong competitive comparison with state-of-the-art attention methods.
Visual perception plays a critical role in obtaining external information. However, processing massive visual information in the real world is challenging for both human brains and computers. Modern psychology and cognitive neuroscience find that the visual attention mechanism is one of the essential keys to a human's good cognitive ability [43], enabling selective and sparse activation of neurons in response to input. In the field of computer vision (CV), various attention methods have also been developed, with successful applications in lots of Convolutional Neural Networks and Transformers, revealing great potential in visual tasks [4, 32, 5, 48, 46, 15]. CV attention methods utilize various calculation functions to generate attention maps, but they can be seen as a general form (refer to Sec. 3.1) [11, 41, 10]. Researchers have recognized several advantages of the attention mechanism, e.g., capturing global long-term information and extracting hierarchical data structures [6]. Moreover, recent research has also highlighted the potential of leveraging sparsity to develop attention algorithms [31, 62, 23]. However, there are still pending questions regarding the efficiency and effectiveness of current attention methods when dealing with regular image features and non-Euclidean point cloud data [57, 35, 9]. Considering the significant memory costs and time complexity [42] associated with current attention methods, it is crucial to explore possibilities of optimizing their internal operations or re-evaluating the current form, especially in the view of sparsity.
*Corresponding author: xwei@sei.ecnu.edu.cn
In this context, recalling the previous works and considering the original motivation behind introducing attention methods remind us of the dictionary learning/sparse coding algorithm, which holds importance but is often overlooked in present discussions. The dictionary learning algorithm decomposes visual data into a dictionary basis and their corresponding sparse representations/codes. Under different conditions, the sparsity of sparse code can be adapted: code with higher sparsity (more zero elements) mostly corresponds to features of smaller scale, higher roughness, or higher hierarchy. Furthermore, this process echoes the nonlinear effects in complex cells and the sparse coding strategies for sensory information in the primary visual cortex and advanced visual region [34, 1, 19, 56]. It has been observed that the structure in a sparse domain could make the hidden patterns more prominent and easier to be captured [29, 52, 55, 60]. All these motivate us to develop an alternative novel attention form, i.e. Dictionary learning-based Attention (Dic-Attn), that offers a visual focus for deep neural networks (DNNs). The  $\text{Dic-Attn}$  module can be conveniently used to replace existing attention modules, and compatibly explain their working principles.
The  $\text{Dic-Attn}$  module employs two selection matrices to exploit channel and spatial attention from dictionary atoms. By selecting the appropriate basis and re-locating the coordinates, we can uncover hidden visual attention information by reconstructing the attention map using the updated basis (dictionary) and coordinates (corresponding sparse representations). Note that, the sparse representations in our proposed module originate from the non-linear decomposition of visual data and are subsequently reconstructed into the attention matrix. The sparsity property of the attention matrix is induced naturally by the sparse prior [52, 50, 60] and is related to the visual attention area within the images. The main contributions of this paper are summarized as follows:
1. We propose a novel visual attention module, namely, the Dictionary learning-based Attention (Dic-Attn) module, which is based on dictionary learning and provides better interpretability of providing contextual information.
2. We demonstrate that the proposed module has great potential in disentangling non-linear structural information by employing a masked dictionary and transformed sparse representations. We construct a deep attention stage composed of the proposed modules. Each module at different depths, ranging from shallow to deep, dynamically updates dictionaries and sparse representations through spatial and channel transformations, step-wise-by reconstructing the final attention map.
3. We apply the proposed module to various visual models since it is designed to be plug-and-play, achieving better performance than various existing attention modules on image classification, semantic segmentation, and point cloud classification, with fewer parameters and faster computation.
Figure 1: The pipeline of our proposed  $\text{Dic-Attn}$  block. First, input features  $\hat{\mathbf{X}}$  into the  $\text{Dic-Attn}$  block. Within the  $\text{Dic-Attn}$  block, the inner  $\text{Dic-Attn}$  module decomposes the normalized input  $\mathbf{X}$  into a dictionary and corresponding sparse representations, which will be transformed in the spatial and channel domains. The  $\text{Dic-Attn}$  module can dynamically adjust the weights to promote the visual focus on influential elements. This process is both data and task-driven, thereby enhancing the effectiveness of the attention mechanism. Finally, the output of the  $\text{Dic-Attn}$  block is obtained by summing the input with the Attention Output. The  $\text{Dic-Attn}$  block can be stacked layer by layer to form a deep attention model, enabling capturing and integrating visual attention hierarchically across multiple levels.
Attention Mechanism in Computer Vision Tasks. CV Attention methods simulate the phenomenon of the human visual system, where selective attention helps neural system analyze and understand complex scenes more efficiently. Recurrent Model [32] is a pioneering work that develops the attention mechanism in DNNs. Later, various attention methods are proposed. Some of them try to improve the pre-processing method of the input images [8, 54], while others modify the way of obtaining the attention matrix, resulting in a series of variants, e.g., spatial transformer network (STN) [16], self-attention (SA) [46], squeeze-and-excitation network (SENet) [14], dual attention (DA) module [5], double attention  $(\mathrm{A}^2)$  module [4], hamburger (HAM) module [6], etc. [21, 10, 48, 25, 39]. The SA module is proposed in [46] and first introduced to computer vision in [49]. It achieved great success in various vision tasks and made the Vision Transformer (ViT) successful in academia. [6] constructs the long-term global context learning as a low-rank recovery problem and proposes the matrix decomposition (MD) module, focusing only on the low-rank property prior assumption of feature representations. It may have limitations in fully capturing the complexity and richness of global structure information in the data. Most recent related work [39] was inspired by the two prominent features of the human visual attention mechanism, i.e., recurrence and sparsity. The so-called concurrency attribute in [39] naturally comes from the ordinary differential equation description of the sparse coding update process, which is actually inside the attention module and the sparse representations require a multi-step update. It also lacks the analysis of attention modules with different depths. Moreover, the residual connections between multi-layer attention modules still remain under explanation. [23] measures sparsity as the percentage of nonzero entries in the activation maps of each layer. It finds that sparsity emerges in the parameters of all layers of the Transformer. However, it does not explicitly state whether the output of the attention module exhibits sparsity or not. Some studies [37, 51, 17, 33] have tried to post-hoc-interpret these developed attention methods, and utilize white-box approaches to do the analysis [13, 41, 58, 31].
Dictionary Learning in Computer Vision Tasks. Sparsity is an important and noticeable prior in the field of computer vision, e.g., the statistically redundant property of natural images is served for dictionary learning. Hence, image sample  $\mathbf{x}$  is reconstructed with the dictionary  $\mathbf{D}$  and its corresponding sparse code  $\phi$ , i.e., samples can be decomposed into a dictionary and corresponding sparse representations. The dictionary  $\mathbf{D}$  captures representative and critical features from the inputs and reveals the interference of insignificant information. Dictionary learning has been extensively studied and verified to be an outstanding approach for exploiting the underlying sparse structures hidden in images [29, 52, 50, 55, 60], efficiently transferring and expressing information from high-dimensional image data into sparse space. It plays an important role in various computer vision tasks, e.g., image reconstruction [1, 29, 3, 38, 64], image classification[52, 47, 55, 24, 60], visual tracking [61] and others.[39]. Note that, compared with non-negative matrix factorization and vector quantization applied in MD module [6] mentioned above, the dictionary learning algorithm performs better in the case of natural images [26]. Technical details about how we learn the dictionary is shown in Sec. 3.2.
In this section, we present the proposed dictionary learning-based visual attention (Dic-Attn) method. First, we give a brief review and summary of the CV attention form and dictionary learning. Then, we describe the main idea of the  $D_{ic}$ -Attn module to obtain visual attention. Finally, we analyze the core concept behind the  $D_{ic}$ -Attn module. The overall structure of the proposed  $D_{ic}$ -Attn method is illustrated in Fig. 1.
3.1 Attention Form
Generally speaking, we can summarize a CV attention form [11, 41].
Input:X,
Attention Module:  $f(\mathbf{X})$ .
Attention Block:  $g(f(\mathbf{X}),\mathbf{X})$
The input features  $\mathbf{X}$  are firstly pre-processed. Subsequently, the attention module obtains an attention map of the input. Finally,  $g(\cdot)$  is introduced to further connect attention map and input features, e.g., the normalization operation, drop operation, skip-connection, etc.. All these constitute an Attention Block. Among them, the function  $f(\mathbf{X}):\mathbf{X}\mapsto$  Attention Output in the attention module is the most critical mapping. Attention blocks can be added to the middle of the network, or form a feed-forward network. For example, in the ViT encoder with depth  $N$ , there will be  $N$  attention blocks. Take ViT as an example, suppose there is a feature tensor  $\mathbf{Z}$  with shape  $b\times n\times h\times w$ , where  $n$  denotes the number of channels. The feature tensor is vectorized in the spatial domain and the last two dimensions are transposed. Hence, the attention block takes a sequence of image features as input, i.e.,  $\mathbf{X}\in \mathbb{R}^{b\times s\times n}$ , where  $s = h\times w$ . The Self-Attention module is introduced and subsequently outputs a spatial attention map, essentially operating in the spatial domain. The formulation of the standard Self-Attention module in ViT [46] is as follows:
$$
[ \mathbf {Q}, \mathbf {K}, \mathbf {V} ] = [ \mathbf {W} _ {\mathbf {Q}} \mathbf {X}, \mathbf {W} _ {\mathbf {K}} \mathbf {X}, \mathbf {W} _ {\mathbf {V}} \mathbf {X} ],
$$
$$
\mathbf {A} = \operatorname {s o f t m a x} \left(\mathbf {W} _ {\text {s c a l e}} \left(\mathbf {Q} \mathbf {K} ^ {T}\right)\right), \tag {1}
$$
Attention Output  $= \mathrm{AV}$
where  $\mathbf{X} \in \mathbb{R}^{b \times s \times n}$  is the input,  $\mathbf{Q} \in \mathbb{R}^{b \times s \times n}$ ,  $\mathbf{K} \in \mathbb{R}^{b \times s \times n}$  and  $\mathbf{V} \in \mathbb{R}^{b \times s \times n}$  denote the query matrix, the key matrix, and the value matrix, respectively. They are obtained by full-connected layers.  $\mathbf{W}_{\mathbf{Q}}, \mathbf{W}_{\mathbf{K}}$  and  $\mathbf{W}_{\mathbf{V}}$  are all square matrices with the shape  $n \times n$ . A represents the attention matrix. Finally, through the residual structure, i.e., skip-connection, attention blocks have the final output result, i.e., Attention Output + X.
The implicit variables or the spatial attention map results of these attention modules can be visualized, e.g., traditional SA blocks in ViT [46], and External-Attention (EA) blocks [10]. These works indicate that the attention module focuses on the discriminative regions of each image. Recalling the CV attention form we summarized at the beginning of this section, we can find that  $f(\cdot)$  is the most important function when obtaining visual attention in DNNs. In this paper,  $f(\cdot)$  contains the process of dictionary learning, the transformation of both learned dictionary & corresponding sparse representations, as well as reconstruction, which is to be introduced in Sec. 3.2.
3.2 Dictionary Learning
Formally, given data  $\mathcal{X} \coloneqq [x_1, \dots, x_n] \in \mathbb{R}^{n \times s}$ , the standard dictionary learning algorithm is to find a dictionary  $\mathbf{D}$  and the corresponding sparse code  $\Phi$ . The dictionary  $\mathbf{D} = [\mathbf{d}_1, \dots, \mathbf{d}_k]$  has a shape of  $n \times k$ , the sparse code  $\Phi = [\phi_1, \dots, \phi_k] \in \mathbb{R}^{k \times s}$  selects atoms from the dictionary to reconstruct the input  $\mathcal{X}$ . The objective function of dictionary learning can be formulated as follows:
$$
\underset {\mathbf {D}, \boldsymbol {\phi} _ {i}} {\operatorname {a r g m i n}} h (\mathbf {D}) + \sum_ {i = 1} ^ {s} \frac {1}{2} \| x _ {i} - \mathbf {D} \boldsymbol {\phi} _ {i} \| _ {2} ^ {2} + g \left(\boldsymbol {\phi} _ {i}\right), \tag {2}
$$
$$
\begin{array}{l} \text {s . t .} \mathbf {D} \in \mathcal {S} (n, k), \phi_ {i} \in \mathbb {R} ^ {k}, \end{array}
$$
where  $h(\mathbf{D}) = \alpha \sum_{i \neq j} \left\| \mathbf{d}_i^T \mathbf{d}_j \right\|_F^2$  denotes the regularization term of the dictionary. It leads to mutual incoherence between pair atoms in the dictionary. Besides,  $\mathbf{D} \in S(n,k) := \{\mathbf{D} \in \mathbb{R}_*^{n \times k} : \text{diag}(\mathbf{D}^T \mathbf{D}) = \mathbf{I}_k\}$ . Namely, the set of dictionaries is a product manifold of  $s$  times the  $(k-1)$ -dimensional unit sphere, i.e.,  $S(n,k)$  restrict all atoms  $\mathbf{d}_i \in \mathbb{R}^k$  to have unit norm.  $g(\phi_i)$  usually utilizes norm constraint to control the sparsity of the sparse code  $\phi_i$ . The  $l_1$  norm regularization is a popular way to replace  $l_0$ -norm by  $l_1$ -norm convex relaxation, known as the Lasso [44] regression. It has numerous variants, e.g., elastic net. Elastic net inherits the stability of Ridge and the sparsity brought by Lasso, for combining both regularization terms of  $l_1$  norm and  $l_2$  norm. Hence in this paper, we have  $g(\phi_i) = \beta \| \phi_i \|_1 + \frac{\lambda}{2} \| \phi_i \|_2^2$ .  $\beta$  and  $\lambda$  are regularization weights.
Eq. (2) can be regarded as an iterative optimization problem, where dictionary  $\mathbf{D}$  and sparse representation alternately updates. Each sparse code is an implicit function about the temporarily fixed dictionary  $\mathbf{D}$ . With Eq. (2), the incoherent assumption of dictionary atoms and the sparse prior  $g(\phi_i)$  [50], the close solution of sparse representation can be expressed as follows:
$$
\phi_ {i} ^ {*} = \left(\mathbf {D} _ {\Lambda} ^ {T} \mathbf {D} _ {\Lambda} + \lambda \mathbf {I}\right) ^ {- 1} \left(\mathbf {D} _ {\Lambda} ^ {T} x _ {i} + \beta \mathbf {v} _ {\Lambda}\right), \tag {3}
$$
where  $\Lambda := \{i \in \{1, \dots, s\} | \phi_{ij}^* \neq 0\}$  denotes the set of indexes of the non-zero entries of the solution  $\phi_i^* = [\phi_{i1}^*, \dots, \phi_{is}^*]$ .  $\mathbf{D}_{\Lambda}$  is the subset of the dictionary in which the indexes of atoms fall into  $\Lambda$ .  $|\Lambda|$  denotes the sparsity of  $\phi_i^*$ ,  $\mathbf{v}_{\Lambda} \in \{\pm 1\}^{|\Lambda|}$  carries the signs of  $\phi_i^*$ . Technically, it is easy to realize the calculation of the above polynomial formula and obtain the sparse representations.
Constructing or learning a dictionary is also crucial for achieving the goal of highlighting significant objects while suppressing distracting or noisy elements. In the early stage, the dictionary is constructed by Cosine functions, Fourier, Wavelets, Contourlets, Gabor, a set of complete bases, etc. These approaches were widely employed in signal processing due to their mathematical simplicity. However, such a fixed dictionary is manually designed under some mathematical constraints and is not flexible enough to represent complex natural image structures. Recently, researchers have turned to directly initializing and learning the dictionary from image data, which can also improve the sparsity of  $\Phi$  and the performance of downstream tasks [52, 60, 64]. The dictionary contains prototype underlying features of vision data, based on which each sample can find its unique optimal sparse solution.
3.3 The Proposed  $\text{Dic-Attn}$  for Visual Attention
3.3.1 Dic-Attn Method
The  $\text{Dic-Attn}$  module first decomposes the input into a dictionary and sparse representation. Then its primary objective is to provide visual attention to the model by reconstructing the attention map with residuals. Its process is as follows,
$$
\mathbf {X} = \mathbf {D} \boldsymbol {\Phi},
$$
$$
\mathbf {V} = \mathbf {D} \mathbf {W} _ {\mathbf {D}},
$$
$$
\mathbf {A} = \operatorname {s o f t m a x} \left(\mathbf {W} _ {\text {s c a l e}} \left(\mathbf {W} _ {\boldsymbol {\Phi}} \boldsymbol {\Phi}\right)\right), \tag {4}
$$
$$
\text {A t t e n t i o n} = \mathrm {V A},
$$
where  $\mathbf{X} = [\mathbf{X}_1^T, \mathbf{X}_2^T, \dots, \mathbf{X}_b^T] \in \mathbb{R}^{b \times s \times n}$ . Each atom  $\mathbf{d}_i \in \mathbb{R}^n$  in the dictionary  $\mathbf{D} \in \mathbb{R}^{n \times k}$  will be re-weighted by  $\mathbf{W}_{\mathbf{D}} \in \mathbb{R}^k$ . The elements of  $\mathbf{W}_{\mathbf{D}} \in \mathbb{R}^k$  are trainable mask weights for dictionary atoms.  $\mathbf{W}_{\Phi} \in \mathbb{R}^{k \times k}$  denotes the selection mask matrix for sparse representations  $\Phi = [\Phi_1, \dots, \Phi_b] \in \mathbb{R}^{b \times k \times s}$ .  $\mathbf{W}_{scale}$  denotes the scaling matrix. Note that  $k = \alpha n$ ,  $\alpha$  is the hyper-parameter.
Firstly, the dictionary is learned on the entire training-dataset, aiming to be adaptive to all inputs  $\mathbf{X}$ . It contains diverse and important features (in atoms) of the input, so it can solve the sparse coding problem and obtain the corresponding code  $\Phi$ . We have:
$$
\mathbf {D} = \left[ \begin{array}{c c c c} \mathbf {d} _ {1} ^ {c} & \mathbf {d} _ {2} ^ {c} & \dots & \mathbf {d} _ {k} ^ {c} \end{array} \right] = \left[ \begin{array}{c c c c} \mathbf {d} _ {1} ^ {r} & \mathbf {d} _ {2} ^ {r} & \dots & \mathbf {d} _ {n} ^ {r} \end{array} \right] ^ {T},
$$
$$
\boldsymbol {\Phi} _ {m} = \left[ \begin{array}{l l l l} \phi_ {1} & \phi_ {2} & \dots & \phi_ {s} \end{array} \right] = \left[ \begin{array}{l l l l} \psi_ {1} & \psi_ {2} & \dots & \psi_ {k} \end{array} \right] ^ {T},
$$
$$
\mathbf {X} _ {m} = \left[ \begin{array}{c c c c} \mathbf {x} _ {1 1} & \mathbf {x} _ {1 2} & \dots & \mathbf {x} _ {1 n} \\ \mathbf {x} _ {2 1} & \mathbf {x} _ {2 2} & \dots & \mathbf {x} _ {2 n} \\ \vdots & \vdots & \vdots & \vdots \\ \mathbf {x} _ {s 1} & \mathbf {x} _ {s 2} & \dots & \mathbf {x} _ {s n} \end{array} \right] = \left[ \begin{array}{c c c c} \mathbf {d} _ {1} ^ {r} \phi_ {1} & \mathbf {d} _ {1} ^ {r} \phi_ {2} & \dots & \mathbf {d} _ {1} ^ {r} \phi_ {n} \\ \mathbf {d} _ {2} ^ {r} \phi_ {1} & \mathbf {d} _ {2} ^ {r} \phi_ {2} & \dots & \mathbf {d} _ {2} ^ {r} \phi_ {n} \\ \vdots & \vdots & \vdots & \vdots \\ \mathbf {d} _ {s} ^ {r} \phi_ {1} & \mathbf {d} _ {s} ^ {r} \phi_ {2} & \dots & \mathbf {d} _ {s} ^ {r} \phi_ {n}, \end{array} \right] \tag {5}
$$
$$
\mathbf {x} _ {i j} = \mathbf {d} _ {i} ^ {r} \boldsymbol {\phi} _ {j} = \sum_ {t = 0} ^ {k} \left(\mathbf {d} _ {t} ^ {c} \boldsymbol {\psi} _ {t}\right) _ {i j},
$$
where  $\mathbf{X}_m$  denotes the  $m$ -th batch of  $\mathbf{X}$ ,  $\Phi_m = [\phi_1, \dots, \phi_n] = [\psi_1, \dots, \psi_k]^T \in \mathbb{R}^{k \times s}$ ,  $\mathbf{d}_t^c$  and  $\mathbf{d}_i^r$  denotes the  $t$ -th column and the  $i$ -th row of  $\mathbf{D}$ . The optimization process of dictionary learning maximizes the descriptiveness of the atoms  $\mathbf{d}_t^c$  while minimizing the redundancy between original inputs  $\mathbf{X}$  and  $\mathbf{D}\Phi$ . Then, transformations  $\mathbf{W}_{\mathbf{D}}$  and  $\mathbf{W}_{\Phi}$  for dictionaries and sparse representations are introduced, respectively, which facilitate the integration of spatial and channel attention into the model by assigning weights in both the spatial and channel domains.
The diagonal transformed matrix  $\mathbf{W}_{\mathbf{D}}$  and transformed matrix  $\mathbf{W}_{\Phi}$  are specifically designed to re-weight the dictionary and transform its corresponding sparse representations, respectively. Through
Algorithm 1: Dictionary Learning-Based Attention Module
Input:  $\mathbf{X}\in \mathbb{R}^{b\times s\times n}$  , depth  $N$    
Parameter:  $\mathbf{D}\in S(n,k),\mathbf{W}_{\Phi}\in \mathbb{R}^{k\times k},\mathbf{W}_{\mathbf{D}}\in \mathbb{R}^{k},\mathbf{W}_{scale}$    
Latent Variables:  $\Phi \in \mathbb{R}^{k\times s}$ $f(\Phi^{\mathrm{T}})$  V   
Output: Attention Output   
# Initialize  $\mathbf{D}\in S(n,k)$  by randomly sampling from  $\mathcal{N}(0,1)$  and atom-wise normalization; # Initialize  $\mathbf{W}_{\Phi},\mathbf{W}_{\mathbf{D}},\mathbf{W}_{scale}$  with reference to kaiming-init method [12];   
# forward   
## Update  $\Phi$  by Eq. (3);   
#Obtain Attention Output by Eq. (4);   
# backward   
## Update D,  $\mathbf{W}_{\Phi},\mathbf{W}_{\mathbf{D}},\mathbf{W}_{scale}$  by minimizing loss function of the network ;