The ubiquity of 2-homogeneity, how its implicit bias selects features, and other stories

Authors: Matus Telgarsky, Danny Hyun Ho Son
Primary area: optimization
Keywords: generalization, optimization, margins, deep learning theory

Abstract:
This work studies the optimization and generalization consequences of a seemingly
  innocuous design choice in many modern architectures:
  they end with a composition of affine parameters belonging to a normalization layer and a linear layer,
  resulting in a fundamentally $2$-homogeneous architecture.
  The first set of results are abstract, showing how any architecture satisfying this type of 2-homogeneity and a
  few regularity conditions on the gradients of the inner layers obtain large margins
  and low test error. As technical byproducts, this part of the story provides an implicitly
  biased gradient flow guarantee and also a nondecreasing margin lemma for inhomogeneous networks. The second set of results instantiate this framework for shallow normalized ReLU networks,
  establishing large margin and low test error via feature selection
  purely from random initialization and standard
  gradient flow. As a corollary, the paper obtains good test error for $k$-bit parity problems,
  in particular passing
  below sample complexity lower bounds from linearized analyses such as the  Neural Tangent Kernel.

Decision: reject
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729445028712

thread_type:
official_review

<thread_object>
creator: reviewer_gwqr

<content>
summary:
The paper studies the gradient flow dynamics on networks where the last layer includes normalization followed by a linear transformation. Formally, it means that they consider networks where the last layers are a two-layer linear diagonal network, but the motivation for studying this architecture comes from normalization layers. They show implicit bias towards ell_1 margin maximization of the linear transformation induced by the last layer (with its normalization). This result includes abstract assumptions, but they also consider two-layer normalized ReLU networks as a concrete instantiation, obtain generalization guarantees, and use it to achieve a generalization guarantee for sparse parities which goes beyond the NTK bounds.

strengths:
The paper considers the challenging problem of analyzing the implicit bias of non-homogenous networks. They give a general margin-maximization guarantee (under certain assumptions) and show that this general result can be used to obtain more concrete bounds for two-layer normalized ReLU networks and specifically for learning sparse parities. Unlike many existing results on the implicit bias in the rich regime, their result gives a bound on the margin rather than just showing convergence to a KKT point of the max-margin problem. The sparse parities learning problem has attracted much interest, and several prior works considered learning sparse parities using shallow neural networks, and thus the corollary (Example 4.1) is of independent interest. The technique seems interesting and non-trivial, and might also be of independent interest.

weaknesses:
Theorem 3.1 is very hard to parse. It includes assumptions that involve the margin induced by the random features at initialization, and the KL divergence between the ell_1 max margin predictor at initialization and the last layer’s initialization. It also includes an assumption called “the slow growing derivatives property” which might be hard to verify. So, it’s hard for me to determine to what extent Theorem 3.1 can be applied in different settings. The good news is that the author/authors prove that it can be applied to sufficiently wide two-layer normalized ReLU networks.

Typos:
- Title is missing
- Line 433: xi_t —> \xi_t
- Line 479: “of of”

questions:
Can you compare your result on two-layer networks to Theorem 3.3 in Telgarsky 2022?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730699970057

thread_type:
official_review

<thread_object>
creator: reviewer_xxtj

<content>
summary:
The paper describes the dynamics of a network with a 2-homogeneous head and establish the margin guarantee for the models in late training phases. Under a simpler setting of adding 2-homogeneous head to a one-hidden-layer network, the paper can further establish generalization guarantees based on the margin guarantees.

strengths:
The paper establishes theorem 3.6, a generalization of result from previous work in a clear form. The paper has clear writing that exhibits the proof ideas.

weaknesses:
1. Theorem 3.1 & 4.1 has assumptions that are hard to verify, and I believe these assumptions should be given out more explicitly in the texts. For theorem 3.1, as the line 260 RHS bounds depends on $t$, the choice of $t$ depends on $R$ and $n$, and the LHS quantity also depends on $n, R$ and $t$, it is essentially hard to tell whether these assumptions can be satisfied. For theorem 4.1, $m$ depends on $t$, $t$ depends on $R$ and $R$ may depend on $m$ as $m$ changes the initialization scheme. Thus the authors should at least provide the scale estimate of $R$ to show that their assumption is reasonable.   

2. The proof technique for Theorem 3.6 is standard and may be lacking of novelty compared to former work (https://arxiv.org/abs/1906.05890) .

3. The correlation of part1 and part2 of theorem 1.1 are unclear from text. It seems like two independent results. The main claim of the current paper - selecting feature from space $F_1$ is more efficient than from $F_2$ using 2-homogeneous head - does not require a result in representation learning, thus is basically already inferred from previous work (Telgarsky 2022 + Lyu 2019, as mentioned in Example 4.1). Also it is unclear from text why 2-homogeneity is unique - from my perspective $k$-homogeneity may also be provable using similar arguments for $k>1$.

Furthermore, the title of the paper should be changed in the pdf file.

questions:
Regarding weakness 1, can the authors provide an instance of choices of the parameters such that the assumptions in theorem 3.1&4.1 are satisfied? And please provide more clarifications on the strength of these assumptions?

Also can the author further clarify weakness 2 & 3?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730773517698

thread_type:
official_review

<thread_object>
creator: reviewer_pv2i

<content>
summary:
This paper shows how any architecture satisfying 2-homogeineity and a few regularity conditions on the gradients of the inner layers obtain large margins and low test error. Instantiating this framework to shallow normalized ReLU networks, the paper gives an end-to-end result showing that it obtains large margin and low text error. As a corollary, the paper obtains good test error for k-bit parity problems.

strengths:
1. The instantiation of the framework to shallow normalized ReLU network is novel and interesting. This improved upon results in [Chizat and Bach 2020] to be quantitative, though the computational complexity may be double exponential. 

2. For result has concrete implications to the support-only k-parity problem, demonstrating neural networks could learn support-only k-parity much more sample-efficient than kernel methods.

weaknesses:
The first set of results (Theorem 1.1) are a bit abstract and contain non-varifiable assumptions. It is not clear whether it is useful beyond the more concrete shallow normalized ReLU network.

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1731117702314

thread_type:
official_review

<thread_object>
creator: reviewer_nazb

<content>
summary:
This paper proposes an expansion of theoretical results about homogeneous networks to those that only *end* in a 2-homogeneous network (with inhomogeneous layers before that). This includes the majority of common present-day architectures, which end in the combination of a normalization layer with learnable scales and shifts by another linear layer. Under certain assumptions, the paper uses a correspondence of gradient descent on these parameters to mirror descent in a re-parameterized space, showing that gradient flow yields large margins, for a certain definition of margin mostly aligning with prior work.

strengths:
The problem addressed by the paper is important. Identifying that the combination of a normalization layer with a linear layer at the end of these networks might be important to their optimization seems like it might be an important observation. The framing of the particular approach with the connection to mirror descent is clever.

Unfortunately I did not have time to carefully follow the proofs in the supplementary material, but the proof sketches in the main body are clear and the approaches seem reasonable to me.

weaknesses:
I found some of the initial discussion a little misleading in that e.g. the discussion around line 140 made me think that the analysis would apply to arbitrary homogeneous subsets of parameters $\hat w$. While this is true for the margin definition given there, almost all of the paper applies only to the doubled-linear-layer structure of equation (2). This is fine; I just found that bit of discussion confusing, particularly since taking e.g. only the last layer would also yield a 1-homogeneous structure, and I found myself wondering if it would really be true that it would increase the margin defined in terms of each possible homogeneous part of the network. (To be clear, I don't think any statements in the paper are wrong, I just found this brief discussion confusing.)

Much more importantly: if the feature map $F$ were frozen, the setting of this paper would become a "diagonal network" of depth two, for which the implicit bias of gradient descent was already determined by [Gunasekar et al. (NeurIPS 2018)](https://arxiv.org/abs/1806.00468) to be, indeed, an l1 minimizer. The difference with the present paper's results is that $F$ is not required to stay frozen, and of course a frozen $F$ would be a substantial difference from actual training. The setting of the present paper, though, requires that although $F$ can change, it must satisfy "slow-growing derivatives" and "initial good features." Together, this seems to allow the optimization on the last two layers to basically pick out the good features before they change too much to no longer be good. This is interesting, but perhaps not as interesting as it might seem without citing the (here-uncited) previous work for frozen $F$, and with the caveats about initial good features only coming several pages into the paper.

The paper also implicitly argues that the 2-homogeneous structure at the end of transformer-like architectures is very important to their optimization geometry. To be more convincing on this point, since the theoretical results are admittedly a "sanity check," it would be interesting to also run some experiments at least on toy problems to see if this seems to line up. That is, do people leave the learnable parameters on in their normalization layers that are immediately next to a linear layer because it actually helps, or just because they didn't think about it?

questions:
- In the case of frozen $F$, would direct application of the results of Lyu and Li plus Gunasekar et al. give roughly the same results as you obtain?

- Is my characterization of "slow-growing derivatives" + "initial good features" -> "the features can't break before last-layer optimization finds the good ones" correct?

- I'm curious about the relationship of your "initial good features" assumption to the "initial alignment" results from the series of papers by Emmanuel Abbe and coauthors, particularly [Abbe et al. (ICML 2022)](https://arxiv.org/abs/2202.12846). It seems, without thinking about it very much, that your "initial good features" assumption implies a high "initial alignment"; as they show, this is in some sense necessary for learning, but in a stronger form than they make.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732714092447

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Revision Summary

comment:
We thank the reviewers for their comments and time. In order to improve clarity,
we have revised the paper and the main changes are as follows.

* We provided a more thorough explanation of the uniqueness of 2 homogeneity 
following the statement of lemma 3.2. As a summary, only a 2 homogeneous head induces 
a reparameterized flow that explicitly optimizes the superior $F_1$ space as opposed 
to the $F_2$ space.

* In response to concerns regarding the novelty of theorem 3.6, 
we revised the discussion following theorem 3.6 by providing a more in depth comparison with https://arxiv.org/pdf/1906.05890 with regards to generality of the result as
well as proof techniques.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper considers the implicit bias of neural networks where the last two layers form a two-layer linear diagonal network; such networks can be motivated from networks having normalization layers with post-normalization diagonal scaling parameters in their last hidden layer. Under technical assumptions that can be roughly described as "slow-growing derivatives" and "initial good features,” the paper proves that the network achieves a good normalized margin early in training, and then the smoothed margin is non-decreasing until the end of training. The first results are given for an abstract family of networks (Section 3), while Section 4 presents a more concrete instantiation of the theorems under a 2-layer normalized ReLU network and an additional assumption on the data distribution.

The paper brings up interesting viewpoints, focusing on the common structure of the last two layers in networks employing normalization layers. The main results are developed for a general framework and they can be specialized to more concrete results for two-layer normalized ReLU networks and for sparse parity problems. Many reviewers, and myself, found the proof technique novel and interesting. The technique to view the evolution of $a \odot b$ as mirror flow on $u$ and using online convex optimization to handle the time-dependency of $V(s)$ is intriguing.

However, the main criticism from all reviewers came from the fact that the main result (Theorem 3.1) relies on complicated and hard-to-verify assumptions and conditions. The “initial good features” assumption requires that the initialized parameter $V(0)$ already yields linearly separable hidden representations $F(x_i;V(0))$. The “constants” $C_1$ and $C_2$ are in fact dependent on the entire trajectory up to time $t$, which makes it impossible to verify a priori. With $t$ growing double exponentially, the conditions on $\\lVert \\bar{p} \\rVert_\infty$ and $C_2$ also become restrictive. Although the authors demonstrate an example in which the theorem can be applied, the application requires unrealistically large width and training time (albeit not infinite). These strong and hard-to-verify conditions make me worry if the results will really be useful/applicable in other setups.

Also, it looks to me that the presentation has some room for improvement, including:
- As noted, simplifying or removing the assumptions and conditions in the main theorem will improve the paper.
- More concrete examples beyond 2-layer normalized ReLU networks will be helpful.
- As Reviewer NazB suggested, the authors could consider running some (toy) experiments corroborating the theory.
- The paper looks difficult to follow if the reader doesn’t have prior knowledge on $F_1$ and $F_2$. They are not formally defined in the paper, except at the very end (in Line 500). Some formal description of these classes should be given.
- The paper claims that the normalized ReLU networks can be viewed as an idealization of BatchNorm, but I question why this network bears resemblance to BatchNorm. There is no batch-wise normalization and the input is not centered at zero.
- The authors claim regarding Theorem 3.6 that any model $h$ that has a subset of homogeneous parameters satisfies the theorem, but one needs to also mention the condition $\sup_\theta \sup_{\\lVert \hat w \\rVert \leq 1} \\lVert \nabla_{\hat w} h(x; (\hat w,\theta)) \\rVert < \infty$ because the claim can be misleading when stated without the condition.
- As noted by Reviewer NazB, the implicit bias of linear diagonal networks on binary classification was studied in Gunasekar et al (2018). Of course, the paper goes beyond the setting of this work, but the authors must contextualize the current submission with respect the this existing result and its subsequent follow-ups.

Overall, although the review scores were borderline positive, I believe the paper sits slightly below the acceptance threshold and requires some major revision before getting published. At this time, I recommend rejection.