Task-Relevant Covariance from Manifold Capacity Theory Improves Robustness in Deep Networks

Authors: william  yang
chi-ning  chou
sueyeon  chung
Venue: NeurIPS_cc_2024_Workshop_UniReps
Published: 2024-09-24 01:59:54

Abstract: Analysis of high-dimensional representations in neuroscience and deep learning traditionally places equal importance on all points in a representation, potentially leading to significant information loss. Recent advances in manifold capacity theory offer a principled framework for identifying the computationally relevant points on neural manifolds. In this work, we introduce the concept of *task-relevant class covariance* to identify directions in representation-space supporting class discriminability. We demonstrate that scaling representations along these directions markedly improves simulated accuracy under distribution shift. Building on these insights, we propose AnchorBlocks, architectural modules that use task-relevant class covariance to align representations with a task-relevant eigenspace. By appending one AnchorBlock onto ResNet18, we achieve competitive performance in a standard domain adaptation benchmark (CIFAR-10C) against much larger robustness-promoting architectures. Our findings provide insight into neural population geometry and methods to interpret/build robust deep learning systems.

Main paper: 
 introduction: 
 In computational neuroscience and deep learning, geometric methods such as canonical correlation analysis (CCA) [1] , representational similarity analysis (RSA) [2] and centered-kernel alignment (CKA) [3] have emerged as an effective approach for the analysis of high-dimensional neural representations. In such methods, point-cloud representations are often analyzed with a uniform prior, i.e., without explicitly privileging certain points over others. However, in their seminal work, Averbeck et al. [4] demonstrate how such uniformity assumptions on data points could lead to misestimation of information content relevant for downstream decoding. Overcoming this issue by optimally weighting data points is challenging in practice, which leaves us without a systematic framework for incorporating task-relevant information into the geometric analysis of neural representations. Manifold capacity theory is a data-driven framework for quantifying task-relevant structure in neural representations through the lens of downstream readout by perceptrons [5, 6] , and has found applications in both neuroscience [7, 8, 9, 6] and machine learning [10, 11, 12, 13, 14, 15, 16] . Manifold capacity generalizes the classic notion of storage capacity (e.g., Cover's theorem [17] , Gardner's formula [18] ) and quantifies the coding efficiency of neural representations in terms of the number of decodeable manifolds 1 stored per neuron. At the core of the theory is a mathematical result relating manifold capacity to a statistical distribution of anchor points (see Equation 2) which are the representative support vectors of each class manifold for a given decision boundary. Intuitively, the anchor distribution can be understood as forming a task-relevant geometry space seen by an optimal linear classifier [6] . In particular, the task-relevant geometry produces class covariance structures that differ from those of the original point-cloud. As a toy example, Figure 1a depicts spherical class manifolds globally arranged in a ring structure 2 . For each class as well as globally, the principal directions of covariance are uninformative. However, if the goal of downstream neurons is to distinguish classes, then the task-relevant directions (red) are tilted, and would differ between class manifolds. This intuition is reflected by the covariance structure of the class anchor points (yellow). In this work, we apply these geometric insights to deep learning to investigate whether task-relevant class covariance from manifold capacity theory helps identify directions in representation-space useful for attenuating irrelevant or amplifying discriminative features. We focus on the setting of distribution shift, in which DNNs are vulnerable to even minor non-stationarities in input data, such as corruptions. We choose this setting due to extensive work, empirical and theoretical, attributing this vulnerability to a DNN's under-emphasis of auxillary class-invariant features, or over-emphasis on easily-degraded low-level features (see [22, 23, 24, 25, 26, 27, 28] ). Using the standard CIFAR-10C domain adaptation task [29] , we demonstrate the usefulness of task-relevant class covariance to identify linear transformations which improve robustness to distribution shift (see Figure 1b ).Our contributions can be summarized as follows:• In Section 2.2, we formally introduce the concept of task-relevant class covariance, based on manifold capacity theory, as a way to identify directions in representation-space which support discriminability between classes.• In Section 3.1, we find that manually augmenting ResNet18 representations with taskrelevant class covariance improves simulated accuracy in domain adaptation on CIFAR-10C.In contrast, augmenting with the original point-cloud covariance is detrimental.• In Section 3.2, we introduce AnchorBlocks, which are architectural modules designed to scale DNN representations along the eigenspace of task-relevant class-covariance matrices. We show that one AnchorBlock fine-tuned on a ResNet18 backbone achieves competitive domain adaptation on CIFAR-10C against larger robustness-promoting ResNet variants. 
manifold capacity theory and anchor geometry: 
 Given a dataset {(x i , y i )} where x i is the i-th input data and y i is its label, which comes from a finite discrete set Y. We denote (x i ) 2 R N as the representation of x i (e.g. the hidden activation of a DNN) and N as the representation dimensionality. For each class y 2 Y, the corresponding manifold M y = conv({ (x i ) : y i = y}) is defined as the convex hull of the representations of the stimuli with label y. From manifold capacity theory [5, 6] , the manifold capacity of (•) is denoted as ↵, which exhibits a closed-form formula depending only on the structure of the manifolds as follows:EQUATIONwhere (y, y 0 ) = 1 if y = y 0 and 1 otherwise, Unif(•) is the uniform distribution of a finite set, vector g is drawn from the multivariate Gaussian distribution N (0, I N ), and are Lagrange multipliers enforcing constraints of the convex optimization problem:min V 2R N kV gk 2 2 : (y, y 0 )V > h  0, 8y 0 , 8h 2 M y 0The term inside Eq. ( 1) has several interpretations, including the average error made by a random classifier. See [6] for further discussion and details. The key relevant observation is that Eq.(1) suggests a joint distribution over the manifolds {M y } y2Y : {s y 0 (y, g)} y 0 2Y ⇠ D anchor where y ⇠ Unif(Y) and g ⇠ N (0, I N ). As such, we have:EQUATIONwhere f (•) : ⌦ y2Y M y ! R 0 is a simple function as defined in Eq. ( 1). This establishes an analytical connection between manifold capacity ↵ and the distribution of anchor points D anchor . (3)Model. All DNN experiments use a publicly-available pretrained ResNet18 [30] obtained from the Huggingface model repository [31] trained on CIFAR-10 [32] (see Section A.1).Data. CIFAR-10C is used to evaluate robustness, and extends CIFAR-10 by algorithmically applying 15 types of common corruptions at 5 severity levels (1-5), resulting in 75 different test sets.Corruptions belong to four broad categories: noise, blur, weather and digital transformations [29] .Robustness Metrics. Per convention in [33] , for each corruption we benchmark the Corruption Error (CE), the classification error rate, and Mean Corruption Error (mCE) which averages error over all 75 corruptions. We also report Clean Error, which is the error rate on the CIFAR-10 test set. 
results: 
 We demonstrate the application of task-relevant covariance in two settings. In Section 3.1, we show that scaling the representations of a class along the eigenspace of their task-relevant class covariance matrix causes all class examples to be correctly classified even under distribution shift. In Section 3.2, we incorporate this geometric insight into a trainable DNN architecture which achieves competitive robustness on CIFAR-10C.If anchor geometry captures task-relevant class-specific features, then one prediction is that scaling representations along the eigenspace of their task-relevant class covariance should improve performance. To test this, we perform a controlled representation manipulation on pretrained ResNet18.Implementation. Let H i µ 2 R N be the representation of the i-th sample of the µ-th class on a test set, where i 2 [M ], µ 2 [P ]. We first run pretrained ResNet18 in inference mode on the CIFAR-10 training set and extract activations used to compute a task-relevant covariance matrix for each class (see Eq. 3). We then run an inference pass on CIFAR-10C and augment each H i µ with the µ-th task-relevant class covariance matrix C eff µ , which can be summarized as H Aug µ := H µ C eff µ (see Algorithm 1). To test the goodness of this augmented representation, we pass it through the model's readout to calculate simulated accuracy (Figure 2 ). As an experimental control, we perform the same procedure using the µ-th point-cloud class covariance matrix from the CIFAR-10 training set. Augmentation with task-relevant class covariance results in perfect simulated accuracy. Note that the manipulation used here is "label-aware", as the label determines which of the P task-relevant class covariance matrices to use for a given test example. We use the term "simulated accuracy" to indicate that it is a measure used for analysis rather than to represent naturalistic classifier performance. We found representations achieve perfect simulated accuracy when scaled in the eigenspace of their task-relevant class covariance. Meanwhile, when representations are scaled by the point-cloud class covariance, simulated accuracy is worse than baseline. This is consistent with our intuition that taskrelevant class covariance captures directions in feature-space supporting discriminability, whereas point-cloud class covariance merely captures directions of high within-class variation.Although the procedure in Section 3.1 offers interesting insight relating task-relevant covariance to robustness, it is label-aware and thus would introduce a leakage if used in real-world settings. Here we introduce AnchorBlocks, architectural modules which allow models to learn to scale their representation by the correct task-relevant class covariance without label-assistance. Architectural summary. We replace ResNet18's final layer with an AnchorBlock (see Figure 3a ). For a P -class problem, the AnchorBlock has P parallel class covariance heads with weights as task-relevant class covariance matrices, followed by a shared linear readout. The representation h 2 R N is processed by all heads, yielding augmented vectors hi , i 2 are reported in [33] , while mCE for ResNeXt29 + TENET is reported in [27] .weights of each head are set to a class covariance matrix and kept fixed by disabling gradients. Then, we replace the readout of ResNet18 with one AnchorBlock, and fine-tune the combined model on the CIFAR-10 training set, keeping the objective function the same. This architecture is label-unaware since the representation is scaled by all P task-relevant class covariance matrices in parallel.ResNet18+AnchorBlock achieves competitive performance compared to robust ResNet variants. The task objective is to minimize error on CIFAR-10C using only CIFAR-10 training data. Following the convention of [33] , we report the robustness of our classifier architecture using mCE (mean corruption error), which is simply the average classification error across all 75 CIFAR-10C test sets (see Figure 3b ). The architectures we compare with are much larger, while DenseNet/RexNeXt implement enhanced feature aggregation -both of these properties purportedly promote robustness (see [29] for good discussion). In addition, our method also outperforms the combination of ResNeXt and TENET, an architectural modification designed to promote robustness and feature diversity [27] . 
conclusion: 
 In this work, we introduce the novel concept of task-relevant covariance derived from the statistics of anchor points in manifold capacity theory, and propose that they identify directions in representation space supporting class discriminability. In Section 3.1, we use a manual representation augmentation experiment to test our prediction about these discriminability properties, and find that simple linear transformations which scale representations along the eigenspace of task-relevant class covariance matrices markedly improve simulated accuracy on CIFAR-10C. In Section 3.2, we use these insights to create AnchorBlocks, trainable architectural modules which perform these linear transformations in a label-unaware manner. We find that fine-tuning one AnchorBlock on a ResNet18 backbone achieves competitive performance on CIFAR-10C compared to larger robustness-promoting architectures.Our approach suggests a geometric interpretation of distribution shift as data non-stationarity that causes greater manifold tangling, in which case simple linear transformations may improve the robustness of representations via manifold untangling (see Figure 1b ). In this geometric interpretation, task-relevant class covariance is a way to identify such linear transformations based on the geometry space seen by a downstream classifier. In a broader sense, our work emphasizes the importance of considering task-relevant structure when analyzing high-dimensional representations, and also highlights the differences of viewing representational statistics from the lens of an ideal observer versus a downstream perceptron. This perspective may have significant implications for the analysis of representations in both computational neuroscience and deep learning. 
 
 Appendix: Although a toy example, ring-structured representations have been empirically observed, e.g.,[19,20,21] 


Decision: In light of the positive reviewers' feedback and relevancy of the submission, we are pleased to accept this paper for presentation at UniReps 2024. We kindly ask the authors to incorporate the reviewers' suggestions and feedback in the final camera-ready version of the manuscript.

Reviews: review: title: title: Review on TRC from MCT Improves Robustness in DN
main_review: main_review: This paper presents a novel method to enhance the robustness of deep neural networks when faced with distribution shifts by utilizing task-relevant class covariance based on manifold capacity theory. 

Strengths
1. Application of Manifold Capacity Theory: The paper applies manifold capacity theory to identify task-relevant features, creating a connection between theoretical neuroscience concepts and deep learning.

2. Adaptability and Simplicity: The proposed AnchorBlocks represent a novel architectural feature that can be easily integrated into existing models, improving robustness without major architectural changes. Also it is easy to implement and do not require complex training processes, making the approach more feasible for practical applications.

3. Empirical Validation: The method demonstrates its effectiveness on a standard robustness benchmark (CIFAR-10C), delivering competitive performance against larger and more complex models focused on robustness.


Weaknesses
1. Comparison with State-of-the-Art Methods: The paper lacks comparisons with recent state-of-the-art robustness techniques, such as adversarial training or advanced data augmentation strategies, which could provide a better context for the method's performance.

2. Potential Computational Overhead: The introduction of multiple parallel class covariance heads could increase computational complexity, but the paper does not analyze or address these potential costs. An evaluation of computational efficiency during training and inference would be helpful.

3. Dependence on Label Availability for Training Covariance Matrices: While the AnchorBlock does not rely on labels during inference, it requires task-relevant class covariance matrices, which are computed using labels during training. This could limit the applicability of the method in scenarios with limited labeled data.



Overall, given the innovative approach, practical significance, and promising results, I recommend this paper for publication.
title: title: Review of submission #50
main_review: main_review: The main idea of this paper: introduces task-relevant class covariance to improve robustness in deep networks under distribution shift.
An important highlight, which I found unique was the idea they proposed of the AnchorBlock acheiving competitive performance to many benchmarks.

Strengths:
- Clear and well motivated writing -- task-relevant covariance derived from manifold capacity theory to enhance class discriminability.
- For initial stage results, they showed that ResNet18 + AnchorBlock outperforms larger architectures designed for robustness, such as ResNeXt + TENET.

Weaknesses:
- Would still like to see behaviour on a different dataset something more complex like CIFAR100.
- Incorporating some real-world distribution shifts (e.g., adversarial examples, real-world corruptions) would provide more practical insights into the effectiveness of task-relevant covariance for robustness. (this is more of question/suggestion to authors)

Question out of curiosity: is the AnchorBlock scalable to larger models and datasets? Did authors try any experiments regarding scalability?
title: title: Review of Submission50
main_review: main_review: This paper proposes to exploit insight coming from manifold capacity theory to identify directions of maximal discriminability in representation space. Experimental results show that scaling representations along such directions produces promising results in classification under distribution shifts.

Strengths:

- The proposed method is interesting and well grounded in theory.
- Experimental results on domain adaptation, although preliminary, are very promising and already competitive with some previous solutions.
- Writing is clear, and experimental sections are well detailed.

Weaknesses:

- Results are reported in terms of mCE against other ResNet variants, showing that ResNet18 with AnchorBlock improves performance with respect to other models, in their standard version or trained with TENET. However, it would have helped to have at least a comparison with the current state of the art performances that can be obtained with models of this family. Moreover, it would be useful to assess the proposed method on additional datasets.
- The mathematical formulation in Eq. 1 is not very clear: e.g. the terms $\mathbf{g}$ and $\lambda$ that appear in the manifold capacity equation are never discussed in the text and their interpretation is not straightforward.
- (Minor) there is a typo in line 45 (auxillary -> auxiliary).