Optimization inspired Multi-Branch Equilibrium Models

1 Introduction . Whereas Deep Neural Networks ( DNNs ) have achieved great success in many real-world tasks such as computer vision and neural language process , the limited interpretability of DNNs greatly hinders their further development . However , many traditional machine learning methods , e.g. , matrix recovery ( Zhang et al. , 2018b ; 2015 ; Liu & Li , 2016 ) , subspace clustering ( You et al. , 2016 ) , image deblurring ( Liu et al. , 2014 ) and so on , can be interpreted into minimizing the following objective functions : min Z , E f ( Z ) + g ( E ) , s.t . X = AZ + BE . where A ∈ Rm×d1 , B ∈ Rm×d2 , X ∈ Rm×n , and f ( · ) and g ( · ) are convex functions designed by the modelling of their properties . For this count , the interpretability of such methods is much better than DNNs . Furthermore , these methods can also enjoy state-of-the-art and robust performance on these tasks . We call such interpretability these methods enjoy as `` mathematical interpretability '' , i.e. , whether the whole network structure can be summarized as a compact mathematical model that can be mathematically analyzed . Except for these methods , the Optimization Induced Equilibrium Networks ( OptEqs ) proposed by Xie et al . ( 2021 ) recover its whole system for recognition to an optimization problem . The forward propagation for OptEqs tries to solve Eqn.1 to get output y , z∗ ∈ Rd for input x ∈ Rdin . We call the front part of Eqn.1 as OptEqs ’ equilibrium equation , which is the central part for equilibrium models . And its forward procedure can be regarded as solving the hidden objective functions shown as Eqn.2 , z∗ = W > σ ( Wz∗ +Ug ( x ) + b ) , y = Wcz ∗ , ( 1 ) min z G ( z ; g ( x ) ) = min z [ 1 > f ( W−1 > z ) − 〈 Ug ( x ) + b , W−1 > z 〉 + 1 2 ‖W−1 > z‖2 − 1 2 ‖z‖2 ] . ( 2 ) W , U , Wc are learnable weights , g is the convolution layers projecting input x from real-world domain to feature space like other neural architectures , and f is the penalty for the outputs . Although the interpretability for OptEqs is good , it still remains some weaknesses worth exploring . First , although OptEqs performs better than some implicit models , these models only consider a single view ( or resolution ) of the input in their implicit parts . We call these models single-view implicit models in the following . However , nearly all the state-of-the-art pattern recognition systems ( Lee et al. , 2009b ; Wang et al. , 2019a ; Huang et al. , 2017 ; He et al. , 2016 ; Burt & Adelson , 1987 ) benefit from the multi-layer or multi-resolution feature extractors in domains like computer vision and audio processing . Secondly , new modules induced by the objective function they proposed are limited , which mainly consider the math properties of the output features like sparsity , but whether these properties can benefit image recognition tasks is unexplored . Besides OptEqs , MDEQ ( Bai et al. , 2020 ) constructs its equilibrium equation z∗ = F ( z∗ ; x ) and block F with the inspiration of explicit models especially HRNet Wang et al . ( 2019a ) . Then they solve the equilibrium equation by the accelerated algorithm for output z∗ . Although it shows the state-of-the-art performance with efficient memory cost as other implicit models , such strategies will make the model lose `` interpretability '' since the MDEQ is a black box because its output is solved in an implicit way ( root-finding method ) , which makes analysis on its features of intermediate layers almost impossible . Furthermore , its complicated structure also hinders its analysis from mathematical systems . More efficient blocks with good interpretability still needs exploring . Motivated by the limitations of the above models , we would like to design a multi-scale DEQ structure with state-of-the-art performance and preserve mathematical interpretability from the inspiration of OptEqs and multi-scale models . The contributions for our work are listed below : • We propose a multi-branch implicit model , called Multi-branch OptEqs ( MOptEqs ) , which efficiently utilizes different scales inputs with better performance and smaller model size . Furthermore , it still retains its connection to an optimization problem . • We propose some properties modelled as new terms for the model ’ s hidden objective function from our analysis on the relationships between branches and their hierarchical dependencies . With new problem ’ s formulation , we obtain the fusion module in our MOptEqs , called Hierarchical Heritage and Diversity Module ( HH & D Module ) . • Apart from the new module design , we also propose the Perturbation Enhanced training strategy ( PE ) for our MOptEqs from our analysis on the model ’ s hidden objective function . The new strategy can not only enhance the performance of our MOptEqs , but also improve the robustness of our models . 1.1 Related Works . Implicit Models . Nearly all modern deep learning approaches use explicit models , which provide an explicit computation graph for the forward propagation . In contrast , the computation graphs for implicit models are `` flexible '' or can be assumed as having `` infinite '' depth . For example , Neural ODEs ( Chen et al. , 2018 ; Massaroli et al. , 2020 ) encode their neural architectures by a differential system with learnable parameters . Then the implicit ODE solvers they used is equivalent to a continuous ResNet taking infinitesimal steps . Furthermore , the training process of Neural ODEs can be depicted as finding a differential system of a certain type ( like heat equations ) by updating its learnable parameters which demonstrate that Neural ODE also enjoy interpretability to a certain extent . Moreover , DEQ ( Bai et al. , 2019 ; Winston & Kolter , 2020 ) is another class of implicit models . The central part of DEQ is the designation of the equilibrium equations z∗ = F ( z∗ ; x ) and block F . Its forward procedure is trying to solve the equilibrium equations with input x to get the equilibrium state z∗ as output by accelerated algorithms . Since the z∗ is fixed given F and x , its inference can be regarded as forwarding an explicit network stacked by the same block F for infinite times . For example , DEQ chooses their F by one set of Conv+ReLU while MDEQ constructs an HRNet-like block as theirs . However , no evidence shows that the best block constructions in explicit models can perform the best in the DEQ scheme . The construction of DEQ blocks or equilibrium equations is an open question worth exploring . Furthermore , most DEQ blocks do not have any mathematical insights ( like the diffusion process in Neural ODE ) and perform totally a black box with limited interpretability . Since implicit models usually adopt accelerated root-finding algorithms for their forward outputs and backward gradients , they enjoy the advantages of constant and more efficient memory costs compared with DNNs . Due to the above advantages , the design of implicit models draws much attention these days ( Ghaoui et al. , 2019 ; Gould et al. , 2019 ) . Apart from the models illustrated above , many kinds of other implicit models have been proposed , including differentiable physics engines ( Qiao et al. , 2020 ; de Avila Belbute-Peres et al. , 2018 ) , logical structure learning ( Wang et al. , 2019b ) and implicit neural blocks ( Li et al. , 2020 ) . Model Interpretaiblity . Many researchers nowadays are trying to make their proposed method more interpretable . Although there are many kinds of approaches to achieve this goal , we marginally divide them into two parts : `` Empirical '' and `` Mathematical '' interpretability . Many works , such as ( Zhang et al. , 2018a ; Zhang & Zhu , 2018 ; Bau et al. , 2018 ) , attempt to empirically disassemble the black box by characterizing some statistical or structural information like the outputs and gradients of neural networks ’ middle layers . However , these works can not be directly implemented on DEQ models since implicit ones do not have explicit depths like DNNs , which makes analyzing such models by dissecting each layer ’ s behavior almost impossible . Apart from that , researchers ( Djolonga & Krause , 2017 ; Amos & Kolter , 2017 ; Xie et al. , 2019 ; Chan et al. , 2020 ) also managed to understand the neural architecture by linking it to a mathematical problem . In this way , researchers can analyze the black box by its corresponding mathematical problem . Furthermore , new components can be proposed due to the analysis of these problems . Our model aims to design a new DEQ architecture with proper mathematical interpretability . Compared with the former works , our model can deal with the classification tasks with better performances with novel components . 2 Multi-Branch Optimized Induced Equilirium Models . 2.1 The proposed architecture for the Multi-Branch OptEqs . Inspired by OptEqs , we first design an objective function for our tasks to solve and then use its first-order stationary conditions as MOptEqs ’ equilibrium equation . The base of our function is formed by summarizing several objective functions of different scales defined for OptEqs ( Details are stated in Appendix.A.1 ) . However , such a model can not obtain satisfying results since these branches are independent . For this count , we need to design some terms describing the dependencies of each branch in the hidden objective function , as shown in the following paragraphs . Then we can obtain the equilibrium equations for our MOptEqs and our MOptEqs ’ block F by analyzing the first-order stationary condition of our designed optimization problem . Hierarchical Heritage Modeling . State-of-the-art explicit models for image tasks are explicitly structured into sequential stages and process different resolutions hierarchically ( He et al. , 2016 ; Shelhamer et al. , 2017 ; Lee et al. , 2009a ) , which implies that features should be extracted hierarchically . In other words , the posterior branches should inherit from their prior ones , which we call such property hierarchical heritage . In our work , we are going to formulate such correlation of neighbouring branches in the hidden objective function by adding the following inner-product term into its origin : H ( zi , zj ) = z > i Pi : jzj ( 3 ) where zi ∈ Rdi , zj ∈ Rdj , Pi : j ∈ Rsi×sj denotes the Average Downsample ( its transpose can be regarded as a weighted nearest upsample ) or Identity matrix suited to the shape of zi and zj ( i < j or j = 1 , i = L ) . This term only estimate the summation for the similarity of i-th and j-th branch ’ s channels with the same channel index , which we call them corresponding channels , as shown in the above equation . When zi and zj are similar , the inner product will be large . Otherwise , the result will become small , which implies that the relationship of corresponding channels is weak . In this way , we can ensure the similarity of corresponding channels of the near branches . Except for the relations between different branches ’ corresponding channels , the relationships between dis-corresponding channels are explored in the following section . Diversity Modeling . In addition to the hierarchical relations for corresponding channels , works ( Pang et al. , 2019 ; Amada et al. , 2021 ) have shown enhancing the diversities between branches can also improve the model ’ s improvements . For this account , we also consider making different branches extract various features to improve the representation abilities of different branches . To achieve such a goal , we add a diversity term in the objective function Eqn . ( 4 ) if we assume that i < j , D ( zi , zj ) = C∑ k1=1 C∑ k2=1 k2 6=k1 |vec−1 ( zi ) ( k1 ) > vec−1 ( Pi : jzj ) ( k2 ) | , ( 4 ) where vec−1 is the inverse vectorization operator convert the vectorized feature zi ∈ Rdi×1 to a matrix in R di C ×C with C channels and vec−1 ( zi ) ( k ) ∈ R di C ×1 denotes the k-th channel of i-th branch in practice . Since we have already considered the relationships between the corresponding channels in the hierarchical term , we only tries to estimate the diversities of dis-corresponding channels of the i-th and j-th branches in this term . As the diversity term goes small , the diversities between the branches become stronger . The architecture of MOptEqs and its hidden problem . With the two terms we proposed , we can reformulate the hidden objective for our problem by replacing the L2 norm with our new modelling as follows : min z1 , ... , zL G ( z1 , ... , zL ; g ( x ) ) = min z1 , ... , zL L∑ i=1 [ 1 > f ( W−1 > i zi ) − 〈 Uig ( x ) + bi , W −1 > i z 〉 + 1 2 ( λD ( zi , zi+1 ) + ‖W−1 > zi‖2 −H ( zi , zi+1 ) ) ] . ( 5 ) g ( x ) is the input feature for the raw input x , λ > 0 is a hyperparameter and zi ∈ Rdi×1 are the final outputs of MOptEqs . We can choose f to constrain zi on our demand and will influence model ’ s activation function . If we choose f ( x ) = I { x ≥ 0 } to ensure the outputs to be positive , then the activation function is ReLU . We set D ( zL , zL+1 ) = D ( zL , z1 ) and H ( zL , zL+1 ) = H ( zL , z1 ) when i = L to complete the loop . The two terms we added tries to make the corresponding channels of different branches correlated by maximizing the H term and enhancing the diversities of different branches by minimizing D term . We note that H and D are not conflict since they are handling different pairs of channels for branches . As the following proposition shows , we finally get the equilibrium equations ( Eqn . ( 6 ) ) for our MOptEqs by calculating the first-order stationary conditions ∇ziG = 0 for problem G. Proposition 1 The proposed multi-branch structure induced by Eqn ( 5 ) can be depicted as solving the equilibrium points z̃∗ : = [ z > ∗1 , ... , z > ∗L ] > ∈ R ∑L i=1 di for the following equations : z̃ = W̃ > σ ( W̃h ( z̃ ) + Ũg ( x ) + b̃ ) where W̃ = W1 . . . WL , Ũ = [ U > 1 , ... , U > L ] > , b̃ = [ b > 1 , ... , b > L ] > . ( 6 ) And h ( z̃ ) = [ h1 ( z1 , zL , z2 ) > , .. , hi ( zi , zi−1 , zi+1 ) > , .. , hL ( zL , zL−1 , z1 ) > ] > and each hi is defined as the mapping from Rdi × Rdi+1 × Rdi−1 to Rdi , h ( zi , zi−1 , zi+1 ) = 1 2 ( Pi : i+1zi+1 +P > i−1 : izi−1 ) − λ 2 vec ( vec−1 ( Pi : i+1zi+1 ) sign [ Mi : i+1 − diag ( Mi : i+1 ) ] ) − λ 2 vec ( vec−1 ( P > i−1 : izi−1 ) sign [ M > i−1 : i − diag ( Mi−1 : i ) ] ) , Mi : i+1 =vec−1 ( Pi : i+1zi+1 ) > vec−1 ( zi ) , Mi−1 : i =vec−1 ( zi ) > vec−1 ( P > i−1 : izi−1 ) . The equilibrium points for the above formula are also the first-order stationary points for the optimization problem Eqn ( 5 ) . The proof of the proposition is listed in Appendix.A.3 . In this manner , we obtain the proposed architecture MOptEqs , its forward propagation is equivalent to find the equilibrium points of Eqn ( 6 ) , and such process can also be regarded as solving the stationary points of the optimization problem ( 5 ) . The whole structure is also shown in Figure.1 ( a ) and we also draw figures to illustrate the practical process of hi ( HH & D ) on its right ( Figure.1 ( b ) ) . We also evaluate the effectiveness of such modelling in Section . 3.2 ’ s experiments .

Summary: The paper introduces a multi-scale network (called Multi-Branch OptEqs or MOptEqs) that extends on recent results of [1]. Specifically, the authors propose to use several DEQ models to incorporate samples of multi-resolutions. The authors further propose a modified training strategy (called perturbation enhanced training) to improve the performance of the model. Empirically, the proposed model shows better performances compared to the previous DEQ models on CIFAR-10, CIFAR-100, and Imagenette datasets.  [1] Xie, Xingyu, et al. "Optimization Induced Equilibrium Networks." arXiv preprint arXiv:2105.13228 (2021).