Understanding Nonlinear Implicit Bias via Region Counts in Input Space

We extend our gratitude to the reviewer for the comments and suggestions. Below, we address the primary concerns that have been raised.

Q1: Including Table 5 in Appendix B to the main text.

A1: Thank you for the suggestion. In the next version, we will move Table 5 from Appendix B to the main text to highlight the robustness of our conclusions to the number of sampling trials, further validating the method’s reliability.

Q2: It is necessary to show that a large learning rate or small batch size leads to better generalization.

A2: Extensive prior research [1][2][3] has demonstrated that large learning rates and small batch sizes typically lead to better generalization, both in terms of generalization gap and top-1 accuracy. However, this is not the primary focus of our paper. Our study aims to explore the relationship between model generalization and region count, providing a novel perspective that complements existing findings.

Q3: The range of generalization gap appears large compared to typical CIFAR-10 performance.

A3: We appreciate the reviewer’s observation. To clarify, data augmentation is not used during the training of the neural networks in our study, whereas [4][5] applied data augmentation when training ResNet models, likely contributing to their higher test accuracy. In Section 7, we also conduct experiments with data augmentation. As shown in Figure 7, with optimal hyperparameters and data augmentation, the test accuracy reaches approximately 93%. Importantly, our correlation between region count and generalization gap remains consistent regardless of whether data augmentation is applied.

Q4: Does region count maintain a consistent correlation after PCA?

A4: We thank the reviewer for suggesting this direction. In our paper, the subspace is generated using convex combinations of training data, allowing each point on the subspace to be expressed as convex combinations of the training datapoint coordinates. These coordinates can then be fed into the neural network for label prediction, enabling us to compute the region count based on the classification of labels across the subspace.

However, when using PCA to generate a subspace, we can determine the coordinates of points within the reduced-dimensional space but cannot map these points back to their original coordinates in the input space. Without the original coordinates, it is not possible to input these points into the neural network for label prediction. Therefore, we have not yet identified a direct method for calculating the region count on a PCA subspace.

Q5: Dose the variation in the dimensionality $R^d$ affects the results of sampled subspace in region count?

A5: The region count increases with subspace dimensionality. However, as shown in Table 2, the correlation between region count and generalization remains consistently high, unaffected by dimensionality changes.

Q6: Is this paper first to understand the correlation between generalization gap and hyper-parameters?

A6: This paper does not directly analyze the correlation between generalization gap and hyperparameters. Instead, it studies the relationship between the generalization gap and region count, as well as region count and hyperparameters.

Q7: How is 'Connectedness' validated in practice?

A7: Algorithm 1 in Appendix B details the computation of region count. The input space is divided into smaller grids, and Breadth-First Search (BFS) is used to traverse adjacent grids with the same label, identifying and counting connected components.

Q8: It is unclear how to interpret the results of the random flip.

A8: Figures 6 and 7 present observations rather than interpretations, showing that mixup and data augmentation affect region count in different ways. These observations suggest that the mechanisms through which mixup and data augmentation improve model generalization may differ. Further exploration of this phenomenon is a promising direction for future research.

Q9: In Appendix B, the third paragraph mentions Table 4 twice.

A9: Thank you for catching this typo. It has been corrected in the revised version.

We thank the reviewer once again for the valuable and helpful suggestions.

References

[1] Keskar, Nitish Shirish, et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).

[2] Jastrzębski, Stanisław, et al. "Three factors influencing minima in sgd." arXiv preprint arXiv:1711.04623 (2017).

[3] Hoffer, Elad, Itay Hubara, and Daniel Soudry. "Train longer, generalize better: closing the generalization gap in large batch training of neural networks." Advances in neural information processing systems 30 (2017).

[4] He, Kaiming, et al. "Deep residual learning for image recognition." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

[5] https://github.com/kuangliu/pytorch-cifar

Thank you to the authors for their efforts. Based on their responses, I have decided to maintain my score. I hope the authors can incorporate these additional points into the final version of the paper.

We thank the reviewer for the comments and constructive suggestions. In the following, we address the main concern raised. Please find the details below.

Q1: Starting with the seminal work of Montufar which also studies the impact of number of linear regions on performances and its tie with the architecture. Discuss the difference.

A1: In the "Related Work" section, specifically under "Region Counts of Neural Networks," we discuss prior works on linear regions and have already clarified the distinction between linear regions and our definition of region count.

Linear regions, as defined in prior works [1][2][3][4][5], refer to the set of inputs corresponding to the same activation pattern in the network. In contrast, our definition of decision regions refers to connected areas in the input space that correspond to the same label. Unlike linear regions, which are independent of the labels output by the neural network and primarily focus on the network's representation capability, our approach is more closely tied to the network's generalization ability. Therefore, the analysis of linear regions fundamentally differs from our work.

Q2: Discuss the relationship with the work of [6][7][8].

A2: We thank the reviewer for suggesting these papers [6][7][8], and we have carefully read them.

Paper [6] introduces a theoretical framework representing deep networks as max-affine spline operators (MASOs) via spline functions and operators. It proposes a regularization term based on this theory, improving classification performance. Paper [7] uses spline partition geometry to characterize and analyze the geometry of neural network decision boundaries, while paper [8] investigates training dynamics, such as grokking and delayed robustness, through the number of spline regions.

These works provide an interesting perspective on partitioning the input space of neural networks using geometric properties, specifically leveraging the continuous piecewise linear (CPWL) activation functions. Their method divides the input space into linear regions determined by hyperplanes derived from each network layer, reflecting the geometric structure of the network rather than directly correlating with its output labels.

In contrast, our method partitions the input space based explicitly on the network's predicted labels, ensuring that each region corresponds to a specific label. This distinction ties our approach more directly to the network’s labeling complexity and allows us to focus on its correlation with generalization ability. We will incorporate these points into the discussion of related work in the revised paper.

Finally, we thank the reviewer once again for the effort in providing us with valuable suggestions. We will continue to provide clarifications if the reviewer has any further questions.

References

[1] Montufar, Guido F., et al. "On the number of linear regions of deep neural networks." Advances in neural information processing systems 27 (2014).

[2] Xiong, Huan, et al. "On the number of linear regions of convolutional neural networks." International Conference on Machine Learning. PMLR, 2020.

[3] Hanin, Boris, and David Rolnick. "Deep relu networks have surprisingly few activation patterns." Advances in neural information processing systems 32 (2019).

[4] Hanin, Boris, and David Rolnick. "Complexity of linear regions in deep networks." International Conference on Machine Learning. PMLR, 2019.

[5] Serra, Thiago, Christian Tjandraatmadja, and Srikumar Ramalingam. "Bounding and counting linear regions of deep neural networks." International conference on machine learning. PMLR, 2018.

[6] Balestriero, Randall. "A spline theory of deep learning." International Conference on Machine Learning. PMLR, 2018.

[7] Humayun, Ahmed Imtiaz, et al. "Splinecam: Exact visualization and characterization of deep network geometry and decision boundaries." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023.

[8] Humayun, Ahmed Imtiaz, Randall Balestriero, and Richard Baraniuk. "Deep networks always grok and here is why." arXiv preprint arXiv:2402.15555 (2024).

Thanks again for your valuable feedback! Could you please let us know whether your concerns have been addressed? We are happy to make further updates if you have any other questions or suggestions.

We thank the reviewer for the time spent on reviewing our work and for the very detailed comments. We add experiments comparing against more metrics such as sharpness-based and Pac-Bayesian based. Although our work aims to identify an interpretable implicit bias for nonlinear models rather than predicting the generalization gap, we are happy to add more empirical results if that could help strengthen our findings. Please find the details below.

Q1: Link to Certified Radius [5] and Generalization Measures.

A1: We have carefully reviewed the paper on Certified Radius [5]. The margin-to-Lipschitz constant proposed in these works is primarily designed to measure a neural network's robustness against adversarial perturbations, rather than its generalization ability. Adversarial robustness aims to reduce vulnerability to intentionally crafted small input perturbations, while generalization focuses on the network's performance on unseen data. Since the input distributions in these two scenarios differ significantly, there is no direct connection between the two concepts.

Q2: How does the computational cost of estimating region count scale with model size and dimensionality, especially in NLP domains?

A2: In Appendix B, we outline the method for calculating region count, which has a time complexity that does not scale with model size but instead grows exponentially with the dimension of the selected subspace. Our paper focuses on classification tasks, where regions are defined based on label consistency. Extending the definition of regions to NLP tasks remains an open challenge.

To explore the applicability of our approach to transformer-based models, we conduct an experiment using Vision Transformers (ViT) on CIFAR-10. The results show a strong correlation of 0.84 between region count and the generalization gap, reinforcing the validity of our measure in this setting. Details of the hyperparameters used are summarized below:

Q3: Can the theoretical framework be feasibly extended to deeper or more sophisticated architectures?

A3: We appreciate the reviewer’s insightful question. Computing region count for more complex architectures is indeed a challenging task, as it cannot be directly bounded using activation numbers. At present, we have not identified a feasible method for theoretical analysis in these cases. We consider this an important avenue for future research.

Q4: Could region count metrics be adapted to other types of tasks, such as regression or structured prediction?

A4: We thank the reviewer for this thoughtful suggestion. To adapt region count metrics to regression tasks, a key step would be to define connectedness. One potential approach is to partition the regression error into small intervals and analyze the connectivity of input space points falling within each interval. We believe this is a promising direction for future exploration.

Q5: Beyond learning rate and batch size, how do other hyperparameters—like weight decay, optimizer choice, and normalization layers affect region count and generalization?

A5: We thank the reviewer for this insightful question. Our preliminary experiments indicate that the relationships between these hyperparameters and region count are highly complex, making it challenging to capture them systematically within the scope of the current work. For this reason, we did not include detailed analyses in the main text. We consider this an important area for future research.

Q6: How does the generalization potential of region counts compare with other metrics such as Lipschitz continuity and flatness?

A6: We thank the reviewer for this question. To compare the effectiveness of region count with other metrics, we conduct additional experiments on CIFAR-10 using ResNet-18. We measured the correlation between the generalization gap and spectral norm [1], PAC-Bayesian flatness [2], and region count:

The results demonstrate that region count exhibits a significantly stronger correlation with the generalization gap.

Q7: High learning rates (linked to weight decay effects with normalization) and small batch sizes are already associated with implicit regularization. The paper suggests that this is explained through sharp minima, but a more precise explanation could clarify this link.

A7: Numerous prior studies [6][7][8] have investigated the effects of learning rate and batch size on generalization, showing that large learning rates and small batch sizes often lead to better generalization. We acknowledge that our explanation could be more precise, and we will include additional clarification and discussion on this topic in the next version of the paper.

Q8: Discuss the relationship between region count and certified radius[5], which uses the ratio of margin and Lipschitz constant as a robustness measure.

A8: We thank the reviewer for this insightful question. Region count and certified radius are indeed related, as both quantify changes in a neural network's predictions within the input space. The key difference lies in their scope: region count is a global measure, capturing the network's overall prediction behavior across the entire input space, while certified radius is a local measure, focusing on the extent of perturbation required to alter a prediction. A promising direction for future work could involve exploring whether certified radius provides an upper bound for region count.

Q9: Would a comparison with sharpness metrics provide valuable insights?

A9: We thank the reviewer for the question. To address this, we conduct additional experiments on CIFAR-10 using ResNet-18, comparing sharpness metrics from PAC-Bayesian bounds, including those using the origin and initialization as reference tensors (PB-I and PB-O), as well as PAC-Bayesian Magnitude-aware Perturbation Bounds (PB-M-I and PB-M-O) [2][3][4]. The results are summarized below:

The results demonstrate that region count shows a significantly stronger correlation with the generalization gap compared to these sharpness metrics.

Q10: Does the slicing method scale effectively to larger models?

A10: In our paper, we already explored the correlation across neural networks with different parameter sizes, finding it to be largely consistent. To further address this question, we conduct additional experiments using ResNet models with varying depths on CIFAR-10. The results are summarized below:

These results indicate that the correlation remains consistently high as the number of parameters increases.

Q11: The scaling technique to invalidate the measure does not seem effective for the margin-to-Lipschitz-based ratio. Are there alternative approaches to validate or refine this measure?

A11: We thank the reviewer for raising this question. The Lipschitz constant of a neural network is upper-bounded by the product of the spectral norms of its weight matrices across all layers. To evaluate the effectiveness of the margin-to-Lipschitz-based ratio, we conduct experiments to verify the correlation between the sum of spectral norms over margin [4] and the generalization gap. The results are summarized below:

The results indicate that the margin-to-Lipschitz-based ratio shows a weak correlation, suggesting it may not be a reliable predictor of generalization compared to region count.

Q12: Could you clarify what is meant by region count when it is not an integer?

A12: Region count is defined as an expectation. As detailed in Appendix B, we calculate it by randomly sampling 100 subspaces and averaging the region counts across these subspaces. While the region count for each individual subspace is an integer, the averaged value is not necessarily an integer.

We thank the reviewer once again for the valuable and helpful suggestions. We would be happy to provide further clarifications if the reviewer has any additional questions.

[1] Bartlett, Peter L., Dylan J. Foster, and Matus J. Telgarsky. "Spectrally-normalized margin bounds for neural networks." Advances in neural information processing systems 30 (2017).

[2] Keskar, Nitish Shirish, et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).

[3] Neyshabur, Behnam, et al. "Exploring generalization in deep learning." Advances in neural information processing systems 30 (2017).

[4] Jiang, Yiding, et al. "Fantastic generalization measures and where to find them." arXiv preprint arXiv:1912.02178 (2019).

[5] Tsuzuku, Yusuke, Issei Sato, and Masashi Sugiyama. "Lipschitz-margin training: Scalable certification of perturbation invariance for deep neural networks." Advances in neural information processing systems 31 (2018).

[6] Keskar, Nitish Shirish, et al. "On large-batch training for deep learning: Generalization gap and sharp minima." arXiv preprint arXiv:1609.04836 (2016).

[7] Jastrzębski, Stanisław, et al. "Three factors influencing minima in sgd." arXiv preprint arXiv:1711.04623 (2017).

[8] Hoffer, Elad, Itay Hubara, and Daniel Soudry. "Train longer, generalize better: closing the generalization gap in large batch training of neural networks." Advances in neural information processing systems 30 (2017).

Thanks again for your valuable feedback! Could you please let us know whether your concerns have been addressed? We are happy to make further updates if you have any other questions or suggestions.

We greatly appreciate the reviewer's comments and valuable suggestions. We address the reviewer's questions in more detail as follows:

Q1: Omission of important details from introduction.

A1: We appreciate the reviewer's feedback on the missing details. In the updated version of our paper, we have revised the introduction to include the definition of region count and its distinction from the concept of linear regions. The changes are highlighted in blue.

Q2: The section on 'Motivation' in Section 3 is misleading. There are also some typos in Section 3.

A2: The motivation section emphasizes that commonly used measures for characterizing implicit bias, such as norm and margin, are valid in specific cases (e.g., linear or homogeneous networks) [1][2][3]. However, these measures fail to perform well for more general nonlinear networks like ResNet-18. Previous research [4] also shows that these measures do not correlate with the generalization gap. Our experiment further demonstrates their limitations, highlighting the need for a new approach to characterizing implicit bias in general neural networks. This led us to introduce region count, which effectively captures implicit bias. We have addressed the typos and revised Section 3 to improve clarity, with changes highlighted in blue.

Q3: Insufficient theoretical link to generalization.

A3: Thank you for raising this question. The question and suggested theorem are more closely related to optimization rather than generalization. A classifier with a very small region count could indeed perform poorly. however, it does not necessarily contradict the possibility of having a small generalization gap. For instance, a naive classifier with region count 1 will perform poorly on both the training and test datasets, resulting in a small generalization gap.

The suggested theorem could potentially be achieved by proving that large learning rates still allow convergence to a minimizer under the setup considered in Section 6.2. We agree with the reviewer that this paper primarily focuses on the empirical and generalization aspects. A more rigorous theoretical analysis of the tradeoff between generalization and optimization in relation to region counts is a valuable direction for future work, and we plan to explore this in subsequent research.

Q4: Minor typos.

A4: We thank the reviewer for pointing out the typos. These have been corrected in the updated version, with changes highlighted in blue. Specifically, for line 988 in the proof of Lemma 2, the condition for the inner products involving $i$ and $i+2$ is not a typo. We have added a more detailed explanation in the appendix to clarify this point.

Finally, we thank the reviewer once again for the effort in providing us with valuable and helpful suggestions. We will continue to provide clarifications if the reviewer has any further questions.

Reference

[1] Soudry, Daniel, et al. "The implicit bias of gradient descent on separable data." Journal of Machine Learning Research 19.70 (2018): 1-57.

[2] Ji, Ziwei, and Matus Telgarsky. "The implicit bias of gradient descent on nonseparable data." Conference on learning theory. PMLR, 2019.

[3] Lyu, Kaifeng, and Jian Li. "Gradient descent maximizes the margin of homogeneous neural networks." arXiv preprint arXiv:1906.05890 (2019).

[4] Jiang, Yiding, et al. "Fantastic generalization measures and where to find them." arXiv preprint arXiv:1912.02178 (2019).

Thank you very much for your reply and for taking the time to answer my questions.

Most of the changes in the manuscript look very good. Specifically,

Specifically, for line 988 in the proof of Lemma 2, the condition for the inner products involving and is not a typo. We have added a more detailed explanation in the appendix to clarify this point.

I see. Thank you for explaining this further. There is a minor typo in the included equation: In lines 992 and 995, it should be $\tilde{x}$, instead of $x$.

To be honest, my concerns with the very problematic Section 3 have not been addressed, since I find the train of thought very misleading. As I stated in my review, "the norm and margin quantities considered have no reason to be correlated with generalization for a ResNet18". They are generalization measures for a different learning problem. It is akin to using generalization guarantees for decision trees and claiming that the theory is wrong if they do not correlate with generalization in transformers with skip connections and relative positional encoding. However, since I agree with the high-level idea of this section and the authors do cite [Jiang et al., 2019], where a more thorough discussion is included, I will not insist on this further. I leave it to the readers to think critically. I updated my score to reflect this.