Deep Networks Learn Features From Local Discontinuities in the Label Function

Authors: Prithaj Banerjee, Harish Guruprasad Ramaswamy, Mahesh Lorik Yadav, Chandra Shekar Lakshminarayanan
Primary area: interpretability and explainable AI
Keywords: Deep Learning, Feature learning, Interpretable, Local Discontinuities, Deep learning theory, Deep neural architectures, Supervised learning

Abstract:
Deep neural networks outperform kernel machines on several datasets due to feature learning that happens during gradient descent training. In this paper, we analyze the mechanism through which feature learning happens and use a notion of features that corresponds to discontinuities in the true label function. We hypothesize that the core feature learning mechanism is label function discontinuities attracting model function discontinuities during training. To test this hypothesis, we perform experiments on classification data where the true label function is given by an oblique decision tree. This setup allows easy enumeration of label function discontinuities, while still remaining intractable for static kernel/linear methods. We then design/construct a novel deep architecture called a Deep Linearly Gated Network (DLGN), whose discontinuities in the input space can be easily enumerated.  In this setup, we provide supporting evidence demonstrating the movement of model function discontinuities towards the label function discontinuities during training. The easy enumerability of discontinuities in the DLGN also enables greater mechanistic interpretability. We demonstrate this by extracting the parameters of a high-accuracy decision tree from the parameters of a DLGN. We also show that the DLGN is competitive with ReLU networks and other tree-learning algorithms on several real-world tabular datasets.

Decision: accept (poster)
Reviewer scores: 5, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730167423870

thread_type:
official_review

<thread_object>
creator: reviewer_xq4v

<content>
summary:
This paper proposes a theory to explain how feature learning happens in neural networks. The authors posit that neural networks align (or rather, get attracted to) the discontinuites in the way the target label changes over the input domain. As a pilot test for this theory, the authors consider a setting where the target function to learn is an Oblique Decision Tree (ODT). ODTs correspond to decision trees where each internal node is a linear threshold. Thus, the hyperplanes associated to the nodes in the decision tree split up the input space into different labeled regions (see Figure 1). The hyperplane corresponding to the root is "most discontinuous" in terms of how the label changes on both sides of it (owing to further and further splits by its many descendants). The authors posit that the training of non-linear neural networks procedurally aligns the model with these discontinuities.

As a tractable model to further empirically study this theory, the authors propose a novel neural network architecture which they term Deep Linearly Gated Networks (DLGNs). DLGNs are somewhere in between standard nonlinear (e.g., ReLU activated) deep networks and deep linear networks. The function computed by a DLGN can be writen as a large summation of terms of the form $f_\pi \cdot g_\pi$. Here, $f_\pi$ is a product of the signs of activations of neurons in a deep linear network (a single neuron from each layer). $g_\pi$ is a product of the activations themselves (a slight detail is that the weight matrices producing the activations in $f_\pi$ and $g_\pi$ are different). Thus, we can think of $f_\pi$ as an indicator of the intersection of halfspaces, while $g_\pi$ is the weight we add up if the indicator turns out to be true.

The authors fix the target function to be an ODT, and train the DLGN architecture on synthetic data labelled by the ODT. The suprising observation is given in Figure 2. At the end of training, if one plots the hyperplanes corresponding to the linear threshold at every neuron, one observes that the hyperplanes (at least in the later layers of the DLGN) align very well with the hyperplane splits in the ODT! Furthermore, Table 1 shows that most of these linear thresholds align with some hyperplane in the ODT. That is, the neurons in the DLGN architecture are "getting attracted" to aligning with some internal node in the ODT.

The authors next use this empirical observation to extract a decision tree out of a trained DLGN. Namely, they plot the linear thresholds corresopnding to all the activations after training, and then cluster these thresholds. The center of the largest cluster is chosen to be a root node hyperplane, and then we recurse this procedure on data on either side of this hyperplane. This procedure is pictorally well illustrated in Figure 3. Again, one can see that on performing such clustering-based decision tree generation, the first cluster center aligns very well with the root node hyperplane, and the phenomenon continues as we recurse. The upshot is that one can extract an interpretable model from the DLGN..

Finally, the authors perform experiments on real-world classification data to illustrate that the classification accuracy of DLGNs is somewhere in between standard ReLU networks and other simpler non-neural network algorithms. Thus, DLGNs have the added benefit of being interpretable, but more powerful than stanard decision trees.

strengths:
The proposed theory is intriguing and indeed very interesting. The plots in Figures 2 and 3 are indeed striking---they illustrate how faithfully the DLGN neuron activations are increasingly aligning with the ODT hyperplanes over the process of training. These observations support the theory put forth by the authors about neural network architectures possibly picking up on the discontinuities in the labelling function. The clustering procedure to extract a decision tree is also very interesting, and empirically seems to work well (as suggested by Figure 3). Overall, I find the theory proposed by the authors, along with the striking empirical illustrations, very interesting. These could motivate further theoretical investigations for such phenomena.

weaknesses:
The pilot experiments done by the authors are admittedly specialized. Namely, they don't actually consider standard ReLU networks, but only the DLGNs they propose. Furthermore, they only consider cases where the target function is conveniently an ODT. While this is totally okay as an initial starting point, it does raise the question about whether such empirical phenomena of the neurons aligning with the discontinuities also arise in cases where the target function has discontinuities of a different nature (like curvy discontinuities, etc). But this is not a significant weakness, as it seems beyond the scope of a pilot study. But it would be really interesting to visualize the activations of the neurons (say with quadratic thresholds) for when the target function is also composed of curvy discontinuities.

questions:
1) Why do you even introduce a manifold $M$ in your notation in line 133? It is never used 
2) On line 192, I don't think I agree that "no other hyperplane other than the internal nodes has this property". e.g. consider any other hyperplane that is not one of the internal nodes in the plot in Figure 1; doesn't it also have points of both labels on either side of it?
3) You introduce this formal notion of $\gamma(R, f)$ as the local discontinuity coefficient in Section 3, but then never mention it anywhere later in the paper. How would you explain the results in say Figure 2 and 3 in the context of this quantity? Could you elaborate on this a bit?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730585721954

thread_type:
official_review

<thread_object>
creator: reviewer_n82d

<content>
summary:
In this paper, the authors propose a mechanism to intuitively explain why neural networks can surpass kernel methods through feature learning. They hypothesize that the feature learning process involves the alignment of model function discontinuities with label function discontinuities during training. To explore this, they introduce a new network architecture called the Deep Linearly Gated Network (DLGN), designed as a surrogate for ReLU networks. They argue that this architecture retains similarities to ReLU networks while offering easier interpretability. Under this framework, they provide empirical evidence showing how model function discontinuities move toward label function discontinuities during training, facilitated by feature learning.

strengths:
This paper is a bold attempt to deepen our understanding of feature learning in deep learning. The hypothesis, data setup, network architecture, and approach to interpretability are unconventional, and they bring a fresh perspective to the literature. This work has the potential to inspire new ideas and serve as a valuable starting point for further exploration.

weaknesses:
From my perspective, while this work is intriguing, it is not yet ready for formal publication. My concerns are as follows:

1. The authors reference previous studies examining the dynamics of single hidden layer models under specialized data and settings to push beyond kernel methods or deep linear models (Damian et al., 2022; Ba et al., 2022). They appropriately note that these analyses, often focused on specific data settings like the parity function, fall short of addressing the needs or behaviors of deeper networks. However, numerous works also investigate complex feature learning with deep neural networks, such as https://arxiv.org/abs/2305.06986 and https://arxiv.org/pdf/2311.13774. These studies should be acknowledged and compared with the current work.

2. I am skeptical about the extent to which the new architecture resembles a ReLU network. Additionally, it is unclear how this intuition or design could extend to CNNs or transformers. Besides, the current version only applies to binary classification, whereas a universal feature learning mechanism should ideally apply across setups, such as binary classification, multi-class classification, and regression. I am unsure how the proposed intuition extends to these broader contexts. For example, existing papers (such as the two mentioned above) demonstrate that neural networks can efficiently learn $h = g \circ p$ with $p$ quadratic and $g$ nonlinear via feature learning. How would this intuition explain such cases?

3. In Section 3, the introduction to ODT progresses too quickly. The paper would benefit from a more mathematically detailed introduction to the new concepts presented in this section.

questions:
See the weakness part.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730642398948

thread_type:
official_review

<thread_object>
creator: reviewer_d8rp

<content>
summary:
The authors introduce a model called the Deep Linearly Gated Network (DLGN) to study feature learning, specifically in binary classification tasks defined by an oblique decision tree labeling function. They use DLGN to test the hypothesis that during training, the model’s discontinuities move towards label function's discontinuities. The paper includes evaluations on dozens of open tabular datasets to compare DLGN with ReLU networks and tree-learning algorithms.

strengths:
1. Clear writing. The authors study a specific class of problems with great clarity. 
2. The authors' persistence in tackling a challenging yet manageable problem setting is commendable.
3. Like a black-box learner, the DLGN is able to learn non-linear features. Yet, it still provides mechanistic interpretability.
4. DLGN outperforms both tree-based and non-tree algorithms, as well as ReLU networks, in the oblique decision tree setting, while maintaining strong competitiveness on real-world tabular datasets.
5. The authors provide a framework that paves the way for future research and development.

weaknesses:
I'm not seeing effective weaknesses.

questions:
1. What's the purpose of defining the manifold $\mathcal{M}$ in line 133?
2. In line 239, Equation (4), is there a missing transpose on $\mathbf{u}_{i_1}^1$?
3. I find it difficult to understand how the computational cost of a forward pass for Equation (1) is less than twice that of a ReLU network with $mL$ nodes. Could the authors provide further clarification on this?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730651372328

thread_type:
official_review

<thread_object>
creator: reviewer_tlyw

<content>
summary:
This paper proposes a novel model architecture (deep linearly gated network, DLGN) and studies the way in which this model seeks out “label discontinuities” in the data during training. This analysis is enabled the fact that we can enumerate such label discontinuities in a DLGN. The synthetic datasets with known discontinuities are generated by another model — an oblique decision tree (ODT). The authors also show how an ODT can be constructed from a trained DLGN, for the purpose of interpretability. Finally, the paper presents results of fitting a DLGN to several UCI regression tasks, comparing performance to several tree-based, kernel-based, and NN-based baselines.

strengths:
- Originality. The DLGN is an interesting architecture that combines deep linear networks with the gating mechanism to construct a novel class of non-linear models. One could also treat DLGN as a novel decision tree parametrization which, when relaxed using a sigmoid, can be learned by back-propagation. To the best of my knowledge, DLGN is a novel, original model architecture, though its connection to soft decision trees should be studied more carefully.

weaknesses:
- Lack of focus. At the moment, the paper’s focus is split between a study of feature learning using DLGN (chapter 5), and a study of the DLGN itself as an expressive, yet interpretable model architecture (chapter 6). To me these are two orthogonal contributions, and the paper would be stronger if authors focused on one of these.
- Significance. The significance of the proposed model and the feature learning study presented in the paper is not clear to me. 
  1. Capturing “label discontinuities” (are these not simply decision boundaries?) is at the core of solving a classification problem, hence it is not surprising that a model which works well on the task has to discover such hyperplanes — I don’t see an alternative way that a model can solve a task. The real question is how an over-parametrized model can correctly identify high-dimensional decision boundaries given limited data — this is the main mystery in the theory of deep learning at the moment, and one that this paper doesn’t shed much light on.
  2. To understand the significance of the findings for deep learning, we would need to understand the relationship between a DLGN and a DNN. While it is nice (albeit, not surprising) to see how a DLGN uncovers the true “label discontinuities”, how can we know that a DNN will demonstrate the same behavior?
  3. While the proposed architecture is appealing due to potentially being both expressive and interpretable, results on the real (UCI) datasets suggests that DLGN is comparable in performance to standard tree algorithms. There is little evidence that we should prefer the proposed architecture to e.g. the well-studied random forests, which we could also argue to be “interpretable”. (I do not consider results on synthetic data to be good evidence, given the connection between ODTs used to generate the data and DLGN.)
  4. While the authors claim that the proposed architecture is interpretable, no interpretations of the models fit to real data are given. If interpretability if the main selling point of the architecture, I would expect a deeper analysis focused on interpretability.
- Presentation. Please consider moving the DLGN model diagram to the main text: it’s difficult to understand the model architecture from the formulas alone. Also, please try to stick to academic language, and avoid informal phrases like “well nigh impossible”, “handily outperform”, “succeed comfortably”, “about the same”, etc.

questions:
- Authors suggest that they use a sigmoid instead of an indicator function to make the DLGN differentiable for training. Have authors considered using a temperature parameter for the sigmoid (potentially annealed during training), as common in other continuous relaxation methods?
- On lines 406-407 authors suggest that they use DBScan for clustering hyperplanes due to this algorithm being “robust to outliers” — why do authors expect outliers to be a significant issue here?
- How do authors explain the observation that certain layers of the DLGN are more prone to matching the true discontinuities than others  (Figure 2)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732692988075

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of changes to the revised version

comment:
We thank the reviewers for the detailed comments. We have made the changes asked for by the reviewers in the revised version. Here is a summary of the changes:

1. Put appendix and main paper in the same pdf.
2. Reduced some of the informal terminology in the main paper.
3. Fixed some errors in the figure in Appendix A.1 and made it more legible.
4. Added minor clarification to the local discontinuity coefficients and added a figure and discussion of the coefficient in Appendix A.4
5. Added proof for DLGN computation, i.e. equivalence of equation 1 and 5 in Appendix A.5
6. Added a short section in the appendix detailing the relation of DLGN to ReLU networks in Appendix A.6
7. Added a new section in the Appendix (A.7) discussing some deep learning phenomena and the candidate answers given by the feature learning narrative developed in the paper.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732713222520

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Feature Learning Narrative Directions from DLGN. (Added in Appendix A7 of revision)

comment:
We believe that the feature learning insights drawn from DLGN are applicable in general. More importantly, we believe that it can shed light on several phenomena that remain mysteries with deep networks. Some of the mysteries and candidate answers motivated from this paper are given below.

**Q1: Is neural network training using gradient descent parameter efficient?**

A1: No. The greedy nature of gradient descent, pushes multiple interconnected parts towards the same feature when it is potentially not required. e.g. In Table 4 the distances of learned DLGN hyperplanes (with width m = 100) to the ODT hyperplanes are given. While the root node hyperplane is discovered separately by multiple gating hyperplanes, some of the other internal nodes are left in the lurch. This causes poor generalisation. Potentially, this narrative can be exploited to make gradient descent less greedy.

**Q2: How does gradient descent discover the label function discontinuities in high dimensional input space?**

A2: The full labelling function may be far from being a linear separator, but when data is restricted within a small enough scope, one of the manifolds making up the decision boundary could very well be a good classifier (this argument uses the insight that even in high dimensional data, if it is linearly separable, picking up the separator can be done in a compute and data efficient manner). This causes that manifold to be picked up by one of the components of the deep network, and thereby making the rest of the learning problem easier and kickstarting a virtuous cycle.

**Q3: Why does neural network pruning enable better learning of smaller architectures than training the smaller architecture from scratch?**

A3: We assume effective learning happens only in problems where the number of label function discontinuities is small. Significant parts of the neural network are simply not necessary to represent the discontinuities, which can be done by a small fraction of the trained deep network model. A large width network is still better for training, because (say) doubling the number of neurons per layer increases the number of paths by a factor of $2^L$. This increases the chances of some path picking a right scope during training enabling the learning of the appropriate label function discontinuity. Once the learning is complete, a significant fraction of the network which were unlucky to not have gotten a good scope can be removed.



**Q4: What is the role of layers in deep networks? Is increasing the number of layers always beneficial?**

A4: The main role of layers is to give context/scope to other layers. For example, with a depth 4 ODT labelling function, a DLGN with 3 or lesser number of hidden layers would not be able to give a good scope to any neuron. A DLGN with 5 or more layers is just unnecessary. This is reflected in our experiments as well, where depth 4 DLGN performed best for a depth 4 ODT labelling function.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes the deep linearly gated network to analyze feature learning in deep neural networks. It highlights how label function discontinuities guide model function discontinuities during training, using a novel architecture for interpretable results. Contributions include mechanistic insights into feature learning, a direct link to decision tree extraction, and competitive performance on synthetic and real-world datasets.

The novelty of the concept proposed by this paper is highly appreciated. The clear presentation also adds to its contribution. Further extensions and more rigorous experimental validation are noted, but the paper is judged to be in satisfactory form at this time.