Title: ENERGY-BASED SPHERICAL SPARSE CODING

Abstract: In this paper, we explore an efficient variant of convolutional sparse coding with unit norm code vectors where reconstruction quality is evaluated using an inner product (cosine distance). To use these codes for discriminative classification, we describe a model we term Energy-Based Spherical Sparse Coding (EB-SSC) in which the hypothesized class label introduces a learned linear bias into the coding step. We evaluate and visualize performance of stacking this encoder to make a deep layered model for image classification.

1. INTRODUCTION:

Sparse coding has been widely studied as a representation for images, audio and other vectorial data. This has been a highly successful method that has found its way into many applications, from signal compression and denoising (Donoho, 2006; Elad & Aharon, 2006) to image classification (Wright et al., 2009) , to modeling neuronal receptive fields in visual cortex (Olshausen & Field, 1997) . Since its introduction, subsequent works have brought sparse coding into the supervised learning setting by introducing classification loss terms to the original formulation to encourage features that are not only able to reconstruct the original signal but are also discriminative (Jiang et al., 2011; Yang et al., 2010; Zeiler et al., 2010; Ji et al., 2011; Zhou et al., 2012; Zhang et al., 2013) .

While supervised sparse coding methods have been shown to find more discriminative features leading to improved classification performance over their unsupervised counterparts, they have received much less attention in recent years and have been eclipsed by simpler feed-forward architectures. This is in part because sparse coding is computationally expensive. Convex formulations of sparse coding typically consist of a minimization problem over an objective that includes a least-squares (LSQ) reconstruction error term plus a sparsity inducing regularizer.

Because there is no closed-form solution to this formulation, various iterative optimization techniques are generally used to find a solution (Zeiler et al., 2010; Bristow et al., 2013; Yang et al., 2013; Heide et al., 2015) . In applications where an approximate solution suffices, there is work that learns non-linear predictors to estimate sparse codes rather than solve the objective more directly (Gregor & LeCun, 2010) . The computational overhead for iterative schemes becomes quite significant when training discriminative models due to the demand of processing many training examples necessary for good performance, and so sparse coding has fallen out of favor by not being able to keep up with simpler non-iterative coding methods.

In this paper we introduce an alternate formulation of sparse coding using unit length codes and a reconstruction loss based on the cosine similarity. Optimal sparse codes in this model can be computed in a non-iterative fashion and the coding objective lends itself naturally to embedding in a discriminative, energy-based classifier which we term energy-based spherical sparse coding (EB-SSC). This bi-directional coding method incorporates both top-down and bottom-up information where the features representation depends on both a hypothesized class label and the input signal. Like Cao et al. (2015) , our motivation for bi-directional coding comes from the "Biased Competition Theory", which suggests that visual processing can be biased by other mental processes (e.g., topdown influence) to prioritize certain features that are most relevant to current task. Fig. 1 illustrates the flow of computation used by our SSC and EB-SSC building blocks compared to a standard feed-forward layer.

Our energy based approach for combining top-down and bottom-up information is closely tied to the ideas of Larochelle & Bengio (2008) 

1.1. NOTATION:

Matrices are denoted as uppercase bold (e.g., A), vectors are lowercase bold (e.g., a), and scalars are lowercase (e.g., a). We denote the transpose operator with , the element-wise multiplication operator with , the convolution operator with * , and the cross-correlation operator with . For vectors where we dropped the subscript k (e.g., d and z), we refer to a super vector with K components stacked together (e.g., z = [z 1 , . . . , z K ] ).

2. ENERGY-BASED SPHERICAL SPARSE CODING:

Energy-based models capture dependencies between variables using an energy function that measure the compatibility of the configuration of variables (LeCun et al., 2006) . To measure the compatibility between the top-down and bottom-up information, we define the energy function of EB-SSC to be the sum of bottom-up coding term and a top-down classification term:

E(x, y, z) = E code (x, z) + E class (y, z).

(1) The bottom-up information (input signal x) and the top-down information (class label y) are tied together by a latent feature map z.

2.1. BOTTOM-UP RECONSTRUCTION:

To measure the compatibility between the input signal x and the latent feature maps z, we introduce a novel variant of sparse coding that is amenable to efficient feed-forward optimization. While the idea behind this variant can be applied to either patch-based or convolutional sparse coding, we specifically use the convolutional variant that shares the burden of coding an image among nearby overlapping dictionary elements. Using such a shift-invariant approach avoids the need to learn dictionary elements which are simply translated copies of each other, freeing up resources to discover more diverse and specific filters (see Kavukcuoglu et al. (2010) ).

Convolutional sparse coding (CSC) attempts to find a set of dictionary elements {d 1 , . . . , d K } and corresponding sparse codes {z 1 , . . . , z K } so that the resulting reconstruction, r = K k=1 d k * z k accurately represents the input signal x. This is traditionally framed as a least-squares minimization with a sparsity inducing prior on z:

arg min z x - K k=1 d k * z k 2 2 + β z 1 .

(2)

Unlike standard feed-forward CNN models that convolve the input signal x with the filters, this energy function corresponds to a generative model where the latent feature maps {z 1 , . . . , z K } are convolved with the filters and compared to the input signal (Bristow et al., 2013; Heide et al., 2015; Zeiler et al., 2010) .

To motivate our novel variant of CSC, consider expanding the squared reconstruction error x -

r 2 2 = x 2 2 -2x r + r 2 2 .

If we constrain the reconstruction r to have unit norm, the reconstruction error depends entirely on the inner product between x and r and is equivalent to the cosine similarity (up to additive and multiplicative constants). This suggests the closely related unit-length reconstruction problem:

arg max z x K k=1 d k * z k -β z 1 (3) s.t. K k=1 d k * z k 2 ≤ 1

In Appendix A we show that, given an optimal unit length reconstruction r * with corresponding codes z * , the solution to the least squares reconstruction problem (Eq. 2) can be computed by a simple scaling r * = (x r * -β 2 z * 1 )r * . The unit-length reconstruction problem is no easier than the original least-squares optimization due to the constraint on the reconstruction which couples the codes for different filters. Instead consider a simplified constraint on z which we refer to as spherical sparse coding (SSC):

EQUATION

In 2.3 below, we show that the solution to this problem can be found very efficiently without requiring iterative optimization.

This problem is a relaxation of convolutional sparse coding since it ignores non-orthogonal interactions between the dictionary elements1 . Alternately, assuming unit norm dictionary elements, the code norm constraint can be used to upper-bound the reconstruction length. We have by the triangle and Young's inequality that:

k d k * z k 2 ≤ k d k * z k 2 ≤ k d k 1 z k 1 ≤ D k z k 2 (5)

where the factor D is the dimension of z k and arises from switching from the 1-norm to the 2-norm.

EQUATION

However, this relaxation is very loose, primarily due to the triangle inequality. Except in special cases (e.g., if the dictionary elements have disjoint spectra) the SSC codes will be quite different from the standard least-squares reconstruction.

2.2. TOP-DOWN CLASSIFICATION:

To measure the compatibility between the class label y and the latent feature maps z, we use a set of one-vs-all linear classifiers. To provide more flexibility, we generalize this by splitting the code vector into positive and negative components:

z k = z + k + z - k z + k ≥ 0 z - k ≤ 0

and allow the linear classifier to operate on each component separately. We express the classifier score for a hypothesized class label y by:

EQUATION

)

The classifier thus is parameterized by a pair of weight vectors (w + yk and w - yk ) for each class label y and k-th channel of the latent feature map.

This splitting, sometimes referred to as full-wave rectification, is useful since a dictionary element and its negative do not necessarily have opposite visual semantics. This splitting also allows the classifier the flexibility to assign distinct meanings or alternately be completely invariant to contrast reversal depending on the problem domain. For example, Shang et al. (2016) found CNN models with ReLU non-linearities which discard the negative activations tend to learn pairs of filters which are related by negation. Keeping both positive and negative responses allowed them to halve the number of dictionary elements.

We note that it is also straightforward to introduce spatial average pooling prior to classification by introducing a fixed linear operator P used to pool the codes (e.g., w + y Pz + k ). This is motivated by a variety of hand-engineered feature extractors and sparse coding models, such as Ren & Ramanan (2013) , which use spatially pooled histograms of sparse codes for classification. This fixed pooling can be viewed as a form of regularization on the linear classifier which enforces shared weights over spatial blocks of the latent feature map. Splitting is also quite important to prevent information loss when performing additive pooling since positive and negative components of z k can cancel each other out.

2.3. CODING:

Bottom-up reconstruction and top-down classification each provide half of the story, coupled by the latent feature maps. For a given input x and hypothesized class y, we would like to find the optimal activations z that maximize the joint energy function E(x, y, z). This requires solving the following optimization:

EQUATION

where x ∈ R D is an image and y ∈ Y is a class hypothesis. z k ∈ R F is the k-th component latent variable being inferred; z + k and z - k are the positive and negative coefficients of z k , such that  and w - yk ∈ R F are the dictionary filter, positive coefficient classifier, and negative coefficient classifier for the k-th component respectively. A key aspect of our formulation is that the optimal codes can be found very efficiently in closedform-in a feed-forward manner (see Appendix B for a detailed argument).

z k = z + k + z - k . The parameters d k ∈ R M , w + yk ∈ R F ,

2.3.1. ASYMMETRIC SHRINKAGE:

To describe the coding processes, let us first define a generalized version of the shrinkage function commonly used in sparse coding. Our asymmetric shrinkage is parameterized by upper and lower thresholds - Fig. 2 shows a visualization of this function which generalizes the standard shrinkage proximal operator by allowing for the positive and negative thresholds. In particular, it corresponds to the proximal operator for a version of the 1 -norm that penalizes the positive and negative components with different weights

β -≤ β + shrink (β + ,β -) (v) =    v -β + if v -β + > 0 0 otherwise v + β - if v + β -< 0 (9) (a) -β -≤ 0 ≤ β + (b) 0 ≤ -β -≤ β + (c) -β -≤ β + ≤ 0 (d) β -≤ 0 ≤ -β +

|v| asym = β + v + 1 + β -v - 1 .

The standard shrink operator corresponds to shrink (β,-β) (v) while the rectified linear unit common in CNNs is given by a limiting case shrink (0,-∞) (v). We note that -β -≤ β + is required for shrink (β + ,β -) to be a proper function (see Fig. 2 ).

2.3.2. FEED-FORWARD CODING:

We now describe how codes can be computed in a simple feed-forward pass. Let

EQUATION

be vectors of positive and negative biases whose entries are associated with a spatial location in the feature map k for class y. The optimal code z can be computed in three sequential steps:

1. Cross-correlate the data with the filterbank d k x

2. Apply an asymmetric version of the standard shrinkage operator

EQUATION

where, with abuse of notation, we allow the shrinkage function (Eq. 9) to apply entries in the vectors of threshold parameter pairs β + yk , β - yk to the corresponding elements of the argument.

3.. Project onto the feasible set of unit length codes:

EQUATION

)

2.3.3. RELATIONSHIP TO CNNS::

We note that this formulation of coding has a close connection to single layer convolutional neural network (CNN). A typical CNN layer consists of convolution with a filterbank followed by a nonlinear activation such as a rectified linear unit (ReLU). ReLUs can be viewed as another way of inducing sparsity, but rather than coring the values around zero like the shrink function, ReLU truncates negative values. On the other hand, the asymmetric shrink function can be viewed as the sum of two ReLUs applied to appropriately biased inputs:

shrink (β + ,β -) (x) = ReLU(x -β + ) -ReLU(-(x + β -)),

SSC coding can thus be seen as a CNN in which the ReLU activation has been replaced with shrinkage followed by a global normalization.

3. LEARNING:

We formulate supervised learning using the softmax log-loss that maximizes the energy for the true class label y i while minimizing energy of incorrect labels ȳ.

arg min

EQUATION

)

where α is the hyperparameter regularizing w + y , w - y , and d. We constrain the relationship between β and the entries of w + y and w - y in order for the asymmetric shrinkage to be a proper function (see Sec. 2.3.1 and Appendix B for details).

In classical sparse coding, it is typical to constrain the 2 -norm of each dictionary filter to unit length. Our spherical coding objective behaves similarly. For any optimal code z * , there is a 1-dimensional subspace of parameters for which z * is optimal given by scaling d inversely to w, β. For simplicity of the implementation, we opt to regularize d to assure a unique solution. However, as Tygert et al. (2015) point out, it may be advantageous from the perspective of optimization to explicitly constrain the norm of the filter bank.

Note that unlike classical sparse coding, where β is a hyperparameter that is usually set using crossvalidation, we treat it as a parameter of the model that is learned to maximize performance.

3.1. OPTIMIZATION:

In order to solve Eq. 13, we explicitly formulate our model as a directed-acyclic-graph (DAG) neural network with shared weights, where the forward-pass computes the sparse code vectors and the backward-pass updates the parameter weights. We optimize the objective using stochastic gradient descent (SGD).

As mentioned in Sec. 2.3 shrinkage function is assymetric with parameters β + yk or β - yk as defined in Eq. 10. However, the inequality constraint on their relationship to keep the shrinkage function a proper function is difficult to enforce when optimizing with SGD. Instead, we introduce a central offset parameter and reduce the ordering constraint to pair of positivity constraints. Let

ŵ+ yk = β + yk -b k ŵ- yk = β - yk + b k (14)

be the modified linear "classifiers" relative to the central offset b k . It is straightforward to see that if β + yk and β - yk that satisfy the constrain in Eq. 13, then adding the same value to both sides of the inequality will not change that. However, taking b k to be a midpoint between them, then both β + yk -b k and β - yk + b k will be strictly non-negative. Using this variable substitution, we rewrite the energy function (Eq. 1) as

EQUATION

)

where b is constant offset for each code channel. The modified linear "classification" terms now take on a dual role of inducing sparsity and measuring the compatibility between z and y.

This yields a modified learning objective that can easily be solved with existing implementations for learning convolutional neural nets:

EQUATION

where ŵ+ and ŵare the new sparsity inducing classifiers, and b are the arbitrary origin points. In particular, adding the K origin points allows us to enforce the constraint by simply projecting ŵ+ and ŵonto the positive orthant during SGD.

3.1.1. STACKING BLOCKS:

We also examine stacking multiple blocks of our energy function in order to build a hierarchical representation. As mentioned in Sec. 3.1.1, the optimal codes can be computed in a simple feedforward pass-this applies to shallow versions of our model. When stacking multiple blocks of our energy-based model, solving for the optimal codes cannot be done in a feed-forward pass since the codes for different blocks are coupled (bilinearly) in the joint objective. Instead, we can proceed in an iterative manner, performing block-coordinate descent by repeatedly passing up and down the hierarchy updating the codes. In this section we investigate the trade-off between the number of passes used to find the optimal codes for the stacked model and classification performance.

For this purpose, we train multiple instances of a 2-block version of our energy-based model that differ in the number of iterations used when solving for the codes. For recurrent networks such as this, inference is commonly implemented by "unrolling" the network, where the parts of the network structure are repeated with parameters shared across these repeated parts to mimic an iterative algorithm that stops at a fixed number of iterations rather than at some convergence criteria. In Fig. 3 , we compare the performance between models that were unrolled zero to four times. We see that there is a difference in performance based on how many sweeps of the variables are made. In terms of the training objective, more unrolling produces models that have lower objective values with convergence after only a few passes. In terms of testing error, however, we see that full code inference is not necessarily better, as unrolling once or twice has lower errors than unrolling three or four times. The biggest difference was between not unrolling and unrolling once, where both the training objective and testing error goes down. The testing error decreases from 0.0131 to 0.0074. While there is a clear benefit in terms of performance for unrolling at least once, there is also a trade-off between performance and computational resource, especially for deeper models.

4. EXPERIMENTS:

We evaluate the benefits of combining top-down and bottom-up information to produce classspecific features on the CIFAR-10 ( Krizhevsky & Hinton, 2009) dataset using a deep version of our EB-SSC. All experiments were performed using MatConvNet (Vedaldi & Lenc, 2015) framework with the ADAM optimizer (Kingma & Ba, 2014) . The data was preprocessed and augmented following the procedure in Goodfellow et al. (2013) . Specifically, the data was made zero mean and whitened, augmented with horizontal flips (with a 0.5 probability) and random cropping. No weight decay was used, but we used a dropout rate of 0.3 before every convolution layer except for the first. For these experiments we consider a single forward pass (no unrolling).

Base Network block kernel, stride, padding activation conv1

3

× 3 × 3 × 96, 1, 1 ReLU/CReLU conv2 3 × 3 × 96/192 × 96, 1, 1 ReLU/CReLU pool1 3 × 3, 2, 1 max conv3 3 × 3 × 96/192 × 192, 1, 1 ReLU/CReLU conv4 3 × 3 × 192/384 × 192, 1, 1 ReLU/CReLU conv5 3 × 3 × 192/384 × 192, 1, 1 ReLU/CReLU pool2 3 × 3, 2, 1 max conv6 3 × 3 × 192/384 × 192, 1, 1 ReLU/CReLU conv7 1 × 1 × 192/384 × 192, 1, 1 ReLU/CReLU

Table 1: Underlying block architecture common across all models we evaluated. SSC networks add an extra normalization layer after the non-linearity. And EB-SSC networks insert class-specific bias layers between the convolution layer and the non-linearity. Concatenated ReLU (CReLU) splits positive and negative activations into two separate channels rather than discarding the negative component as in the standard ReLU.

4.1. CLASSIFICATION:

We compare our proposed EB-SSC model to that of Springenberg et al. (2015) , which uses rectified linear units (ReLU) as its non-linearity. This model can be viewed as a basic feed-forward version of our proposed model which we take as a baseline. We also consider variants of the baseline model that utilize a subset of architectural features of our proposed model (e.g., concatenated rectified linear units (CReLU) and spherical normalization (SN)) to understand how subtle design changes of the network architecture affects performance.

We describe the model architecture in terms of the feature extractor and classifier. Table 1 shows the overall network architecture of feature extractors, which consist of seven convolution blocks and two pooling layers. We test two possible classifiers: a simple linear classifier (LC) and our energy-based classifier (EBC), and use softmax-loss for all models. For linear classifiers, a numerical subscript indicates which of the seven conv blocks of the feature extractor is used for classification (e.g., LC 7 indicates the activations out of the last conv block is fed into the linear classifier). For energy-based classifiers, a numerical subscript indicates which conv blocks of the feature extractor are replace with a energy-based classifier (e.g., EBC 6-7 indicates the activations out of conv5 is fed into the energy-based classifier and the energy-based classifier has a similar architecture to the conv blocks it replaces). The notation differ because for energy-based classifiers, the optimal activations are a function of the hypothesized class label, whereas for linear classifiers, they are not. (Shang et al., 2016) , and to intermediate variants. The baseline models all perform very similarly with some small reductions in error rates over the baseline CReLU+LC 7 . However, CReLU+LC 7 reduces the error rate over ReLU+LC 7 by more than one percent (from 11.40% to 10.17%), which confirms the claims by Shang et al. (2016) and demonstrates the benefits of splitting positive and negative activations. Likewise, we see further decrease in the error rate (to 9.74%) from using spherical normalization. Though normalizing the activations doesn't add any capacity to the model, this improved performance is likely because scale-invariant activations makes training easier. On the other hand, further sparsifying the activations yielded no benefit. We tested values β = {0.001, 0.01} and found 0.001 to perform better. Replacing the linear classifier with our energy-based classifier further decreases the error rate by another half percent (to 9.23%).

Model:



4.2. DECODING CLASS-SPECIFIC CODES:

A unique aspect of our model is that it is generative in the sense that each layer is explicitly trying to encode the activation pattern in the prior layer. Similar to the work on deconvolutional networks built on least-squares sparse coding (Zeiler et al., 2010) , we can synthesize input images from activations in our spherical coding network by performing repeated deconvolutions (transposed convolutions) back through the network. Since our model is energy based, we can further examine how the topdown information of a hypothesized class effects the intermediate activations. The first column in Fig. 4 visualizes reconstructions of a given input image based on activations from different layers of the model by convolution transpose. In this case we put in zeros for class biases (i.e., no top-down) and are able to recover high fidelity reconstructions of the input. In the remaining columns, we use the same deconvolution pass to construct input space representations of the learned classifier biases. At low levels of the feature hierarchy, these biases are spatially smooth since the receptive fields are small and there is little spatial invariance capture in the activations. At higher levels these class-conditional bias fields become more tightly localized.

Finally, in Fig. 5 we shows decodings from the conv2 and conv5 layer of the EB-SSC model for a given input under different class hypotheses. Here we subtract out the contribution of the top-down bias term in order to isolate the effect of the class conditioning on the encoding of input features.

As visible in the figure, the modulation of the activations focused around particular regions of the image and the differences across class hypotheses becomes more pronounced at higher layers of the network.

5. CONCLUSION:

We presented an energy-based sparse coding method that efficiently combines cosine similarity, convolutional sparse coding, and linear classification. Our model shows a clear mathematical connection between the activation functions used in CNNs to introduce sparsity and our cosine similarity convolutional sparse coding formulation. Our proposed model outperforms the baseline model and we show which attributes of our model contributes most to the increase in performance. We also demonstrate that our proposed model provides an interesting framework to probe the effects of class-specific coding.

:

We note that our formulation is also closely related to the dynamical model suggested byRozell et al. (2008), but without the dictionary-dependent lateral inhibition between feature maps. Lateral inhibition can solve the unit-length reconstruction formulation of standard sparse coding but requires iterative optimization.

Figures and Tables:

Figure 1: Building blocks for coding networks explored in this paper. Our coding model uses non-linearities that are closely related to the standard ReLU activation function. (a) Keeping both positive and negative activations provides a baseline feed-forward model termed concatenated ReLU (CReLU). (b) Our spherical sparse coding layer has a similar structure but with an extra bias and normalization step. Our proposed energy-based model uses (c) energy-based spherical sparse coding (EB-SSC) blocks that produces sparse activations which are not only positive and negative, but are class-specific. These blocks can be stacked to build deeper architectures.

Figure 2: Comparing the behavior of asymmetric shrinkage for different settings of β + and β -. (a)-(c) satisfy the condition that -β -≤ β + while (d) does not.

Figure 3: Comparing the effects of unrolling a 2-block version of our energy-based model. (Best viewed in color.)

Figure 4: The reconstruction of an airplane image from different levels of the network (rows) across different hypothesized class labels (columns). The first column is pure reconstruction, i.e., unbiased by a hypothesized class label, the remaining columns show reconstructions of the learned class bias at each layer for one of ten possible CIFAR-10 class labels. (Best viewed in color.)

Comparison of the baseline ReLU+LC 7 model, its derivative models, and our proposed model on CIFAR-10. The results shown in Table 2 compare our proposed model to the baselines ReLU+LC 7 (Springenberg et al., 2015) and CReLU+LC 7

Table content: <table><tr><td/><td colspan="3">Train Err. (%) Test Err. (%) # params</td></tr><tr><td>ReLU+LC 7</td><td>1.20</td><td>11.40</td><td>1.3M</td></tr><tr><td>CReLU+LC 7</td><td>2.09</td><td>10.17</td><td>2.6M</td></tr><tr><td>CReLU(SN)+LC 7</td><td>0.99</td><td>9.74</td><td>2.6M</td></tr><tr><td>SSC+LC 7</td><td>0.99</td><td>9.77</td><td>2.6M</td></tr><tr><td>SSC+EBC 6-7</td><td>0.21</td><td>9.23</td><td>3.2M</td></tr></table>

Peer review (iclr2017):
First, I*d like to thank the authors for their answers and clarifications. I find, the presentation of the multi-stage version of the model much clearer now. Pros: + The paper states a sparse coding problem using cosine loss, which allows to solve the problem in a single pass. + The energy-based formulation allows bi-directional coding that incorporates top-down and bottom-up information in the feature extraction process. Cons: + The cost of running the evaluation could be large in the multi-class setting, rendering the approach less attractive and the computational cost comparable to recurrent architectures. + While the model is competitive and improves over the baseline, the paper would be more convincing with other comparisons (see text). The experimental evaluation is limited (a single database and a single baseline) ------ The motivation of the sparse coding scheme is to perform inference in a feed forward manner. This property does not hold in the multi stage setting, thus optimization would be required (as clarified by the authors). Having an efficient way of performing a bi-directional coding scheme is very interesting. As the authors clarified, this could not necessarily be the case, as the model needs to be evaluated many times for performing a classification. Maybe an interesting combination would be to run the model without any class-specific bias, and evaluation only the top K predictions with the energy-based setting. Having said this, it would be good to include a discussion (if not direct comparisons) of the trade-offs of using a model as the one proposed by Cao et al. Eg. computational costs, performance. Using the bidirectional coding only on the top layers seems reasonable: one can get a good low level representation in a class agnostic way. This, however could be studied in more detail, for instance showing empirically the trade offs. If I understand correctly, now only one setting is being reported. Finally, the authors mention that one benefit of using the architecture derived from the proposed coding method is the spherical normalization scheme, which can lead to smoother optimization dynamics. Does the baseline (or model) use batch-normalization? If not, seems relevant to test. Minor comments: I find figure 2 (d) confusing. I would not plot this setting as it does not lead to a function (as the authors state in the text).