A Differentiable Rank-Based Objective for Better Feature Learning

Authors: Krunoslav Lehman Pavasovic, Giulio Biroli, Levent Sagun
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: ranks, statistics, machine learning, feature learning, rank ordering correlation, fairness

Abstract:
In this paper, we leverage existing statistical methods to better understand feature learning from data. We tackle this by modifying the model-free variable selection method, Feature Ordering by Conditional Independence (FOCI), which is introduced in Azadkia & Chatterjee (2021). While FOCI is based on a non-parametric coefficient of conditional dependence, we introduce its parametric, differentiable approximation. With this approximate coefficient of correlation, we present a new algorithm called difFOCI, which is applicable to a wider range of machine learning problems thanks to its differentiable nature and learnable parameters. We present difFOCI in three contexts: (1) as a variable selection method with baseline comparisons to FOCI, (2) as a trainable model parametrized with a neural network, and (3) as a generic, widely applicable neural network regularizer, one that improves feature learning with better management of spurious correlations. We evaluate difFOCI on increasingly complex problems ranging from basic variable selection in toy examples to saliency map comparisons in convolutional networks. We then show how difFOCI can be incorporated in the context of fairness to facilitate classifications without relying on sensitive data.

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

Reviews and discussion:
<thread>
creation date:
1730266652654

thread_type:
official_review

<thread_object>
creator: reviewer_p7dz

<content>
summary:
This paper introduces Chatterjee’s coefficient into feature learning by proposing a differentiable formulation to estimate it. The conditional dependence metric proposed by the authors is applicable to feature learning and enhances performance in downstream tasks.

strengths:
1. The proposed correlation estimation term is highly versatile and “plug-and-play,” making it adaptable for feature selection, predictor training, and as a regularizer to mitigate spurious correlations.
2. The paper is well-organized, with clear sections on methodology, relevant theorems, and extensive evaluation experiments that strengthen the contribution.

weaknesses:
1. There is a disconnect between the theoretical foundation and its practical application, leading to a potentially vacuous theorem. The assumption that the coefficient \(\beta\) of the Softmax function is infinite is impractical, as it is only set to values like 5 in the experiments.
2. Symbol definitions are ambiguous, as discussed in the questions below.
3. The comparison of the proposed method focuses on conventional methods such as LDA and PCA. It would be beneficial to include comparisons with more state-of-the-art methods.

questions:
1. The definition of the rank \( r \) in Line 096 is unclear. What is meant by \( Y_{(j)} \leq Y_{(i)} \)?
2. Line 076 declares \(\sigma\) as a Sigmoid function, yet it denotes the Softmax function. Could the authors clarify this discrepancy?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671051401

thread_type:
official_review

<thread_object>
creator: reviewer_uzzo

<content>
summary:
The paper introduces an algorithm (difFOCI) that gives a measure of correlation based on Feature Ordering by Conditional Independence (FOCI). The method is differentiable and can be used in ML methods. It is used in 3 ways: 

dF1) maximize T(Y,fθ(X)); 

dF2) minimize loss + T(X_G,fθ(X)) such that the features are independent of group features

dF3) maximize T(Y,fθ(X)|X_S) such that features are correlated with Y when conditioned on X_s

Synthetic and more real data experiments are performed for all cases.

strengths:
* S1. The proposed method is sound and has good potential for feature selection and feature debiasing 

* S2. The method has mainly good results in synthetic and some real datasets.

weaknesses:
* W1. The experiments on real datasets are not that strong. 

* W1.1 For the spurious correlation experiments, Waterbirds dataset is a small and simple dataset and a successful method should be tested on other datasets besides it. How many seeds were used? Was the same protocol used to select hyperparameters for the baselines and the proposed method? The benchmark used in [A] can be used to evaluate the proposed method more rigorously. The results of the method will be more reliable if multiple datasets from [A] were used, together with similar hyperparameter selection and evaluation.


* W2. For the fairness experiments, some additional explanations are needed. 

* W2.1 Are the explicit features Xs a subset of X? 

* W2.2 The authors state that “the predictor should be independent of one or more sensitive features”. How is this enforced? By maximizing T (Y , θ ⊙ X | Xs ) the model would learn to predict the target when we know Xs. Is this right? Can the authors explain this?

* W2.3 Is there a measure to evaluate if the model relies on sensitive features? At the moment the results only report performance, and it is not clear if these results incorporate sensitive features or not.


 
Minor: 
* In the context of Waterbirds, group refers to the combination of class and spurious variable (background). So X_G should be denoted as a spurious variable, not a group variable.


[A] Yang et al. "Change is hard: A closer look at subpopulation shift." (2023)

questions:
See weak points.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691266347

thread_type:
official_review

<thread_object>
creator: reviewer_clox

<content>
summary:
This paper uses existing statistical methods to better understand feature learning from data, by modifying an existing model-free variable selection method to a trainable version. Experiments on toy examples and real world datasets show the effectiveness of the proposed method.

strengths:
1. The motivation of the paper is well stated.
2. The paper is well-structured and well-written.
3. Providing results on both toy experiments and real world datasets makes the paper more solid.

weaknesses:
1. The real world datasets seem to be out-dated, where the latest one was released in 2019. It would be more convincing to presents results on more recent are more complex datasets, such as those in WILDS benchmark.

2. This paper only considers one model architecture, i.e., ResNet-50. With the increasing usage of Transformer-based models, it is also important to show the effectiveness on more complex models.

3. Simply showing the improved performance on worst group accuracy does not sufficiently suggest the decrease reliance on the spurious features. Note that, the overall performance of the proposed method shows consistently lower performance \(2 - 4%\) for average accuracy. It would be suggested to conduct the synthetic experiments in [1].

4. Another concern is that the experiments are conducted on the case where the train test set share the same distribution. Given that the feature selector is a trained NN, it is interesting to show if such a method maintains its performance of feature selection where the train test set have a different distribution.  For example, CIFAR10 vs CIFAR10.1, DomainNet-real vs DomainNet-Sketch? as in [2].

I would be happy to increase my score if some of my concerns are resolved. 

[1] The Pitfalls of Simplicity Bias in Neural Networks, NeurIPS, 2020.

[2] A CLOSER LOOK AT MODEL ADAPTATION USING FEA- TURE DISTORTION AND SIMPLICITY BIAS, ICLR, 2023.

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

</thread>
<thread>
creation date:
1730696500387

thread_type:
official_review

<thread_object>
creator: reviewer_zqah

<content>
summary:
The paper introduces a new method called difFOCI, which bases on an existing approach called FOCI. FOCI is a tool for choosing important features from data based on their statistical relationships, but it is not easy to use in some deep learning tasks because it does not adapt well to different models.  Motived by this, difFOCI creates a flexible, learnable version of FOCI that can be used in more situations.

This submision uses difFOCI in three ways: as a tool to pick important features, as part of a model that learns using a neural network, and as a method to help models learn better by reducing irrelevant correlations in the data. To validate this, this submission tests difFOCI on a range of tasks, from simple examples to more complex problems like understanding which parts of images are important in neural networks. Also, difFOCI is shown to be helpful in making fairer predictions by letting models classify without depending on sensitive information.

strengths:
- The motivation is clear to me. FOCI is a tool for selecting important features from data based on their statistical relationships. However, it is not differentiable, which hinders its use in deep neural networks. To address this limitation, this submission proposes a differentiable, parametric approximation.
- This submission provides a clear definition of difFOCI. The toy examples offer some intuition into how difFOCI works.
- The three applications are well-chosen. They effectively demonstrate the usefulness of difFOCI.

weaknesses:
- While some real-world datasets are used in the experiments to demonstrate the effectiveness of difFOCI, it is unclear if it can be extended to large-scale datasets. Specifically, the datasets in Section 5.1 are small-scale, and the neural networks or learning algorithms used are relatively simple. The Waterbird task, for example, is simpler compared to multi-class tasks. Please discuss the scalability and generalization potential of difFOCI.
- The fairness study is interesting; however, the dataset and task here are relatively simple. Can difFOCI be applied in large-scale settings?
- The visualization in Figure 2 is not very clear. What is the main takeaway from this figure? Additional visualization examples would be helpful.

questions:
- Please discuss the scalability of difFOCI on large-scale datasets and its applicability to more challenging tasks.
- Additionally, please clarify the purpose of the visualization in Figure 2.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732201100427

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updates to the paper in response to the reviewers

comment:
We would like to express our gratitude for the reviewers’ thorough evaluations and constructive feedback. While the reviewers found the motivation behind difFOCI clear, they collectively identified two key areas for improvement:
- **More comprehensive experiments**: all reviewers recommended including additional architectures and employing larger, more diverse datasets and further benchmarks.
- **Additional fairness evidence**: reviewers ZQah, cLoX, and uzzo requested additional evidence to demonstrate the capabilities of difFOCI in reducing bias/spurious correlations.

*(Edit 27/11: added medical CheXpert dataset.)*
In response to the first point, reviewers’ consensus prompted us to conduct further experiments. We evaluated difFOCI across five additional datasets, which vary in size, number of classes, modality, architecture, and pretraining displayed in Table 1a:

|Dataset|Year|Size|N. Classes|Modality|Architecture|
|--|--|--|---|--|--|
|MultiNLI|2017|300k|3|Text|BERT|
|CivilComments|2019|250k|2|Text|BERT|
|CelebA|2015|200k|2|Image|ResNet-50 w. ImageNet-1K|
|NICO++|2022|90k|60|Image|ViT-B w. DINO|
|MetaShift|2022|3.5k|2|Image|VIT-B w. CLIP|
|CheXpert|2019|225k|2|Image|ResNet-50 w. ImageNet-21K|
**Table 1a: Dataset Overview**

We further added three benchmarks: Just Train Twice (JTT) [1], Mixup [2] and Invariant Risk Minimization (IRM) [3]. Our results show that difFOCI (with ERM and DRO) achieves competitive performance, both in terms of average accuracy (Table 1b) and worst group accuracy (Table 1c). We appreciate the reviewers suggesting this enhancement.

|Average acc.|difFOCI+ERM|difFOCI+DRO|ERM|DRO|JTT|Mixup|IRM|
|--|---|--|--|--|--|--|--|
|MultiNLI|**81.98±0.2**|**81.82±0.5**|81.4±0.1|80.2±0.6|81.2±0.4|80.7±0.1|77.7±0.3|
|CivilComments|**86.28±0.1**|81.90±0.3|85.7±0.4|82.3±0.4|84.3±0.5|84.9±0.3|85.4±0.2|
|CelebA|94.41±1.1|92.9±2.1|94.9±0.2|93.1±0.6|92.4±1.6|**95.7±0.2**|94.5±1.0|
|NICO++|**85.7±0.3**|**85.8±0.5**|84.7±0.6|83.0±0.1|**85.3±0.1**|84.2±0.4|84.7±0.5|
|MetaShift|**92.1±0.2**|**91.8±0.3**|91.3±0.5|89.0±0.2|90.7±0.2|91.2±0.4|**91.5±0.6**|
|CheXpert|87.1±0.3|81.9±0.5|86.5±0.3|77.9±0.4|75.7±1.7|82.2±5.1|**90.0±0.2**|
**Table 1b: Avg. acc.**


|Worst group acc.|difFOCI+ERM|difFOCI+DRO|ERM|DRO|JTT|Mixup|IRM|
|--|--|--|--|--|---|---|--|
|MultiNLI|**77.6±0.1**|**77.5±0.2**|66.9±0.5|77.0±0.1|69.6±0.1|69.5±0.4|66.5±1.0|
|CivilComments|66.32±0.2|**70.3±0.2**|64.1±1.1|**70.2±0.8**|64.0±1.1|65.1±0.9|63.2±0.5|
|CelebA|**89.32±0.4**|**89.8±0.9**|65.0±2.5|**88.8±0.6**|70.3±0.5|57.6±0.5|63.1±1.7|
|NICO++|**47.10±0.7**|46.3±0.2|39.3±2.0|38.3±1.2|40.0±0.0|43.1±0.7|40.0±0.0|
|MetaShift|83.10±0.5|**91.7±0.2**|80.9±0.3|86.2±0.6|82.6±0.6|80.9±0.8|84.0±0.4|
|CheXpert|54.42±3.2|**75.3±0.3**|50.1±3.5|73.9±0.4|61.5±4.3|40.2±4.1|35.1±1.2|
**Table 1c: Worst group acc.**

To address the second point, we expanded analysis in Section 5.3 showing that (dF3) maintains good predictive performance of $y$ (from $X$), while also significantly decreasing predictability about sensitive attributes $X_s$.

Specifically, we train two NNs to predict $y$: the first NN was trained on $X$ (without $X_s$), while the second NN was trained on features $f_\theta(X)$ obtained using (dF3). We then used the final layers of both NNs to predict the sensitive $X_s$ and found that difFOCI (dF3) significantly reduced the predictability of $X_s$ (sometimes to chance level) without significantly impacting accuracy on $y$ (in some cases even slightly improving it). The results are presented in Table 2 below:


|Dataset|Features|Train acc: $y$|Val. Acc: $y$|Test acc: $y$|Train acc: $X_s$|Val. Acc: $X_s$|Test acc: $X_s$|
|--|--|--|--|---|--|--|--|
|Bank marketing|Stand. data|91.32±2.3|93.27±1.2|90.05±2.0|89.09±1.2|72.26±1.5|70.93±0.9|
||(dF3) features|90.81±1.8|92.13±2.6|89.35±1.1|63.12±2.8|62.24±0.7|**63.81±2.1**|
|Student data|Stand. data|88.35±1.7|79.63±0.9|75.67±1.3|95.68±2.1|72.16±2.4|71.21±1.5|
||(dF3) features|80.18±2.9|72.16±1.6|72.73±1.7|59.47±1.1|58.95±1.0|**48.89±1.1**|
|ASCI Income|Stand. data|83.49±2.4|85.10±2.1|81.30±2.7|68.97±1.6|67.67±2.6|66.00±0.7|
||(dF3) features|82.80±0.8|81.99±1.5|82.95±0.9|56.58±1.2|55.01±2.0|**52.73±2.0**|
**Table 2. Fairness experiments**




While these additions did not substantially alter the manuscript's content, apart from rephrasing subsection 5.3 in response to Reviewer uzzo's valid remark, we believe they considerably strengthened the paper. We extend our gratitude to all the reviewers for their contributions to these improvements.

---

[1] Liu, E.Z., Haghgoo, B., Chen, A.S., Raghunathan, A., Koh, P.W., Sagawa, S., Liang, P. and Finn, C., 2021, July. Just train twice: Improving group robustness without training group information.

[2] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. Mixup: Beyond empirical risk minimization. 

[3] Arjovsky, M., Bottou, L., Gulrajani, I. and Lopez-Paz, D., 2019. Invariant risk minimization.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732555551638

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We appreciate the time and effort taken by all reviewers to provide feedback on our paper. We have now responded to all comments and uploaded a revised manuscript. If there are any other questions or concerns, we would be happy to address them.
</content>
</thread_object>

</thread>


Metareview:
metareview:
Here the authors propose a new algorithm called diffFOCI which is a differentiable approximation to a recently proposed variable selection method called FOCI (Azadkia and Chatterjee (2021)). The authors combine diffFOCI with NN end-to-end training and show promising results in three very important areas -- feature selection, domain-shift/spurious-correlation, and fairness.

All the reviewers unanimously gave a final score of 6. Though they appreciated the underlying idea, motivation, and the writing quality; each raised more or less the same concern which was regarding the limited experimental results. However, during the rebuttal period the authors did a wonderful job in answering most of their queries and presented new set of promising results on a variety of new datasets and artchitectures (BERT, DINO, CLIP etc.). There were some concerns regarding the theorem and the independence assumptions which also the authors replied convincingly. 

I think that this paper is a very nice blend of theory and practice and the results are very promising. Of course the proposed method is not an absolute winner in all the tasks presented in this work but still it seems to be the one with the most promising results throughout. Hence, I support the acceptance of this work and would strongly suggest the reviewers to incorporate all the comments by the reviewers in their revised draft. They all makes sense and would surely help in improving the quality/visibility of the work.