SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Mingjie Li<sup>1</sup>, Yisen Wang<sup>1,2</sup>, Zhouchen Lin<sup>1,2,3*</sup>
$^{1}$  National Key Lab of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University
$^{2}$  Institute for Artificial Intelligence, Peking University
3 Peng Cheng Laboratory
lmjat0111@outlook.com, {yisen.wang, zlin} @pku.edu.cn
Despite the connection established by optimization-induced deep equilibrium models (OptEqs) between their output and the underlying hidden optimization problems, the performance of it along with its related works is still not good enough especially when compared to deep networks. One key factor responsible for this performance limitation is the use of linear kernels to extract features in these models. To address this issue, we propose a novel approach by replacing its linear kernel with a new function that can readily capture nonlinear feature dependencies in the input data. Drawing inspiration from classical machine learning algorithms, we introduce Gaussian kernels as the alternative function and then propose our new equilibrium model, which we refer to as GEQ. By leveraging Gaussian kernels, GEQ can effectively extract the nonlinear information embedded within the input features, surpassing the performance of the original OptEqs. Moreover, GEQ can be perceived as a weight-tied neural network with infinite width and depth. GEQ also enjoys better theoretical properties and improved overall performance. Additionally, our GEQ exhibits enhanced stability when confronted with various samples. We further substantiate the effectiveness and stability of GEQ through a series of comprehensive experiments.
Deep Neural Networks (DNNs) show impressive performance in many real-world tasks on various data like graphs [43], images [41, 20], sequences [8], and others. However, most neural network structures are constructed by experience or searching on the surrogate datasets [31, 57]. Therefore, these architectures cannot be interpretable and such a phenomenon hinders further development. Apart from the current neural network models, traditional machine learning methods like dictionary learning [46, 33], subspace clustering [54] and other methods [29, 56, 55, 30] can design their whole procedure by designing optimization problems with specific regularizers customized from their mathematical modeling and requirements. Thus, these models are easily interpreted. However, the traditional machine learning algorithms' whole procedures do not consider the hidden properties of features and labels. Therefore, they usually perform worse on tasks with more data.
To link two types of models, OptEqs [52] tries to recover the model's hidden optimization problem to make their model "mathematically explainable". They claim that the output features  $\tilde{\mathbf{z}}^*$  (we also called the equilibrium state) with respect to input  $\mathbf{x}$ , which are obtained by solving the fixed-point
equation in Eqn (1), is the optimal solution for its hidden optimization problem defined in Eqn (2).
$$
\tilde {\mathbf {z}} ^ {*} = \tilde {\mathbf {W}} ^ {\top} \sigma \left(\tilde {\mathbf {W}} \tilde {\mathbf {z}} ^ {*} + \mathbf {U} \mathbf {x} + \mathbf {b}\right), \hat {\mathbf {y}} = \tilde {\mathbf {W}} _ {c} \tilde {\mathbf {z}} ^ {*} + \mathbf {b} _ {c}, \tag {1}
$$
$$
\underset {\tilde {\mathbf {z}}} {\min} G (\tilde {\mathbf {z}}; \mathbf {x}) = \underset {\tilde {\mathbf {z}}} {\min} \left[ \mathbf {1} ^ {\top} f (\tilde {\mathbf {W}} ^ {- 1 ^ {\top}} \tilde {\mathbf {z}}) - \left\langle \mathbf {U x} + \mathbf {b}, \tilde {\mathbf {W}} ^ {- 1 ^ {\top}} \tilde {\mathbf {z}} \right\rangle + \frac {1}{2} \| \tilde {\mathbf {W}} ^ {- 1 ^ {\top}} \tilde {\mathbf {z}} \| _ {2} ^ {2} - \frac {1}{2} \| \tilde {\mathbf {z}} \| _ {2} ^ {2} \right], (2)
$$
where  $\sigma$  is the ReLU activation function,  $\mathbf{U}$ ,  $\tilde{\mathbf{W}}$ ,  $\mathbf{b}$ ,  $\tilde{\mathbf{W}}_c$ ,  $\mathbf{b}_c$  are learnable parameters. They are trained by optimizing loss functions, such as cross-entropy loss, which are calculated based on the final prediction  $\hat{\mathbf{y}}$  derived from  $\tilde{\mathbf{z}}^*$  as shown in Eqn (1).  $\tilde{\mathbf{W}}^{-1}$  represents an invertible or pseudo-invertible matrix for  $\tilde{\mathbf{W}}$ , and the function  $f$  is a positive indicator function, which outputs 1 when  $x \geq 0$  and outputs  $\infty$  in other cases. By choosing this function, the first-order condition for Eqn 2 will contain a non-linear ReLU function, which is the activation function in our GEQ. The equivalence between  $\mathbf{z}^*$  and the optimal solution for the hidden optimization problem (2) enables researchers to not only gain insights into OptEqs' behavior by understanding the underlying hidden optimization problem but also innovate by designing new models based on different problem formulations tailored to specific tasks. For instance, OptEqs introduces a module that promotes output sparsity, and Multi-branch OptEqs (MOptEqs) incorporates fusion modules to enhance the diversity among its branches. Despite these advancements and the incorporation of various modules inspired by different vision tasks, the performance of OptEqs and related models still falls short when compared to deep neural networks in image classification. This discrepancy suggests the existence of crucial components that limit the performance of equilibrium models.
To identify such a component, we delve into the hidden optimization problem of OptEqs and observe that it can be decomposed into two distinct parts: the regularizer term for the output features and the feature extraction term. While the feature extraction term is crucial as it depends on the input and determines the patterns extracted from the input features, the exploration of the regularizer term has been largely overlooked, with linear kernel functions being the predominant choice for feature extraction. Thereby, we believe that the limitations of previous equilibrium models stem from their feature extraction parts, as linear kernels struggle to capture complex features effectively. Building upon these insights, we take a step forward by leveraging the widely adopted Gaussian kernel for feature extraction in the hidden optimization problem.
Then by calculating the stationary condition for the above new hidden optimization problem, we propose our new type of OptEqs, the Gaussian kernels inspired equilibrium models (GEQ). The model involves a new attentive module induced by its hidden optimization problem and enjoys much better performances on classification tasks even compared with deep models. Furthermore, we also prove that the new model's outputs are equivalent to the outputs for OptEqs with weight-tied "infinite wide" mappings. Therefore, an interesting finding is that our model can be regarded as a "double-infinite" model because the original OptEqs can be regarded as a weight-tied "infinite deep" model. Apart from the above findings, the utilization of Gaussian kernels also makes our proposed model enjoy better generalization abilities. Besides the generalization abilities, we also analyze the stability of our GEQ and find its stability is better on various inputs. We summarize our contributions as follows:
- We first reformulate the OptEqs' hidden optimization problem with Gaussian kernels and propose a new equilibrium model called GEQ. It contains a new attention module induced by its hidden optimization problem and performs better on real-world datasets.
- We find that our GEQ can be regarded as a weight-tied neural network with both infinite width and depth, and better generalization ability through our analysis. Empirical results also confirm the superiority of our GEQ.
- We theoretically demonstrate the advantages of the stability of our GEQ compared with former OptEqs on various inputs. We also conduct experiments to validate such advantages.
2.1 Implicit Models
Most modern deep learning approaches provide explicit computation graphs for forward propagations and we call these models "explicit models". Contrary to these models, recent researchers proposed some neural architecture with dynamic computation graphs and we call them "implicit models". A
notable example of an implicit model is Neural ODEs [7], its architecture is encoded as a differential system and the implicit ODE solvers they used are equivalent to continuous ResNets that take infinitesimal steps. By representing the entire structure using differential systems, implicit models tap into the black box of traditional neural networks while offering increased flexibility and interpretability. Because of the flexibility and the interpretability of implicit models, the design of implicit models [15, 17] draws much attention these days. Many kinds of implicit models have been proposed, including optimization layers [11, 1], differentiable physics engines [40, 9], logical structure learning [50], differential programming [53, 43], and others [25, 49].
Among the various implicit models, OptEqs [52] and its multi-branch version MOptEqs [28] stand out as they not only exhibit superior performance compared to other implicit models but also explicitly establish the relationship between their structure and a well-defined optimization problem. Therefore, exploring better equilibrium models is a promising direction to achieve more interpretable neural architectures. However, it is worth noting that while OptEqs and its variants have shown promising results, their performance is still not entirely satisfactory, particularly when compared to the deep explicit models. Other works [51, 42] also show the connection between their architectures and optimization problems, but their performance is also not satisfying. Besides the above general models, there are many works design equilibrium models from the view of optimization to deal with their specific tasks, like implicit graph models [6, 26], certified robust models [27] and image denoising models [5]. However, these models can only work on their specific domains.
2.2 Infinite Wide Models and Kernel Methods in Deep Learning
By employing kernel methods to estimate the outputs of single-layer networks for various samples, researchers discover that such networks can exhibit characteristics of a Gaussian process (GP) when their parameters are randomly initialized with a large width limit [35]. Building upon this idea, recent researchers have extended these findings to neural networks with multiple layers [24, 10] and other architectures [37, 13]. These studies primarily focus on weakly-trained models, where the network parameters are randomly initialized and kept fixed throughout the training process except for the last classification layer [2]. Despite their "weakly-trained" nature, these models still provide valuable insights applicable to current neural networks. For instance, mean-field theory [4, 16, 19] explains phenomena such as gradient vanishing and exploding during back-propagation, which are relevant not only to single-layer networks but also to other structures like convolutional neural networks (CNNs) and recurrent neural networks (RNNs). Other researchers explore stationary kernels to enhance the interpretability of neural networks by designing different activation functions [34].
In addition to weakly trained models, recent studies [22, 2] introduce the concept of Neural Tangent Kernel (NTK) and its variants. These works have demonstrated that the sample kernel of infinitely wide networks, with appropriate initialization, can converge to a fixed neural tangent kernel when trained using gradient descent with infinitesimal steps (gradient flow). The NTK model is a theoretical construct with strict constraints, and its weights are not learned. It is important to note that although our model can also be seen as an infinitely wide model, there are several key differences between our approach and the aforementioned models. Firstly, our model utilizes kernel methods to operate on input features and output features, while the NTK models employ the kernel method on samples. Secondly, our GEQ model can be viewed as employing a "weight-tied infinite wide" projection that is parameterized by learnable parameters, allowing for updates during the training process. This contrasts with NTKs and NTK-DEQ [12](an equilibrium model constructed with vanilla NTK layers), where the weights are fixed and not learned. Therefore, despite the potential overlap in terminologies used in our paper and NTK-related works, our GEQ model differs significantly.
3 Gaussian kernel Inspired Equilibrium models
3.1 Formulation and Structure of GEQ
Before starting our analysis, we need to reformulate the original formulations of OptEqs' equilibrium equation (1) and hidden optimization problem (2) for convenience. We replace  $\tilde{\mathbf{W}}_c\tilde{\mathbf{W}}^\top$  with  $\mathbf{W}_c$ ,  $\mathbf{W} := \tilde{\mathbf{W}}^\top$ , and replace  $\mathbf{z}$  with  $\tilde{\mathbf{W}}^{-1\top}\tilde{\mathbf{z}}$ . Then the original OptEqs' optimization problem can be
reformulated as:
$$
\min  _ {\mathbf {z}} G (\mathbf {z}; \mathbf {x}) = \min  _ {\mathbf {z}} \left[ \mathbf {1} ^ {\top} f (\mathbf {z}) + \frac {1}{2} \| \mathbf {z} \| _ {2} ^ {2} - \langle \mathbf {U x} + \mathbf {b}, \mathbf {z} \rangle - \frac {1}{2} \| \mathbf {W z} \| _ {2} ^ {2} \right]. \tag {3}
$$
With the new formulation, we can rewrite the equilibrium equation for OptEqs with input  $\mathbf{x}$  by calculating Eqn (3)'s first order stationary condition  $\nabla G = \mathbf{0}$  and then reformulate it as follows:
$$
\mathbf {z} ^ {*} = \sigma \left(\mathbf {W} ^ {\top} \mathbf {W} \mathbf {z} ^ {*} + \mathbf {U} \mathbf {x} + \mathbf {b}\right), \tag {4}
$$
where  $\sigma$  is the ReLU activation function,  $\mathbf{U},\mathbf{W},\mathbf{b}$  are learnable parameters trained by optimizing loss functions (like cross-entropy loss). From problem Eqn (3), one can see that the GEQ's outputs try to extract features by minimizing the similarity term with the input feature  $\mathbf{U}\mathbf{x} + \mathbf{b}$  through a linear kernel function with some constraints defined in its regulation terms to prevent the trivial outputs. Such an explanation can also extend to other DEQs [3, 51] under the symmetric weight constraints. However, as other traditional machine learning mechanisms show, linear kernel functions cannot perform well when processing complex inputs. We deem that this term will also restrict the performance in equilibrium models. We note that the symmetric constraints won't influence the final performance much as many works [32, 21] show.
A natural consideration arises as to whether we can utilize alternative kernel functions to extract input features for the equilibrium state. However, we find that other equilibrium models employing different kernels with inner products, like the polynomial kernel and sigmoid kernel, lead to a similar structure to OptEqs with appropriate weight re-parameterization and lead to similar empirical results. We provide a detailed discussion of the related models in Appendix 4.6. Thereby, we decide to use the Gaussian kernels, and our new hidden optimization equation is formulated as follows:
$$
\min  _ {\mathbf {z}} G (\mathbf {z}; \mathbf {x}) = \min  _ {\mathbf {z}} \left[ \mathbf {1} ^ {\top} f (\mathbf {z}) + \frac {1}{2} \| \mathbf {z} \| _ {2} ^ {2} - \frac {1}{2 \gamma} e ^ {- \gamma \| \mathbf {U x} + \mathbf {b} - \mathbf {W z} \| _ {2} ^ {2}} \right], \tag {5}
$$
where  $\gamma$  is the hyperparameter denoting the reciprocal of Gaussian kernels' variance for scaling. Calculating  $\nabla G = \mathbf{0}$  for new  $G$ , we can get the Gaussian kernel inspired Equilibrium models (GEQ) as the following fixed-point equation:
$$
\mathbf {z} ^ {*} = \sigma \left[ e ^ {- \gamma \| \mathbf {U x} + \mathbf {b} - \mathbf {W z} ^ {*} \| _ {2} ^ {2}} \mathbf {W} ^ {\top} \left(- \mathbf {W z} ^ {*} + \mathbf {U x} + \mathbf {b}\right) \right]. \tag {6}
$$
Compared with linear kernels, Gaussian kernels can easily extract the non-linear relations from the input features and show more stable and powerful performance in SVM and other machine learning methods [21, 44]. We also find that the formulation of our GEQ is similar to adding a new attention module to the original equilibrium models. Therefore, our GEQ is supposed to enjoy more representative abilities than the original OptEqs. In the following parts of this section, we will analyze the theoretical advantages of our GEQ against the vanilla OptEqs. And we also empirically evaluate GEQ's performance in the following sections.
3.2 GEQ equals to the OptEqs with infinite width
Like other Gaussian-related models, our GEQ model can also be regarded as computing similarities by mapping them to an infinite-dimensional space. This allows GEQ to extract input features at the infinite-dimensional level, enabling the capture of non-linear dependencies in the input space. Essentially, our GEQ can be seen as a specialized version of OptEqs operating within the infinite-dimensional space after mapping input features  $\mathbf{x}$  and output embedding  $\mathbf{z}$  to this expanded domain.
Proposition 1. The output of our GEQ (Eqn (6)) is the same as a special OptEqs' output whose hidden optimization problem is defined as follows:
$$
\min  _ {\mathbf {z}} G (\mathbf {z}; \mathbf {x}) = \min  _ {\mathbf {z}} \left[ \mathbf {1} ^ {\top} f (\mathbf {z}) + \frac {1}{2} \| \mathbf {z} \| _ {2} ^ {2} - \lambda \left\langle \Phi_ {\mathbf {U}} (\mathbf {x} + \mathbf {U} ^ {- 1} \mathbf {b}), \Phi_ {\mathbf {W}} (\mathbf {z}) \right\rangle \right], \tag {7}
$$
where  $f$  is the positive indicator function and  $(1 + \partial f)^{-1}$  is the ReLU activation function,  $\lambda = e^{-\gamma \| \mathbf{Ux} + \mathbf{b}\| _2^2}e^{-\gamma \| \mathbf{Wz}\| _2}$ , and  $\Phi_{\mathbf{W}}(\mathbf{z}) = [\mathbf{1},\sqrt{\mathbf{2}\gamma}\Phi_{\mathbf{W}}^{(\mathbf{1})}(\mathbf{z}),\dots,\sqrt{(2\gamma)^i / i!}\Phi_{\mathbf{W}}^{(i)}(\mathbf{z}),\dots]\in \mathbb{R}^{1\times \infty}$  which maps the inputs to the infinite-dimensional space with  $\Phi_{\mathbf{W}}^{(i)}:\mathbb{R}^n\to \mathbb{R}^{in^i}$  defined as follows:
$$
\Phi_ {\mathbf {W}} ^ {(i)} = \left[ \underbrace {\overbrace {(W x) _ {0} (W x) _ {0} \dots (W x) _ {0}} ^ {i} , \overbrace {(W x) _ {0} (W x) _ {0} \dots (W x) _ {1}} ^ {i} , \dots , \overbrace {(W x) _ {j} (W x) _ {k} \dots (W x) _ {m}} ^ {i} , \dots} _ {i n ^ {i}} \right], \tag {8}
$$
where  $(Wx)_{j}$  denotes the  $j$ -th element of vector  $\mathbf{Wx}$ .
Based on the analysis provided above, it becomes evident that the hidden optimization problem of our GEQ exhibits a similar formulation to a specific OptEqs, whose inputs  $\mathbf{x}$  and outputs  $\mathbf{z}$  are mapped to an infinite-dimensional space using the weight-tied infinite wide mapping  $\Phi_{\mathbf{W}}$  and  $\Phi_{\mathbf{U}}$ . Given that both GEQ and OptEqs are derived from their respective hidden optimization problems, the equivalence in these problems implies the existence of the same equilibrium states for both models. Consequently, our GEQ can be considered an extension of the "infinite-depth" OptEqs to the "infinite-width" domain. Since wider neural networks are generally expected to perform better on classification tasks, we can infer that our GEQ outperforms vanilla equilibrium models like OptEqs. We further support this claim with the theoretical analysis illustrated in the subsequent sections.
3.3 Generalization abilities for our GEQ
Apart from the above empirical intuition, we are going to prove our GEQ's generalization advantages over OptEqs using the generalization bound under the PAC-Bayesian framework [36]. For convenience, we use  $f_{geq}(\mathbf{x})$  denotes the equilibrium state  $\mathbf{z}^*$  for input  $\mathbf{x}$ . Then we use the expected margin loss  $\mathcal{L}_{\eta}(f_{geq}^{c})$  at margin  $\eta$  of our GEQ on the data distribution  $\mathcal{D}$  for classification, which is defined as follows,
$$
\mathcal {L} _ {\eta} (f _ {g e q} ^ {c}) = \mathbb {P} _ {(\mathbf {x}, \mathbf {y}) \sim \mathcal {D}} \left[ f _ {g e q} ^ {c} (\mathbf {x}) _ {y} \leq \eta + \max _ {j \neq y} f _ {g e q} ^ {c} (\mathbf {x}) _ {j} \right], \tag {9}
$$
where  $f_{geq}^{c}(\mathbf{x}) = \mathbf{W}_{c}f_{geq}(\mathbf{x}) + \mathbf{b}_{c}$  stands for GEQ's final prediction at input  $\mathbf{x}$  with learnable parameters  $\mathbf{W}_c$  and  $\mathbf{b}_c$ , and the index  $j, y$  here denote the prediction score for certain class. Then we can analyze the generalization bound for our GEQ following the former work's settings [38].
Proposition 2. If input  $\| \mathbf{x}\| _2$  is bounded by  $B$ ,  $\mu \coloneqq \max \left\{\| \mathbf{U}\| _2,\| \mathbf{W}\| _2,\| \mathbf{W}_c\| _2,\| \mathbf{b}\| _2\right\} < 1$ , then we have following results for GEQ and OptEqs with ReLU activation. For  $\forall \delta ,\eta >0$  with probability at least  $1 - \delta$  over the training set of size  $M$ , we have:
$$
\begin{array}{l} \mathcal {L} _ {0} \left(f _ {g e q} ^ {c}\right) \leq \hat {\mathcal {L}} _ {\eta} \left(f _ {g e q} ^ {c}\right) + \sqrt {\frac {1 6 h \ln (2 4 h) \left[ \beta_ {\max } \mu^ {4} B + (2 \mu \beta_ {\max } + 1) (1 - \beta_ {\max } m) \mu B + (1 - \beta_ {\max } m) ^ {2} \right] ^ {2} \mathcal {B} _ {W}}{\eta^ {2} (1 - \beta_ {\max } m) ^ {4} M}} + \frac {\ln (\frac {M \sqrt {M}}{\delta})}{M}, \tag {10} \\ \mathcal {L} _ {0} (f _ {o p t e q} ^ {c}) \leq \hat {\mathcal {L}} _ {\eta} (f _ {o p t e q} ^ {c}) + \sqrt {\frac {1 6 h \ln (2 4 h) [ \mu^ {3} B + (1 - m) \mu B + (1 - m) ^ {2} ] ^ {2} \mathcal {B} _ {W}}{\eta^ {2} (1 - m) ^ {4} M} + \frac {\ln (\frac {M \sqrt {M}}{\delta})}{M}}, \\ \end{array}
$$
where  $\hat{\mathcal{L}}_{\eta}(f_{geo}^{c})$  denotes the empirical margin loss on the training set, the maximum scaling number is defined by  $\beta_{max}\coloneqq \max_{\mathbf{x}\in \mathcal{D}}e^{-\gamma \| \mathbf{Ux} + \mathbf{b} - \mathbf{Wz}\| _2^2},\mathcal{B}_W\coloneqq \| \mathbf{W}^\top \mathbf{W}\| _F^2 +\| \mathbf{U}\| _F^2 +\| \mathbf{b}\| _2^2 +\| \mathbf{W}_c\| _F^2 +\| \mathbf{b}_c\| _2^2,$  and  $m = \| \mathbf{W}^{\top}\mathbf{W}\|_{2}$  is less than 1 to ensure the convergence of equilibrium models.
Remark 1. If  $\beta_{\mathrm{max}} < 0.8$  and  $\mu, m > 0.9$ , we can get  $\frac{\beta_{\mathrm{max}}\mu}{1 - \beta_{\mathrm{max}}m} < \frac{1}{1 - m}$  and  $\frac{2\mu\beta_{\mathrm{max}} + 1}{1 - \beta_{\mathrm{max}}m} \leq \frac{1}{1 - m}$ . In the meanwhile, our GEQ's generalization bound is tighter than the original OptEq.
In practical experiments, we find that the above conditions for  $\beta_{\mathrm{max}}$  and  $\mu, m$  are satisfied in most cases. And the following experiments also support our above theoretical advantages. However, the above bound is not tight, and how to approximate a much tighter bound for equilibrium models still needs exploring.
3.4 GEQ enjoys More Stable Performance
Apart from better performance, Gaussian kernel stands out as one of the most extensively employed kernels in machine learning tasks owing to its stability across various input scenarios. Motivated by this, we aim to investigate whether incorporating Gaussian kernels into our equilibrium models can enhance the model's stability across diverse inputs. Firstly, we are going to estimate output changes with respect to the input perturbations.
Proposition 3. If norms for the inputs and outputs are bounded by  $B$ , the spectral norm for the weight parameter  $\mathbf{W}$ ,  $\mathbf{U}$  of equilibrium models with ReLU activation are bounded by  $\mu < 1$  to ensure
convergence, then we have the conclusions as below:
$$
\left\| f _ {g e q} \left(\mathbf {x} _ {1}\right) - f _ {g e q} \left(\mathbf {x} _ {2}\right) \right\| _ {2} \leq L _ {g e q} \| \mathbf {x} _ {1} - \mathbf {x} _ {2} \| _ {2} = \frac {\beta_ {\max } \mu^ {2} + \sqrt {\gamma} B \mu^ {3}}{1 - \beta_ {\max } \mu^ {2} - \sqrt {\gamma} B \mu^ {3}} \| \mathbf {x} _ {1} - \mathbf {x} _ {2} \| _ {2}, \tag {11}
$$