I S A TTENTION B ETTER T HAN M ATRIX D ECOMPOSITION ?
A BSTRACT
As an essential ingredient of modern deep learning, attention mechanism, especially self-attention, plays a vital role in the global correlation discovery.
However, is hand-crafted attention irreplaceable when modeling the global context?
Our intriguing ﬁnding is that self-attention is not better than the matrix decomposition (MD) model developed 20 years ago regarding the performance and computational cost for encoding the long-distance dependencies.
We model the global context issue as a low-rank completion problem and show that its optimization algorithms can help design global information blocks.
This paper then proposes a series of
Hamburgers, in which we employ the optimization algorithms for solving MDs to factorize the input representations into sub-matrices and reconstruct a low-rank embedding.
Hamburgers with different MDs can perform favorably against the popular global context module self-attention when carefully coping with gradients back-propagated through MDs.
Comprehensive experiments are conducted in the vision tasks where it is crucial to learn the global context, including semantic segmentation and image generation, demonstrating signiﬁcant improvements over self-attention and its variants. Code is available.
I NTRODUCTION
Since self-attention and transformer (Vaswani et al., 2017) showed signiﬁcant advantages over recurrent neural networks and convolutional neural networks in capturing long-distance dependencies, attention has been widely adopted by computer vision (Wang et al., 2018;
Zhang et al., 2019a) and natural language processing (Devlin et al., 2019) for global information mining.
However, is hand-crafted attention irreplaceable when modeling the global context?
This paper focuses on a new approach to design global context modules.
The key idea is, if we formulate the inductive bias like the global context into an objective function, the optimization algorithm to minimize the objective function can construct a computational graph, i.e., the architecture we need in the networks .
We particularize this idea by developing a counterpart for the most representative global context module, self-attention.
Considering extracting global information in the networks as ﬁnding a dictionary and the corresponding codes to capture the inherent correlation, we model the context discovery as low-rank completion of the input tensor and solve it via matrix decomposition.
This paper then proposes a global correlation block, Hamburger, by employing matrix decomposition to factorize the learned representation into sub-matrices so as to recover the clean low-rank signal subspace.
The iterative optimization algorithm to solve matrix decomposition deﬁnes the central computational graph, i.e., Hamburger’s architecture.
Our work takes advantage of the matrix decomposition models as the foundation of Hamburger, including Vector Quantization (VQ) (Gray & Neuhoff, 1998), Concept Decomposition (CD) (Dhillon & Modha, 2001), and Non-negative Matrix Factorization (NMF) (Lee & Seung, 1999).
Additionally, instead of directly applying Back-Propagation Through Time (BPTT) algorithm (Werbos et al., 1990) to differentiate the iterative optimization, we adopt a truncated BPTT algorithm, i.e., one-step gradient, to back-propagate the gradient effectively.
We illustrate the advantages of Hamburger in the fundamental vision tasks where global information has been proven crucial, including semantic segmentation and image generation.
The experiments prove that optimization-designed Hamburger can perform competitively with state-of-the-art attention models when avoiding the unstable gradient back-propagated through the iterative computational graph of MD.
Hamburger sets new state-ofthe-art records on the PASCAL VOC dataset (Everingham et al., 2010) and PASCAL Context dataset (Mottaghi et al., 2014) for semantic segmentation and surpasses existing attention modules for GANs in the large scale image generation on ImageNet (Deng et al., 2009).
The contributions of this paper are listed as follows:
• We show a white-box approach to design global information blocks, i.e., by turning the optimization algorithm that minimizes an objective function, in which modeling the global correlation is formulated as a low-rank completion problem, into the architecture.
• We propose Hamburger, a light yet powerful global context module with O ( n ) complexity, surpassing various attention modules on semantic segmentation and image generation.
• We ﬁgure out that the main obstacle of applying MD in the networks is the unstable backward gradient through its iterative optimization algorithm.
As a pragmatic solution, the proposed one-step gradient facilitates the training of Hamburger with MDs.
M ETHODOLOGY 2. W ARM U P
Since matrix decomposition is pivotal to the proposed Hamburger, we ﬁrst review the idea of matrix decomposition.
A common view is that matrix decomposition factorizes the observed matrix into a product of several sub-matrices, e.g., Singular Value Decomposition.
However, a more illuminating perspective is that, by assuming the generation process, matrix decomposition acts as the inverse of the generation, disassembling the atoms that make up the complex data.
From the reconstruction of the original matrices, matrix decomposition recovers the latent structure of observed data.
Suppose that the given data are arranged as the columns of a large matrix X
=
[ x 1 , · · · , x n ] ∈ R d × n .
A general assumption is that there is a low-dimensional subspace, or a union of multiple subspaces hidden in X .
That is, there exists a dictionary matrix D
=
[ d 1 , · · · , d r ] ∈ R d × r and corresponding codes C
= [ c 1 , · · · , c n ] ∈ R r × n that X can be expressed as
[Equation]
where ¯ X ∈ R d × n is the output low-rank reconstruction, and E ∈ R d × n is the noise matrix to be discarded.
Here we assume that the recovered matrix ¯ X has the low-rank property, such that
[Equation]
Different MDs can be derived by assuming structures to matrices D , C , and E (Kolda & Bader, 2009; Udell et al., 2016).
MD is usually formulated as an objective with various constraints and then solved by optimization algorithms, with classic applications to image denoising (Wright et al., 2009;
Lu et al., 2014), inpainting (Mairal et al., 2010), and feature extraction (Zhang et al., 2012).
2. P ROPOSED M ETHOD
We focus on building global context modules for the networks without painstaking hand-crafted design.
Before starting our discussion, we review the representative hand-designed context block self-attention pithily.
The attention mechanism aims at ﬁnding a group of concepts for further conscious reasoning from massive unconscious context (Xu et al., 2015; Bengio, 2017; Goyal et al., 2019).
As a representative, self-attention (Vaswani et al., 2017) is proposed for learning long-range dependencies in machine translation,
[Equation]
where Q , K , V ∈ R n × d are features projected by linear transformations from the input.
Selfattention extracts global information via attending all tokens at a time rather than the typical one-byone processing of recurrent neural networks. [Figure] 𝑪
[Figure]
𝑑×𝑛 𝑑×𝑟 𝑟×𝑛 𝑑×𝑛
Though self-attention and its variants achieved great success, researchers are confronted with (1) developing new global context modules based on self-attention, typically via hand-crafted engineering, and (2) explaining why current attention models work.
This paper bypasses both issues and ﬁnds a method to easily design global context modules via a well-deﬁned white-box toolkit.
We try to formulate the human inductive bias, like the global context, as an objective function and use the optimization algorithm to solve such a problem to design the module’s architecture.
The optimization algorithm creates a computational graph, takes some input, and ﬁnally outputs the solution.
We apply the computational graph of optimization algorithms for the central part of our context module.
Based on this approach, we need to model the networks’ global information issue as an optimization problem.
Take the convolutional neural networks (CNN) as an example for further discussion.
The networks output a tensor X ∈ R C × H × W after we feed into an image.
Since the tensor can be seen as a set of HW C -dimensional hyper-pixels, we unfold the tensor into a matrix X ∈ R C × HW .
When the module learns the long-range dependencies or the global context, the hidden assumption is that the hyper-pixels are inherently correlated.
For the sake of simplicity, we assume that hyper-pixels are linearly dependent, which means that each hyper-pixel in X can be expressed as the linear combination of bases whose elements are typically much less than HW .
In the ideal situation, the global information hidden in X can be low-rank .
However, due to vanilla CNN’s poor ability to model the global context (Wang et al., 2018; Zhang et al., 2019a), the learned X is usually corrupted with redundant information or incompleteness.
The above analysis suggests a potential method to model the global context, i.e., by completing the low-rank part ¯ X in the unfolded matrix X and discarding the noise part E , using the classic matrix decomposition models described in Eq.
(1), which ﬁlters out the redundancy and incompleteness at the same time.
We thus model learning the global context as a low-rank completion problem with matrix decomposition as its solution.
Using the notion of Sec. 2.1, the general objective function of matrix decomposition is
[Equation]
where L is the reconstruction loss, R 1 and R 2 are regularization terms for the dictionary D and the codes C .
Denote the optimization algorithm to minimize Eq.
(4) as M .
M is the core architecture we deploy in our global context module.
To help readers further understand this modeling, We also provide a more intuitive illustration in Appendix G.
In the later sections, we introduce our global context block, Hamburger, and then discuss detailed MD models and optimization algorithms for M .
Finally, we handle the gradient issue for back-propagation through matrix decomposition.
2.2. H AMBURGER
Hamburger consists of one slice of “ham” (matrix decomposition) and two slices of “bread” (linear transformation).
As the name implies, Hamburger ﬁrst maps the input
Z ∈ R d z × n into feature space with a linear transformation W l ∈ R d × d z , namely “lower bread”, then uses matrix decomposition
M to solve a low-rank signal subspace, corresponding to the “ham”, and ﬁnally transforms extracted signals into the output with another linear transformation W u ∈ R
d
z × d , called “upper bread”,
[Equation]
[Equation]
where M is matrix decomposition to recover the clear latent structure, functioning as a global nonlinearity .
Detailed architectures of M , i.e., optimization algorithms to factorize X , are discussed in Sec. 2.2.2.
Fig.
1 describes the architecture of Hamburger, where it collaborates with the networks via Batch Normalization (BN) (Ioffe & Szegedy, 2015), a skip connection, and ﬁnally outputs Y ,
Y =
Z + BN( H ( Z )) .
2.2. H AMS
This section describes the structure of “ham”, i.e., M in Eq.
As discussed in the previous section, by formulating the global information discovery as an optimization problem of MD, algorithms to solve MD naturally compose M .
M takes the output of “lower bread” as its input and computes a low-rank reconstruction as its output, denoted as X and ¯ X , respectively.
[Equation]
We investigate two MD models for M , Vector Quantization (VQ), and Non-negative Matrix Factorization (NMF) to solve D and C and reconstruct ¯ X , while leaving Concept Decomposition (CD) to
Appendix B. The selected MD models are introduced brieﬂy because we endeavor to illustrate the importance of the low-rank inductive bias and the optimization-driven designing method for global context modules rather than any speciﬁc MD models.
It is preferred to abstract the MD part as a whole, i.e., M in the context of this paper, and focus on how Hamburger can show the superiority in its entirety.
Vector Quantization Vector Quantization (VQ) (Gray & Neuhoff, 1998), a classic data compression algorithm, can be formulated as an optimization problem in term of matrix decomposition:
[Equation]
where e i is the canonical basis vector, e i
= [0 , · · · , 1 , · · ·
[Equation]
[Equation] i th .
The solution to minimize the objective in Eq.
(8) is K-means (Gray & Neuhoff, 1998).
However, to ensure that VQ is differentiable, we replace the hard arg min and Euclidean distance with softmax and cosine similarity, leading to
Alg. 1, where cosine ( D , X ) is a similarity matrix whose entries satisfy cosine ( D , X ) ij =
d (cid:62) i x j (cid:107) d (cid:107)(cid:107) x (cid:107) , and softmax is applied column-wise and T is the temperature.
Further we can obtain a hard assignment by a one-hot vector when T → 0 .
Algorithm 1 Ham: Soft VQ
Input X .
Initialize D , C .
for k from 1 to K do
C ← softmax ( 1 T cosine ( D , X ))
D ← XC (cid:62) diag ( C 1 n ) −end for
Output ¯ X = DC .
Algorithm 2 Ham: NMF with MU
Input X .
Initialize non-negative D , C for k from 1 to K do
[Equation]
end for
Output ¯ X = DC .
Non-negative Matrix Factorization
If we impose non-negative constraints on the dictionary D and the codes C , it leads to Non-negative Matrix Factorization (NMF) (Lee & Seung, 1999):
[Equation]
To satisfy the non-negative constraints, we add a ReLU non-linearity before putting X into NMF.
We apply the Multiplicative Update (MU) rules (Lee & Seung, 2001) in Alg. 2 to solve NMF, which guarantees the convergence.
As white-box global context modules, VQ, CD, and NMF are straightforward and light, showing remarkable efﬁciency.
They are formulated into optimization algorithms that mainly consist of matrix multiplications with the complexity O ( ndr ) , much cheaper than complexity O ( n 2 d ) in self-attention as r (cid:28) n .
All three MDs are memory-friendly since they avoid generating a large n × n matrix as an intermediate variable, like the product of Q and K of self-attention in Eq.
In the later section, our experiments prove MDs are at least on par with self-attention, though the architectures of M are created by optimization and look different from classic dot product self-attention.
2. O NE - STEP G RADIENT
Since M involves an optimization algorithm as its computational graph, a crux to fuse it into the networks is how the iterative algorithm back-propagates gradient.
The RNN-like behavior of optimization suggests Back-Propagation Through Time (BPTT)
algorithm
(Werbos et al., 1990) as the standard choice to differentiate the iterative process.
We ﬁrst review the BPTT algorithm below.
However, in practice, the unstable gradient from BPTT does harm Hamburger’s performances.
Hence we build an abstract model to analyze the drawbacks of BPTT and try to ﬁnd a pragmatic solution while considering MD’s nature as an optimization algorithm.
As shown in Fig.
2, x , y and h t denote input, output and intermediate result at time step t , respectively, while F and G are operators.
At each time step, the model receives the same input x processed by the underlying networks.
[Equation]
[Figure]
The intermediate results h i are all discarded.
Only the output of the last step h t is passed through G for output y ,
[Equation]
In the BPPT algorithm, the gradient from output y to input x is given, according to the Chain rule:
[Equation]
A thought experiment is to consider t → ∞ , leading to a fully converged result h ∗ and inﬁnite terms in Eq.
We suppose that both F and G are Lipschitz with constants L h w.r.t.
h , L x w.r.t.
x , and L G , and L h < 1 .
Note that these assumptions apply to a large number of optimization or numerical methods.
Then we have:
Proposition 1 { h i } t has linear convergence.
Proposition
[Equation]
Proposition
[Equation]
[Table]
It is easy to incur gradient vanishing w.r.t.
h
when L h is close to 0 and gradient explosion w.r.t.
x when L h is close to 1.
The Jacobian matrix ∂ y ∂ x , moreover, suffers from an ill-conditioned term ( I − ∂ F ∂ h ∗ ) − 1 when the largest eigenvalue of ∂ F ∂ h , i.e., the Lipschitz constant of F w.r.t.
h , approachesand its minimal eigenvalue typically stays near 0, thus restricts the capability of the gradient to search the well-generalized solution in the parameter space.
The erratic scale and spectrum of the gradient back through the optimization algorithm indicate the infeasibility to apply BPTT to Hamburger directly, corroborated by the experiments in Tab. 1, using the same ablation settings as Sec. 3.1.
The analysis inspires us a possible solution.
Note that there are a multiplication of multiple Jacobian matrices ∂ h j ∂
h
j − 1 and a summation of an inﬁnite series in BPTT algorithm , leading to uncontrollable scales of gradients.
It enlightens us to drop some minor terms in the gradient while preserving its dominant terms to ensure the direction is approximately right.
Considering terms of Eq.
(12) as a series, i.e., { ∂ y ∂ h t (cid:16)(cid:81) t − 1 j = t − i ∂ h
[Equation]
x
[Equation]
i , it makes sense to use the ﬁrst term of this series to approximate the gradient if the scale of its terms decays exponentially measured by the operator norm.
The ﬁrst term of the gradient is from the last step of optimization, leading to the one-step gradient, (cid:100)
[Equation]
The one-step gradient is a linear approximation of the BPTT algorithm when t → ∞ according to the Proposition 2.
It is easy to implement, requiring a no _ grad operation in PyTorch (Paszke et al., 2019) or stop _ gradient operation in TensorFlow (Abadi et al., 2016) and reducing the time and space complexity from O ( t ) in BPTT to O (1) .
We test adding more terms to the gradient but its performance is worse than using one step.
According to experimental results, one-step gradient is acceptable to back-propagate gradient through MDs.
[Table]
E XPERIMENTS
In this section we present experimental results demonstrating the techniques described above.
Two vision tasks that beneﬁt a lot from global information and attention mechanism attract us, including semantic segmentation (over 50 papers using attention) and deep generative models like GANs (most state-of-the-art GANs adopt self-attention since SAGAN (Zhang et al., 2019a)).
Both tasks are highly competitive and thus enough for comparing Hamburger with self-attention.
Ablation studies show the importance of MD in Hamburger as well as the necessity of the one-step gradient.
We emphasize the superiority of Hamburger on modeling global context over self-attention regarding both performance and computational cost.
3. A BLATION E XPERIMENTS
We choose to conduct all ablation experiments on the PASCAL VOC dataset (Everingham et al., 2010) for semantic segmentation, and report mIoU of 5 runs on the validation set in the form of best(mean) .
ResNet-50 (He et al., 2016) with output stride 16 is the backbone for all ablation experiments.
We employ a 3 × 3 conv with BN (Ioffe & Szegedy, 2015) and ReLU to reduce channels from 2048 to 512 and then add Hamburger, the same location as popular attentions in semantic segmentation.
For detailed training settings, please see Appendix E.1.
[Table]
[Figure] [Figure]
Breads and Hams
We ablate each part of the
Hamburger.
Removing MD (ham) causes the most severe decay in performance, attesting to the importance of MD.
Even if only the parameter-free MD is added (only ham), the performance can visibly improve.
Parameterization also helps the Hamburger process the extracted features.
Bread, especially upper bread, contributes considerable performance.
Latent Dimension d and r
It is worth noting that there is no simple linear relation between d and r with performances measured by mIoU, though d = 8 r is a satisfactory choice.
Experiments show that even r = 8 performs well, revealing that it can be very cheap for modeling the global context.
Iterations K
We test more optimization steps in the evaluation stage.
In general, the same K for training and test is recommended.
K = 6 is enough for CD and NMF, while even K = 1 is acceptable for VQ.
Typically 3 ∼ 6 steps are enough since simple MD’s prior is still biased, and full convergence can overﬁt it.
The few iterations are cheap and act as early stopping .
3. A C LOSE L OOK AT H AMBURGER
To understand the behavior of Hamburger in the networks, we visualize the spectrums of representations before and after Hamburger on the PASCAL VOC validation set.
The input and output tensors are unfolded to R C × HW .
The accumulative ratio of squared largest r singular values over total squared singular values of the unfolded matrix has been shown in Fig.
A truncated spectrum is usually observed in classic matrix decomposition models’ results due to the low-rank reconstruction.
In the networks, Hamburger also promotes energy concentration while preserving informative details via the skip connection.
Additionally, we visualize the feature maps before and after Hamburger in
Fig.
MD helps Hamburger learn interpretable global information by zeroing out uninformative channels, removing irregular noises, and completing details according to the context.
[Figure] [Figure]
3. A C OMPARISON WITH A TTENTION
This section shows the advantages of MD-based Hamburger over attention-related context modules in computational cost, memory consumption, and inference time.
We compare Hamburger (Ham) with self-attention (SA) (Vaswani et al., 2017), Dual Attention (DA) module from DANet (Fu et al., 2019),
Double Attention module from A 2 Net (Chen et al., 2018b), APC module from APCNet (He et al., 2019b), DM module from DMNet (He et al., 2019a), ACF module from CFNet (Zhang et al., 2019b), reporting parameters and costs of processing a tensor Z ∈ R 1 × 512 × 128 × 128 in Tab. 3.
Excessive memory usage is the key bottleneck of cooperating with attention in real applications.
Hence we also provide the GPU load and inference time on NVIDIA TITAN Xp.
In general, Hamburger is light in computation and memory compared with attention-related global context modules
[Table] [Table]
3. S EMANTIC S EGMENTATION
We benchmark Hamburger on the PASCAL VOC dataset (Everingham et al., 2010), and the PASCAL
Context dataset (Mottaghi et al., 2014), against state-of-the-art attentions.
We use ResNet-101 (He et al., 2016) as our backbone.
The output stride of the backbone is 8.
The segmentation head is the same as ablation experiments.
NMF is usually better than CD and VQ in ablation studies (see Tab. 1).
Therefore, we mainly test NMF in further experiments.
We use HamNet to represent ResNet with
Hamburger in the following section.
Results on the PASCAL VOC test set, and the PASCAL Context validation set, are illustrated in
Tab.
[Table]
and Tab. 5, respectively.
We mark all attention-based models with ∗ in which diverse attentions compose the segmentation heads.
Though semantic segmentation is a saturated task
[Table]
and most contemporary published works have approximate performances, Hamburger shows considerable improvements over previous state-of-the-art attention modules.
[Table]
[Table] [Table]
3. I MAGE G ENERATION
[Table]
Attention presents as the global context description block in deep generative models like
GANs.
Most state-of-the-art GANs for conditional image generation integrate self-attention into their architectures since SAGAN (Zhang et al., 2019a), e.g., BigGAN (Brock et al., 2018),
S 3 GAN (Luˇci´c et al., 2019)
[Table]
and LOGAN (Wu et al., 2019).
It is convincing to benchmark MDbased Hamburger in the challenging conditional image generation task on ImageNet (Deng et al., 2009).
Experiments are conducted to compare Hamburger with self-attention on ImageNet 128 × 128.
Selfattention is replaced by Hamburger with NMF ham in both generator and discriminator at feature resolution 32 × 32, named as HamGAN-baby .
HamGAN achieves an appreciable improvement in
Fr ´ echet Inception Distance (FID) (Heusel et al., 2017) over SAGAN.
Additionally, we compare
Hamburger with a recently developed attention variant Your Local GAN (YLG) (Daras et al., 2020) using their codebase and the same training settings, named HamGAN-strong .
HamGAN-strong offers over 5% improvement in FID while being 15% faster for the total training time and 3.6x faster for the module time (1.54 iters/sec of HamGAN, 1.31 iters/sec of YLG, and 1.65 iters/sec without both context modules, averaged from 1000 iterations) on the same TPUv3 training platform.
R ELATED W ORK
General Survey for Attention
The last ﬁve years have witnessed a roaring success of attention mechanisms (Bahdanau et al., 2015; Mnih et al., 2014; Xu et al., 2015; Luong et al., 2015) in deep learning.
Roughly speaking, the attention mechanism is a term of adaptively generating the targets’ weights to be attended according to the requests.
Its architectures are diverse, and the most well-known one is dot product self-attention (Vaswani et al., 2017).
The attention mechanism has a wide range of applications, from a single source (Lin et al., 2017) to multi-source inputs (Luong et al., 2015; Parikh et al., 2016), from global information discovery (Wang et al., 2018; Zhang et al., 2019a) to local feature extraction (Dai et al., 2017; Parmar et al., 2019).
Previous researchers attempt to explain the effectiveness of attention mechanisms from numerous aspects.
Capturing long-range dependencies (Wang et al., 2018), sequentially decomposing visual scenes (Eslami et al., 2016; Kosiorek et al., 2018), inferring relationships between the part and the whole (Sabour et al., 2017; Hinton et al., 2018), simulating interactions between objects (Greff et al., 2017; van Steenkiste et al., 2018), and learning the dynamics of environments (Goyal et al., 2019) are often considered as the underlying mechanisms of attention.
One common idea from biology is that attention simulates the emergence of concerns in many unconscious contexts (Xu et al., 2015).
Some work tries to interpret the attention mechanism by visualizing or attacking attention weights (Serrano & Smith, 2019; Jain & Wallace, 2019;
Wiegreffe & Pinter, 2019), while others formulate attention into non-local operation (Wang et al., 2018) or diffusion models (Tao et al., 2018; Lu et al., 2019) or build attention-like models via Expectation
Maximization (Greff et al., 2017; Hinton et al., 2018; Li et al., 2019) or Variational Inference (Eslami et al., 2016) on a mixture model.
A connection between transformer and graph neural network is discussed as well (Liang et al., 2018; Zhang et al., 2019c).
Overall, discussions towards attention are still far from reaching agreements or consistent conclusions.
Efﬁcient Attention Recent works develop efﬁcient attention modules via low-rank approximation in both computer vision (Chen et al., 2018b; Zhu et al., 2019; Chen et al., 2019; Li et al., 2019) and natural language processing (Mehta et al., 2019; Katharopoulos et al., 2020;
Wang et al., 2020; Song et al., 2020).
Technically, the low-rank approximation usually targets at the correlation matrix, i.e., the product of Q and K after the softmax operation, using a product of two smaller matrices to approximate the correlation matrix and applying the associative law to save the memory cost and computation, where the approximation involves kernel functions or other similarity functions.
Other works (Babiloni et al., 2020; Ma et al., 2019) make efforts to formulate attention into tensor form but may generate large intermediate variables.
In this paper, we do not approximate attention or make it efﬁcient.
This paper formulates modeling the global context as a low-rank completion problem.
The computation and memory efﬁciency is a by-product of the low-rank assumption on the clean signal subspace and optimization algorithms as architectures.
Matrix Decomposition in Deep Learning There is a long history of combining MD with deep learning.
Researchers focus on reducing the parameters in the networks via factorization on the weights, including the softmax layer (Sainath et al., 2013), the convolutional layer (Zhong et al., 2019), and the embedding layer (Lan et al., 2019).
Tariyal et al.
(2016) attempts to construct deep dictionary learning for feature extraction and trains the model greedily.
This paper tries to factorize the representations to recover a clean signal subspace as the global context and provide a new formulation for modeling the long-range dependencies via matrix decomposition.
C ONCLUSION
This paper studies modeling long-range dependencies in the networks.
We formulate learning the global context as a low-rank completion problem.
Inspired by such a low-rank formulation, we develop the Hamburger module based on well-studied matrix decomposition models.
By specializing matrix decomposition’s objective function, the computational graph created by its optimization algorithm naturally deﬁnes ham, Hamburger’s core architecture.
Hamburger learns interpretable global context via denoising and completing its input and improves the spectrum’s concentration.
It is startling that, when prudently coped with the backward gradient, even simple matrix decomposition proposed 20 years ago is as powerful as self-attention in challenging vision tasks semantic segmentation and image generation, as well as light, fast, and memory-efﬁcient.
We plan to extend Hamburger to natural language processing by integrating positional information and designing a decoder like
Transformer, build a theoretical foundation for the one-step gradient trick or ﬁnd a better method to differentiate MDs, and integrate advanced MDs in the future. A CKNOWLEDGMENTS
Zhouchen Lin is supported by NSF China (grant no.s 61625301 and 61731018), Major Scientiﬁc
Research Project of Zhejiang Lab (grant no.s 2019KB0AC01 and 2019KB0AB02), Beijing Academy of Artiﬁcial Intelligence and Qualcomm.
We thank Google’s Tensorﬂow Research Cloud (TFRC) for providing us Cloud TPUs
A T ABLE OF N OTION
[Table]
B H AMS
Additionally
[Table]
we introduce another type of ham adopted by Hamburger, Concept Decomposition.
Concept Decomposition
We ﬁrst enhance Concept Decomposition (Dhillon & Modha, 2001) to the following form:
[Equation]
This problem has a closed solution w.r
[Table]
t. C under a given D , i.e., C =
( D (cid:62) D + β I ) − 1 D (cid:62) X .
Since D (cid:62) D + β
I is a positive deﬁnite matrix with a regularized conditional number, the inverse can be more numerically stable than the original one where a semi-positive deﬁnite matrix D (cid:62) D is given under β = 0 .
In practice, 0.01 or 0.1 makes no difference for β .
Algorithm 3 Ham: Soft CD
Input X .
Initialize D , C for k from 1 to K do
C ← softmax ( 1 T cosine ( D , X ))
D ← normalize ( XC (cid:62) ) end for
C
[Equation]
Output ¯ X [Equation] DC
[Equation]
The dictionary in CD is given by spherical K-means (Dhillon & Modha, 2001) with objective
Q ( D , X ) , as mentioned in Eq.
[Equation]
The same strategy as VQ is adopted to make the whole algorithm differentiable, however, in which each column of D is normalized to be a unit vector and thus differs from VQ.
C P ROOF OF P ROPOSITIONS
We investigate an abstract RNN model inspired by numerical methods to understand the drawbacks of
BPTT algorithm in differentiating the optimization algorithm of MDs, M .
We show the propositions in Sec.
2.3 to illustrate the unstable gradient from M when using BPTT algorithm, considering MDs’ nature as optimization algorithms.
Proposition 1 The iterations of F have linear convergence.
Proof .
It is obvious that F is a contraction mapping w.r.t.
h under arbitrary given x .
We can then conclude { h t } is a Cauthy sequence and F ( ∗ , x ) admits a unique ﬁxed point h ∗ due to Banach Fixed
Point Theorem.
[Equation]
Eq.
(16) shows the linear convergence.
Proposition
[Equation]
Proof .
Note that F ( ∗ , x ) admits a unique ﬁxed point h ∗ under arbitrary given x , i.e.,
[Equation]
By differentiating the above equation, we can obtain
[Equation]
The Jacobian matrix I − ∂ F ∂ h ∗ is invertible, which implies the existence of the implicit function h ∗ ( x ) .
Immediately, we have
[Equation]
(19) which completes the proof.
Proposition
[Equation]
Proof .
[Equation]
Then we have:
[Equation]
Then we have:
[Equation]
D D ATASETS
PASCAL VOC
The PASCAL VOC dataset (Everingham et al., 2010) is a widely used dataset in both semantic segmentation and detection.
For segmentation, it contains 10,582 images for training, 1,449 images for validation and 1,456 images for testing.
PASCAL VOC dataset involves 20 foreground object classes and a background class for segmentation and detection.
PASCAL Context
The PASCAL Context dataset (Mottaghi et al., 2014) is a challenging dataset in semantic segmentation, which provides detailed labels and involves 59 foreground object classes and a background class for segmentation.
It consists of 4,998 and 5,105 images in training and validation set, respectively.
ILSVRC
The ILSVRC 2012 (ImageNet) (Deng et al., 2009) dataset contains 1.3M training samples and 50k test images, categorized into 1000 object classes.
We resize images to resolution 128 × 128, as done in SNGAN with projection (Miyato & Koyama, 2018) and SAGAN (Zhang et al., 2019a).
E D ETAILS OF E XPERIMENTS E. A BALATION E XPERIMENTS
We use dilated ResNet-50 (He et al., 2016) with the output stride 16 as the backbone.
The backbone is pre-trained on ImageNet (Deng et al., 2009).
We apply a poly-learning rate policy under batch size 12 and 30k iterations (about 35 epochs) for fast experiments (less than 12 hours using 1 NVIDIA
TITAN Xp GPU).
The initial learning rate is set to 0.009, multiplied by (1 − iter iter max ) 0 .
9 after each iteration, with momentum 0.9 and weight decay 0.0001.
Hyperparameters of Hamburger are the same as Appendix E.3.
E. A C OMPARISON WITH A TTENTION M ECHANISM
We report MACs according to Molchanov et al.
(2016), using torchproﬁle 1 , a more accurate proﬁler for Pytorch.
Real-time cost is measured by built-in Pytorch memory tools on NVIDIA TITAN
Xp
GPU with a input tensor Z ∈ R 1 × 512 × 128 × 128 .
Inference times are averaged results from 20 repeats of 100 runs.
E. S EMANTIC S EGMENTATION
Architectures We use ResNet-101 (He et al., 2016) with the ouptput strid 8 as our backbone.
We adopt dilated convolution (Chen et al., 2018a) to preserve more detail spatial information and enlarge receptive ﬁeld as done in the backbone of state-of-the-art attention models (Fu et al., 2019;
Li et al., 2019; Zhang et al., 2019b).
We employ a 3 × 3 convolution layer with BN and ReLU to reduce channels from 2048 to 512 and then add Hamburger on the top of the backbone.
Note that the input of
Hamburger is a tensor Z ∈ R C × H × W .
We unfold Z to a matrix Z ∈ R C × HW and set d z = C and n = HW for Hamburger.
Latent dimension d and r , i.e., the column vectors’ dimension of the input matrix X ∈ R d × n to M and the number of atoms in the dictionary D ∈ R r × d , are set to 512 and 64.
The iterations of MD’s optimization algorithm, K , are set to 6.
Non-negative Matrix Factorization (NMF) is our default ham for semantic segmentation.
Data augmentation In the training stage, we apply random left-right ﬂipping, random scaling (from 0.5 to 2), and cropping to augment the training data.
Images are resized to 513 × 513 for the
PASCAL VOC dataset and the PASCAL Context dataset.
In the test stage, the multi-scale and ﬂipping strategy is applied as other state-of-the-art attention-based models (Fu et al., 2019; Yuan & Wang, 2018; Yuan et al., 2020).
Optimization We use mini-batch SGD with momentum 0.9 to train HamNet.
Synchronized Batch
Normalization is adopted in experiments on semantic segmentation.
All backbones are ﬁne-tuned from ImageNet (Deng et al., 2009) pre-training.
Following previous works (Zhao et al., 2017;
Chen et al., 2018a), we apply a poly-learning rate policy.
The initial learning rate is multiplied by (1 − iter iter max ) 0 .
For the PASCAL VOC dataset, learning rate, weight decay, batch size, iterations are set to 0.009, 0.0001, 16, and 60k, respectively.
We ﬁne-tune HamNet on the PASCAL VOC trainval set with the learning rate down to a tenth.
The learning rate, weight decay, batch size, iterations are 0.002, 0.0001, 16, and 25k for the PASCAL-Context dataset.
E. I MAGE G ENERATION
We use the ofﬁcial GAN codebase 2 from Tensorﬂow (Abadi et al., 2016) and TF-GAN to train
HamGAN and evaluate FID.
Architectures Experiments on ImageNet are conducted using the same architecture as
SAGAN (Zhang et al., 2019a), and YLG (Daras et al., 2020), including Spectral Normalization (Miyato et al., 2018) in both the generator and the discriminator, conditional Batch Normalization in the generator, and class projection in the discriminator (Miyato & Koyama, 2018).
Hamburger with NMF ham is placed at feature resolution 32 × 32 in both the generator and the discriminator where self-attention can obtain the best FID according to Zhang et al. (2019a).
We use d = 8 r for
Hamburger, and d is the same as the input channels, while the optimization steps K are 6.
Restricted to expenditures of training GANs on ImageNet, d , r , and K are decided according to the ablation experiments on semantic segmentation without new ablation experiments.
Optimization For all models, we use Adam (Kingma & Ba, 2015) optimizer with TTUR (Heusel et al., 2017).
HamGAN employs the same training settings as SAGAN (Miyato et al., 2018) and
YLG (Daras et al., 2020), respectively.
Evaluation metrics
The quality of images generated by GANs are evaluated by Fr ´ echet Inception
Distance (FID) (Heusel et al., 2017).
Lower FID indicates that the model can generate higher-ﬁdelity images.
In our experiments, 50k images are sampled from the generator to compute FID.
We evaluate
HamGAN for 6 runs and report the best FID to approximately match the convention in the modern
GAN research like Kurach et al.
and CR-GAN (Zhang et al., 2020), reporting top 5%/15% results in the experiments.
F F URTHER R ESULTS FROM A BLATION E XPERIMENTS
[Table]
Initialization We test four types of initialization for the dictionary D , including ﬁxed initialization, learned initialization, random initialization, and warm start with online update.
Usually, random initialization is the best choice that means we can sample each entry of D from a given distribution like Uniform (0 , 1) as the initialization of the optimization algorithm M .
For NMF, after initializing
D , we initialize C = softmax ( 1 T cosine ( D , X )) since K-means is usually applied for initializing
NMF and this initialization for C is equivalent to a single update in Spherical K-means.
A special reminder is that it is not suitable to initialize either D or C to values too close to 0 due to the property of the MU rule.
So the temperature T is recommended to be a higher value like 1 in this initialization for C .
Random initialization also works for C in NMF with scores 77.8(77.6) when sampling C ij ∼ Uniform (0 , 1) .
Note that learned initialization is always the worst one since the
BPTT algorithm is employed to learn the initialization that the gradient from M may impede the training of the backbone, instead of the one-step gradient.
Warm start beneﬁts MD with unit vectors in the dictionary D like CD.
In general, random initialization is good enough for all three selected
MD models.
A possible reason is that it can enforce the network to adapt to the results solved by different initializations during the training process, acting like an inner augmentation .
[Table]
Temperature T
As we have claimed, when T approaches 0, we can get a solution close to the original problem in both VQ and CD.
In VQ and
CD experiments, a relatively low temperature T is more recommended to solve a better D for
MD.
However, it will not receive more gains but increase the variance during training if we further lower T .
Iterations K
We take the iterations K of optimization algorithms M for all three MD models,
NMF, CD, and VQ, into our consideration.
More iterations and even fully converged results for M are tested in the evaluation stage but worse than little optimization steps.
The smaller K , ranging from 1 to 8, can be treated as early stopping for the optimization algorithm M , obtaining satisfactory performances.
For a detailed visualization, see Fig.
7, Fig. 8, Fig. 9.
[Figure]
[Figure] [Figure]
G A N I NTUITIVE I LLUSTRATION
In this section, we hope to give an example to help our readers develop insight into why the low-rank assumption is useful for modeling the representations’ global context.
The low-rank assumption helps because it represents the inductive bias that the low-level representations contain limited and much less high-level concepts than the scale of the representations themselves .
Imagine an image in which a person works on the road.
Many hyper-pixels extracted by the backbone CNN will describe the road.
Note that the road can be considered as repeats of small road patches, which means that we can represent the road via modeling the basic road patches and repeat them.
Mathematically, it is equivalent to ﬁnding a small set of bases D corresponding to different road patches and a coefﬁcient matrix C that captures the relation between the elementary road patches and the hyper-pixels.
This example illustrates that the high-level concepts, i.e., the global context, can be low-rank in the ideal situation.
The hyper-pixels describing the road patches have close semantic attributes.
However, due to the vanilla CNN’s inefﬁciency for modeling the long-range dependencies, the learned representation contains too many local details and incorrect information, lacking global guidance. Imagine that the person in the image wears gloves.
When we see the gloves patch locally, we think that this patch describes gloves.
When we consider the global context, we can understand that this patch is a part of a person.
The semantic information is hierarchical, depending on at which level we hope to comprehend.
This work aims at enabling the networks to understand the context globally via the low-rank completion formulation.
We thus model the incorrect information, namely the redundancies and incompleteness, as a noise matrix.
To emphasize the global context, we decompose the representations into two parts, a low-rank global information matrix and a noise matrix, by employing the optimization algorithm to recover the clean signal subspace, discard the noises, and enhance the global information via the skip connection.
It could be learned from the data on how much global information the networks need for a speciﬁc task.