Discovering Latent Groups for Robust Classification

2026-06-22

AI Summary

This work addresses the performance degradation of machine learning models on minority subpopulations due to reliance on spurious correlations. To mitigate this issue, the authors propose the Neural Classification Tree (NCT) framework, which explicitly embeds latent subgroup structure into the model architecture for the first time. NCT employs a tree-based routing mechanism that dynamically assigns samples to “easy” or “hard” nodes based on prediction correctness, and iteratively refines the model using pseudo-labels derived from routing paths—without requiring ground-truth subgroup annotations. Evaluated across five benchmarks encompassing both binary and multi-class spurious correlations, the method consistently identifies and improves performance on minority groups while offering strong interpretability and robustness comparable to state-of-the-art approaches.

Abstract

Machine learning models exploit spurious correlations, achieving high average accuracy but failing disproportionately on underrepresented subgroups. Existing methods address this by adjusting network parameters, guided either by subgroup annotations or inferred pseudo-group labels. Yet at inference, these methods produce only a class prediction, with no insight into a sample's latent subgroup. We propose neural classification trees (NCT), a framework that achieves robustness by encoding subgroup structure in its tree-shaped architecture. By routing each sample to an "easy" or "hard" node of this tree -- based on prediction correctness -- and reusing these routes as pseudo-labels for the next iteration, NCT disentangles conflicting subgroups, without requiring subgroup supervision. We evaluate NCT on five benchmarks spanning binary and multi-class spurious correlations. Our experiments show that the learned tree topology provides strong interpretability by consistently isolating minority subgroups, which provides a transparent mapping between the model architecture and the data's latent group structure, while yielding competitive robustness with state-of-the-art methods.

Problem

• spurious correlations
• latent subgroups
• robust classification
• subgroup fairness
• model interpretability

Innovation

• neural classification trees
• latent subgroups
• spurious correlations
• robust classification
• interpretable architecture