Paper: Feature emergence via margin maximization: case studies in algebraic tasks (ICLR-2024)

Reviewer comment (Writing):
I noticed a few times that notations/terms were introduced without proper definitions. Maybe these are well-known in the field, but I think they should be restated for outsiders like me. I also found that some things should be deduced from context and would gain in being stated more explicitly.

Author response:
1. **Notations and Section 3** - See general comments regarding our revisions to section 3. We appreciate the specific comments and have updated section 2 to introduce notation more clearly.

Reviewer comment (Writing):
In Section 2, it would be helpful to state that the output of the neural network is d-dimensional somewhere. Same for the neuron. I am used to having neurons outputting a scalar, but here $\phi$ is $\mathbb{R}^d$-valued, right?

Author response:
1. **Notations and Section 3** - See general comments regarding our revisions to section 3. We appreciate the specific comments and have updated section 2 to introduce notation more clearly.

Reviewer comment (Writing):
Section 3: $\Delta(\mathcal{Y})$ designate the set of functions $f: \mathcal{Y} \rightarrow [0,1]$ such that $\sum_{y \in \mathcal{Y}}f(y)$, right? Might be worth making it more explicit.

Author response:
1. **Notations and Section 3** - See general comments regarding our revisions to section 3. We appreciate the specific comments and have updated section 2 to introduce notation more clearly.

Reviewer comment (Theory):
Proof outline of Theorem 4: The part about the constrained optimization leading to sparsity in Fourier space went over my head. Where does this constraint come from? The norm constraint in the max-margin problem? I understand this is a proof sketch, but I think this could be made more transparent without giving all the details.

Author response:
2. **Proof outline of Theorem 4 (Theorem 5 in the revised version)** - Thanks for the suggestions. The norm constraint indeed comes from the max-margin problem, we now state this in the proof outline.

Reviewer comment (Writing):
$\{u,v,w\} \in \mathbb{R}^d$ is unclear. Shouldn’t it be $u,v,w \in \mathbb{R}^d$?

Author response:
3. **Minor comments** - Indeed, each of $u, v, w \in \mathbb{R}^d$ and we have made this revision. We have also fixed the notation for the expected class-weighted margin in the proof outline of Theorem 4 (Theorem 5 in the revised version) to be consistent with section 3.

Reviewer comment (Writing):
Proof outline of Theorem 4: I thought the notation $\mathbb{E}_{a,b}[\psi’(u,v,w)]$ to be a bit confusing and “too different” from the way it was introduced in Section 3.1.

Author response:
3. **Minor comments** - Indeed, each of $u, v, w \in \mathbb{R}^d$ and we have made this revision. We have also fixed the notation for the expected class-weighted margin in the proof outline of Theorem 4 (Theorem 5 in the revised version) to be consistent with section 3.

Reviewer comment (Theory):
The paper can only handle homogeneous functions as it needs to use the conclusion from [1]. Then, the neural network cannot have the bias terms. This is restricted. For example, Theorem 5 needs $\Omega(2^k)$ number of neurons to learn sparse parity, while a 2-layer ReLU NN with bias only needs $\Omega(poly(k))$ neurons to learn structured sparse parity [2,3]. This weakness may be beyond the scope of this paper but it is important.

Author response:
**Homogeneous architectures** - Yes, our approach can only handle homogeneous architectures, but that still allows us to have biases in the first layer (which are still homogeneous in the weights). The requirement for 2^k neurons comes from the fact that we are using a quadratic activation function. We too hope that our theory can be extended in future work to handle more general regularizations and activation functions.

Reviewer comment (Theory):
As the conclusion section mentioned, the paper needs some properties about homogeneous so that it can only handle $L_{2, \nu}$ norm regularization. It would be good to extend to more general settings.

Author response:
**Generalizing Lemma 3** - Yes, Lemma 3 can be generalized to any homogeneous parametric function having an additive structure, i.e, the output is given by the sum of individual homogeneous functions (which do not share parameters). However, a neural network having more than one hidden layer does not have this additive structure, as some of the hidden layers are shared.

Reviewer comment (Experiments):
If the understanding is correct, it seems like the training is done on all of the data in the support so there is no generalization needed and all analysis is done on the representation. This means it does not address one of the most shocking observations in these arithmetic tasks, namely, grokking.

Author response:
3. **Generalization** - Yes, we do agree that the main focus of the work is on representation learning, and not generalization. We do not claim to be providing an explanation for grokking, which is an independently interesting phenomenon.

Reviewer comment (Theory):
There is no optimization analysis, so the whole picture is not complete (although the converged result matches the prediction).

Author response:
4. **Optimization** - Yes, we agree that we have only provided the maximum margin result, which when combined with the result of Wei et al, shows that the global optimum of the training loss is indeed the maximum margin solution for vanishing regularization. We do not prove that the global optimum will be reached by training; as you note, we provide empirical evidence instead.

Reviewer comment (Writing):
What is the relationship between $u,v,w$ and $\omega$? Concatenation?

Author response:
5. **Relation between $\omega$ and $u,v,w$** - Yes $\omega$ is the concatenation of $u,v,w$.

Reviewer comment (Writing):
What is $P(D)$ in Sec 3? Is it an arbitrary distribution?

Author response:
6. **Definition of P(D)** - P(D) is the set of distributions on the samples belonging to the training dataset.

Reviewer comment (Writing):
Where does $\tau(x,y)$ in Sec 3 come from?

Author response:
7. **Intuition about $\tau(x,y)$** - Please look at the updated section 3, where we provide an intuition about $\tau(x,y)$. This is required to convert the network-level search problem to a neuron-level search for multi-class classification.

Reviewer comment (Theory):
Unclear understanding of C1. How can the support of the distribution be the subset of argmin which is presumably a single point or a set of discrete points ($\text{sqt}(q^*)\subseteq \text{argmin}_{(x,y) \in D} g(\theta^*, x, y)$)? Does that mean $q^*$ is a discrete distribution on the $\mathbb{R}^n$?

Author response:
8. **Support of q** - Yes, $q^*$ is discrete as it can only have support on the points in the training dataset $D$.

Reviewer comment (Theory):
One impression from reading this paper is that the provided construction actually seems to be the only sensible solution given the constraints of the function class. Is this a correct assessment? This is important because it is somewhat similar to saying that the answer is already assumed by making all the assumptions, which undermines these frameworks' ability to help understand real LLMs. If not, what other circuits are possible beyond the constructions in the paper? How do we know this is not the only way to do it?

Author response:
9. **Other solutions** - Regarding whether the provided construction is the only sensible solution—please look at Appendix D in the paper. We have provided another construction which does not display any Fourier sparsity; so there are other solutions possible, but margin maximization favors the Fourier-sparse solutions. We have added a reference to it after the informal theorem.