Discovering Latent Groups for Robust Classification

Unknown

Conceptio › Archive › arXiv CS

arXiv CSopen access

Discovering Latent Groups for Robust Classification

Unknown · 2026 · arxiv_cs

arXiv CS · Papers · License: Open Access · 2026

Open Source ↗Direct PDF ↓

machine learning, deep learning, neural networks

Discovering Latent Groups for Robust Classification

Ankur Garg, Ulrich Aïvodji, Samira Ebrahimi Kahou, Vincent Michalski

arXiv:2606.23609v1 [cs.LG] 22 Jun 2026

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. 1

Introduction

Deep neural networks are prone to learning shortcuts or features that are correlated with labels in the training set but have no causal relationship to them. [1]. A classic example is Waterbirds dataset, where a standard model will classify birds based on the background (water vs. land) rather than the bird’s features. When such spurious correlations shift at test time, performance degrades for minority subgroups, such as waterbird in front of a land background. Although the research community has developed effective methods to mitigate this issue, a critical gap still remains: interpretability of the latent subgroup structure. Supervised approaches such as group distributionally robust optimization (group DRO) [2], while robust, require expensive, finegrained annotations for every training sample. Just train twice (JTT) [3], environment inference for invariant learning (EIIL) [4], and deep feature reweighting (DFR) [5] drop training-time labels but still rely on validation group labels, while GEORGE [6] and explainability-guided pseudo-group map (ExMap) [7] operate without any group annotation. Across all three regimes, the model’s parameters are adjusted to improve worst-group metrics, but the final classifier remains opaque at inference: it does not reveal which latent subgroup a specific sample belongs to, nor does it structurally isolate conflicting features. We take a different approach: instead of only relying on parameter updates to handle diverse groups, we make the partition architectural. We introduce NCT, a framework that builds on a well-established observation: training difficulty is a reliable proxy for semantic group identity, with samples aligned with spurious correlations being easy and conflicting counter-examples being hard. NCT iteratively routes samples to easy and hard branches based on this signal, and — unlike methods that use the same signal transiently — preserves the resulting partition as the inference-time architecture itself. Our contributions are: (i) Architecture as partition. Where prior work [3, 8, 9, 7] uses correctness, loss, or attribution signals transiently for reweighting or pseudo-group inference, neural classification 1 https://github.com/agarg-dev/Neural-Classification-Trees/

Preprint.

tree (NCT) preserves the difficulty-based partition as the inference-time architecture, with each leaf encoding both the predicted class and the difficulty path. (ii) Unsupervised depth selection. A pseudo worst-group accuracy (pWGA) criterion with a Wilson-tolerance early-stopping bound decides when to stop deepening the tree without group annotations. (iii) Theoretical motivation. Iteration-1 errors are enriched with minority samples under simplicity bias (Theorem 4.2); structural separation yields a positive minimax-risk gap under feature conflict (Theorem 4.4). (iv) Empirical evidence. Across five benchmarks, NCT concentrates minority subgroups in hard branches (82% of landbird-on-water, 73% of blond-male, 47% of benign-no-patch, 71% of digit-8 and 84% of color-conflicting digits) while matching or approaching state-of-the-art worst-group accuracy against eight baselines.

Related Work

Optimization for Spurious Correlations. A dominant paradigm for mitigating spurious correlations modifies the optimization objective to upweight minority groups. When group labels are available, group DRO [2] minimizes the worst-case group loss. When training labels are unavailable, recent approaches infer pseudo-groups to guide optimization. GEORGE [6] clusters representations from a standard empirical risk minimization (ERM) model to estimate subclass labels, subsequently using them for robust optimization. JTT [3] identifies error sets from an early-stopped ERM model and retrains the final model by upweighting these hard samples. Correct-n-contrast (CnC) [8] extends this by applying contrastive learning to push the inferred groups apart in feature space. Other works focus on the classifier head; DFR [5] freezes the feature extractor and retrains only the last layer on a group-balanced validation set. Inferring Environments and Pseudo-Groups. A parallel line of work treats group identification as an inference problem in its own right. EIIL [4] learns a soft environment partition that maximally violates an invariance penalty, which is then handed to invariant or distributionally robust optimizers. LfF [10] reweights a debiased model using the loss of a deliberately biased one, while SelecMix [11] exploits bias-conflicting pairs. More recent work sharpens the partition: XRM [9] trains twin networks on disjoint halves of the training data and uses confident held-out cross-mistakes to discover environments without group-annotated validation data, and GIC [12] infers groups from a spurious-attribute classifier whose predictions vary across distributional shifts. These methods share a downstream pattern: the discovered partition is consumed by a separate invariant-learning or reweighting algorithm, and the resulting classifier is opaque at inference. NCT is orthogonal: its contribution is the architectural persistence of the partition, rather than a sharper pseudo-group signal. Input-Level and Explainability-Guided Interventions. Rather than altering the loss function, some methods intervene directly on the data or leverage post-hoc explanations. MaskTune [13] forces the model to explore robust features by masking out the most discriminative regions of the input image during fine-tuning. ExMap [7] clusters explainability heatmaps from a pre-trained model into pseudo-groups, exploiting the observation that spurious and robust decisions produce distinct attribution patterns. NCT contributes to this neighborhood by using correctness rather than attributions as the difficulty signal, and by preserving the resulting partition as the inference-time architecture rather than as a one-off clustering step. Structural and Modular Interpretability. A complementary line of work builds interpretability directly into the model architecture. Concept bottleneck models (CBMs) [14] force information to pass through a layer of human-aligned concepts but typically require expensive concept annotations or discovery algorithms like concept balancing technique (CoBalT) [15]. Tree-structured classifiers such as ProtoTree [16] and NBDT [17] also expose inference-time decision paths, but through class-taxonomy hierarchies; neither was designed for spurious-correlation settings. NCT adopts a modular approach akin to mixture of experts (MoE) [18, 19] or neural trees [20]. Unlike standard MoE that routes samples to maximize predictive likelihood—and therefore collapses to spurious shortcuts—NCT routes by correctness, decomposing the spurious-vs-core subgroup structure. Across these four families, NCT is the only entry whose inference-time architecture exposes the spurious-correlation subgroup partition: each leaf records both the predicted class and the difficulty path that produced it, with no post-hoc clustering required. 2

Figure 1: NCT inference. Backbone features propagate through parent-to-child head connections across iterations. The argmax over all leaf node outputs determines the group label, encoding both class and difficulty path. NCT Architecture

Backbone

Parent Heads

K Heads

K Probs

argmax

Predicted

Sample

Yes

2ℓ (easy)

2ℓ+1 (hard)

Match?

Assigned (node ℓ)

(t)

Figure 2: Routing mechanism. A sample’s current node assignment ℓi is updated based on prediction correctness: correctly classified samples proceed to the easy branch, while misclassified samples are routed to the hard branch.

Methodology

3.1

Problem Formulation

We consider a C-way classification task over a dataset D = {(xi , yi )}N i=1 , where xi ∈ X is an input image and yi ∈ {0, . . . , C−1} is the class label. Our objective is to learn a hierarchical mapping that decomposes the data distribution into partitions of varying difficulty without attribute supervision. The training process proceeds in discrete iterations t ∈ {1, . . . , T }, where the model maintains a set of K (t) = C · 2t−1 classification nodes. Each training sample xi is assigned a unique node index (t) (1) ℓi ∈ {0, . . . , K (t) −1}, with ℓi = yi . Each node serves as a specialized expert for its assigned samples. Setting C = 2 recovers the binary case. 3.2

Hierarchical Feature Architecture

The NCT is defined by a shared backbone fθ : X → Rd and a hierarchy of classification heads H(t) as illustrated in Figure 1. Each head, indexed by node j ∈ {0, . . . , K (t) −1}, has internal (t) (t) representation ϕj (x) and output logit hj (x). For t = 1, heads receive backbone features fθ (x); (t−1)

for t > 1, head j receives parent representation ϕ⌊j/2⌋ (x). We denote the full output vector as (t)

(t)

h(t) (x) = [h0 (x), . . . , hK (t) −1 (x)]. 3

For the initial iteration (t = 1), the heads operate directly on the backbone features z = fθ (x). For subsequent iterations (t > 1), we implement a hierarchical feature flow. Rather than discarding previous heads, we retain all heads from prior iterations and utilize them as intermediate feature extractors. A specific child head j at iteration t receives its input not from the raw backbone, but from (t−1) the parent representation ϕ⌊j/2⌋ (x). This ensures that specialized heads in deeper layers operate on features that have already been conditioned by the class-specific semantics of the parent. Asymmetric Head Capacity. For t ≥ 2, each split produces an even-indexed easy child (2j) and (t) an odd-indexed hard child (2j+1). Easy children use a linear ϕj ; hard children use a linear–ReLU– dropout block. The asymmetry reflects what each branch handles: easy partitions contain samples the parent already classifies correctly and need little added capacity, while hard partitions absorb the misclassified samples whose conflicting features demand richer representations. Inference. During inference, backbone features propagate through the hierarchy and all leaf nodes produce outputs in parallel as shown in Figure 1. We define the predicted node assignment as (t) (t) ℓ̂i = arg maxj hj (xi ). The final class is read off from the predicted leaf: ŷi =

(T ) T −1 ℓ̂i / 2 ,

(T )

(1) (T )

which generalizes the binary ℓ̂i < 2T −1 rule to C classes. Beyond the class, ℓ̂i also records the easy/hard trajectory along the route to its leaf, exposing the model’s structural grouping of the input. 3.3

Difficulty-Based Sample Partitioning

We employ a correctness-based partitioning strategy to determine the tree topology (Figure 2). This mechanism divides the training data based on the alignment between the predicted and assigned nodes. The routing is deterministic and driven by training dynamics. We define the binary difficulty (t) (t) (t) indicator as di = I[ℓ̂i ̸= ℓi ]. The node assignment for the next iteration follows (t+1)

ℓi

(t)

= 2 ℓ i + di .

(2)

This update rule guarantees that a parent node j splits into exactly two children: an easy child (2j) containing samples correctly classified at step t, and a hard child (2j + 1) containing misclassified (1) (2) (T −1) (T ) samples. Consequently, the binary string di di · · · di along the route to leaf ℓi encodes the precise history of training difficulty for the samples assigned to it. 3.4

Optimization and Stability

Training a hierarchical model iteratively can lead to catastrophic forgetting [21], where updates for child nodes destabilize the parent representations that children depend on. To mitigate this, we employ a stabilized optimization objective and a two-phase training protocol. Auxiliary Loss. To ensure that the feature representations remain valid for the parent tasks, we enforce an auxiliary loss on the parent nodes during child training. The total loss at iteration t is Ltotal = LBCE σ(h(t) (x)), ℓ(t) + λaux LBCE σ(h(t−1) (x)), ℓ(t−1) ,

(3)

where LBCE is the one-vs-all binary cross-entropy (BCE) computed against the one-hot encoding of the node assignment, ℓ(t−1) = ⌊ℓ(t) /2⌋ is the parent assignment, h(t−1) (x) is the output logits of the parent heads, and λaux controls the trade-off between child specialization and parent stability. Two-Phase Training. Each iteration is trained in two phases to stabilize feature learning. Phase 1 freezes the backbone and parent layers so the new heads adapt to the existing feature space. Phase 2 fine-tunes the entire network so the backbone can resolve hard nodes that were previously inseparable. 4

3.5

Controlling Hierarchy Depth

A correctness-based partition grows multiplicatively in T , but not every dataset benefits from a deeper tree. NCT pairs two mechanisms—sparse-node merging and a pseudo-WGA depth criterion—both operating without group annotations. Sparse-Node Merging. Routing can produce hard children with very few samples, especially after several splits. We merge any hard child j with training count nj < mmin into its easy sibling j − 1, reassigning the corresponding pseudo-labels. At inference, the merged head’s logit is masked to −∞, ensuring the arg max never selects an unused leaf. Setting mmin = 0 disables merging. Pseudo-WGA Criterion. To decide when to stop deepening, at each iteration t ≥ 2 we group (t) (t) validation samples by the iteration-2 ancestor of their predicted leaf, ai = ⌊ℓ̂i /2t−2 ⌋ ∈ (t) {0, . . . , K (2) −1}, and define pWGA2 as the minimum validation accuracy across these K (2) groups. Evaluating at iteration-2 granularity (rather than at the current depth) keeps each group large enough for reliable accuracy estimates as the tree deepens. We keep iteration t as the operating depth only if its pWGA2 stays within a Wilson tolerance of the running best: q ⋆ ⋆ (t) (s) ) pWGA2 ≥ max pWGA2 − zα p n(1−p , (4) worst s<t

where p⋆ is the running best, nworst the worst group’s sample count, and zα = 1.96. If the inequality fails we stop deepening and return the iteration with the highest pWGA2 . Appendix F.3 verifies that pWGA2 tracks true worst group accuracy (WGA) closely across our benchmarks.

Theoretical Analysis

We provide theoretical motivation for NCT’s design by analyzing: (1) when correctness-based routing recovers minority subgroups, and (2) why structural separation can reduce worst-case risk compared to single-model approaches. Setup. Consider a classification problem where inputs x ∈ X have class labels y ∈ {0, 1} and latent attributes a ∈ {0, 1} representing spurious features. Let ρ = P (a = y) > 0.5 denote the spurious correlation strength. This creates four subgroups: two majority groups where a = y (e.g., waterbirds on water) and two minority groups where a ̸= y (e.g., waterbirds on land). Standard ERM fails on minority groups because it exploits the spurious correlation. 4.1

Minority Group Recovery (t)

(t)

NCT routes samples to “easy” and “hard” branches based on prediction correctness: di = I[ℓ̂i ̸= (t) (1) ℓi ]. Since ℓi = yi , we analyze when this error-based signal at the first iteration correlates with minority group membership. Assumption 4.1 (Simplicity Bias). After iteration 1 training, the model’s predictions satisfy P (ℓ̂(1) = a) = 1 − ϵ for some small ϵ ≥ 0. This assumption reflects the empirically observed tendency of neural networks to learn simpler, more prominent features before complex ones [22, 23]. When spurious correlations are strong, the spurious attribute provides an easy solution that gradient descent finds early in training. (1)

(1)

Theorem 4.2 (Minority Enrichment). Under Assumption 4.1, let E = {i : ℓ̂i ̸= ℓi } denote the set of misclassified samples after iteration 1. The proportion of minority samples in this error set is: (1 − ϵ)(1 − ρ) (5) P (a ̸= y | i ∈ E) = ϵρ + (1 − ϵ)(1 − ρ) When ϵ = 0, this equals 1: all misclassified samples are from minority groups. (1)

Proof sketch. Since ℓi = yi , we write y for the true label. Under simplicity bias, the per-group error rates are ϵ on majority groups (a = y) and 1 − ϵ on minority groups (a ̸= y). Applying Bayes’ rule with prior P (a ̸= y) = 1 − ρ yields Equation (5). Full derivation in Appendix I.1. 5

For Waterbirds (ρ = 0.95, ϵ = 0.02), Equation (5) gives ∼72% minority share in the error set—a 14× enrichment over their 5% population fraction, which routes them into the hard branch. The bound extends to C-way classification by replacing ρ with the per-class majority-attribute probability ρy = P (a = a⋆ (y) | y), preserving enrichment whenever ρy > 1/C. 4.2

Benefit of Structural Separation

We now provide motivation for why training separate classifiers for different subgroups can reduce worst-case risk, compared to training a single classifier on all data. Within each class c, suppose the data naturally partitions into subgroups ScE (easy) and ScH (hard) that may require different features for optimal classification—corresponding to samples that will be assigned to nodes 2c (easy) and 2c+1 (hard) after the first split. Since each head performs one-vs-all classification, the head at node 2c distinguishes ScE from ScH ∪ S¬c , while the head at node 2c+1 distinguishes ScH from ScE ∪ S¬c , where S¬c denotes all samples from the opposite class. Let G be the hypothesis class of a classifier head, and let RS (g) denote the risk of classifier g on subgroup S. For notational convenience, let RE (·) = RScE (·) and RH (·) = RScH (·). Define the optimal classifiers: ∗ gE = arg min RE (g)

(6)

∗ gH = arg min RH (g)

(7)

g∈G g∈G

Definition 4.3 (Feature Conflict). Feature conflict occurs when the optimal hypotheses for the easy ∗ ∗ and hard subgroups differ: gE ̸= gH . Feature conflict arises when the features minimizing risk on the majority group differ from those required for the minority group. While a single oracle classifier could in principle achieve low risk on both by relying on causal features, simplicity bias drives ERM toward the simpler majority solution ∗ ∗ gE , conflicting with the minority-optimal gH . Theorem 4.4 (Approximation Gap). Under feature conflict, for any single classifier g ∈ G: max RE (g), RH (g) ≥ ∗ ∗ max RE (gE ),RH (gH ) +∆

(8)

where ∆ > 0 when feature conflict exists. Proof. The gap h i ∗ ∗ ∆ = min max RE (g), RH (g) − max RE (gE ), RH (gH ) g∈G

is non-negative by construction. A strictly positive lower bound is established in Appendix I.2 for a gaussian feature model with spurious magnitude µs and core magnitude µc , taking the closed form p ∆ = Φ(−µc ) − Φ − µ2s + µ2c > 0 whenever µs > 0, with Φ the standard normal CDF. The gap is monotone increasing in spurious strength: stronger spurious correlations make structural separation more beneficial. For a Waterbirds-like regime (µs = 2, µc = 1), ∆ ≈ 14.6%. Connection to NCT. NCT addresses the gap of Theorem 4.4 by training separate heads for each partition. By Theorem 4.2, under simplicity bias, the hard heads specialize on training data enriched with minority samples; at inference, the cross-head arg max (Equation (1)) lets the head specialized on similar samples determine the prediction. The realized benefit depends on (i) routing fidelity (Theorem 4.2) and (ii) per-head learning quality. While the framework does not guarantee optimality, it provides a principled mechanism for specialization that parametric reweighting methods lack. 6

Table 1: Main results. WGA and average accuracy (AUROC for ISIC) across five benchmarks. Property columns: T = uses training group labels, V = uses validation group labels, I = inference-time interpretability of the discovered partition. Cells marked ∗ are sourced from prior baseline papers. “–” indicates baselines not adapted to that benchmark. Method

Properties

Waterbirds

CelebA

ISIC (AUROC)

UMNIST

CMNIST

WGA

Avg

WGA

Avg

Overall

Non-Patch

Histopath.

WGA

Avg

WGA

Avg

Group DRO

✓

✗

90.7∗ ±0.4

92.7∗ ±0.4

89.3∗ ±0.9

92.8∗ ±0.1

.933∗ ±.005

.923∗ ±.003

.875∗ ±.004

96.8∗ ±0.4

98.0∗ ±0.3

78.5∗ ±4.5

90.6∗ ±0.1

JTT CnC EIIL DFR

✗ ✗ ✗ ✗

✓ ✓ ✓ ✓

✗ ✗ ✗ ✗

86.7∗ 88.5∗ ±0.3 87.3∗ 92.1∗

93.3∗ 90.9∗ ±0.1 93.1∗ 96.7∗

81.1∗ 88.8∗ ±0.9 81.3∗ 86.9∗

88.0∗ 89.9∗ ±0.5 89.5∗ 91.1∗

.892 ±.005 .951 ±.008 .958 ±.007 .957 ±.006

.862 ±.006 .915 ±.014 .921 ±.013 .922 ±.008

.827 ±.008 .870 ±.025 .877 ±.017 .876 ±.011

92.2 ±1.0 92.8 ±1.8 94.0 ±1.5 97.1 ±0.1

98.5 ±0.2 97.3 ±0.8 97.6 ±0.2 97.8 ±0.0

74.5∗ ±2.4 77.4∗ ±3.0 72.8∗ ±6.8 –

90.2∗ ±0.8 90.9∗ ±0.6 90.7∗ ±0.9 –

ERM GEORGE ExMap NCT (Ours)

✗ ✗ ✗ ✗

✗ ✗ ✗ ✓

63.3∗ ±1.6 76.2∗ ±2.0 92.5∗ 88.0 ±0.9

97.3∗ ±0.1 95.7∗ ±0.5 96.0∗ 92.8 ±0.7

40.3∗ ±2.3 53.7∗ ±1.3 84.4∗ 86.1 ±0.9

95.7∗ ±0.0 94.6∗ ±0.2 91.8∗ 88.4 ±0.9

.957∗ ±.002 .927∗ ±.008 .957 ±.004 .959 ±.001

.922∗ ±.003 .912∗ ±.005 .923 ±.008 .924 ±.002

.875∗ ±.005 .876∗ ±.006 .878 ±.017 .880 ±.002

93.9∗ ±0.6 95.7∗ ±0.6 96.7 ±0.3 93.7 ±2.5

98.7∗ ±0.1 98.1∗ ±0.3 97.6 ±0.1 97.5 ±0.4

0.0∗ ±0.0 76.4∗ ±2.3 – 72.9 ±2.8

20.1∗ ±0.2 89.5∗ ±0.3 – 88.4 ±1.5

Experimental Setup

Datasets. We evaluate on five spurious-correlation benchmarks: Waterbirds [2] (background), CelebFaces Attributes (CelebA) [24] (demographic), International Skin Imaging Collaboration (ISIC) [25, 26] (acquisition artifacts), Undersampled MNIST (UMNIST) [6] (digit identity), and the 5-class Colored MNIST (CMNIST) [27, 8] (color). Appendix A.1 gives the full description. Evaluation Metrics. Our primary metric is WGA—the lowest accuracy across pre-defined subgroups—reported alongside average accuracy. ISIC uses area under the receiver operating characteristic curve (AUROC) following the GEORGE [6] evaluation protocol. Implementation Details. Backbones are ImageNet-pretrained [28] ResNet-50 [29] for Waterbirds, CelebA, and ISIC, and LeNet-5 [30] for UMNIST and CMNIST. Models train for up to three iterations with AdamW [31], using two phases per iteration (Section 3.4) and the depth criterion of Section 3.5. Hyperparameters and per-dataset node-sampling strategies are tuned via Optuna [32] against validation pWGA2 ; full configuration in Appendices B.3 and G. Baselines. We compare against eight baselines spanning three supervision tiers: group-supervised (group DRO [2]); validation-group-supervised (JTT [3], CnC [8], EIIL [4], DFR [5]); and unsupervised (ERM, GEORGE [6], ExMap [7]). Where published numbers are unavailable for ISIC or UMNIST, we adapt the official implementations under the same Optuna budget as NCT.

Results and Analysis

We evaluate NCT through three lenses: (1) robustness against spurious correlations compared to baselines, (2) quantitative interpretability through node-level sample alignment, and (3) qualitative interpretability through feature attribution. Auxiliary diagnostics — pseudo-WGA proxy quality (Appendix F.3), cross-head calibration (Appendix D), training-time comparison (Appendix E), and component ablations (Appendix H) — are deferred to the appendix. 6.1

Quantitative Performance: Mitigating Spurious Correlations

Table 1 reports the full comparison (per-iteration breakdowns and per-seed depth selection are in Appendix F). On ISIC, all three of NCT’s AUROC numbers lead the table, narrowly above Group DRO despite using no training group labels. Within the unsupervised tier, NCT improves on ERM and GEORGE on almost every benchmark and is competitive with ExMap on Waterbirds and CelebA. Against validation-supervised methods, NCT matches or narrowly exceeds them on Waterbirds, is comparable on CMNIST and CelebA, but trails on UMNIST. Overall, NCT is competitive with the strongest baselines on most benchmarks in unsupervised tier while remaining the only unsupervised entry that exposes the discovered partition structurally. 6.2

Structural Interpretability: Latent Group Discovery

Unlike robust optimization methods that implicitly reweight samples, NCT generates an explicit partition of the data. We quantify this via the iteration-2 capture rate: the percentage of a ground-truth 7

Waterbirds

CelebA

40 20

LB = Landbird WB = Waterbird

rd lig

Ea sy

CMNIST Subgroup

Digit 0-4 Digit 5-9 Non-8 (Easy) Digit 5-9 Digit-8 (Hard)

80 60 40 20

Subgroup

100 % of Subgroup Population Captured

lig

nig Be

UMNIST

asy nE

ard

asy

Ha rd

rd Ha WB

sy Ea WB

Ha rd

Malig = Malignant

NB = Not Blond B = Blond

Benign With Patch (Easy) Benign No Patch (Hard) Malignant No Patch

nH ard

Subgroup

100

nig

Not Blond Male (Easy) Not Blond Female (Hard) Blond Female (Easy) Blond Male (Hard)

ISIC Subgroup

100 % of Subgroup Population Captured

% of Subgroup Population Captured

LB on Land (Easy) LB on Water (Hard) WB on Water (Easy) WB on Land (Hard)

% of Subgroup Population Captured

Subgroup

100

Class 0-1 Class 2-3 Class 4-5 Class 6-7 Class 8-9

80 60 40 20

Left bar: aligned color Right stack: 4 mismatched colors

Ha rd 8-9

sy Ea 8-9

Ha rd 6-7

sy Ea 6-7

Ha rd 4-5

Ea 4-5

2-3

Ha rd

sy Ea 2-3

Ha rd 0-1

sy Ea 0-1

Ha rd 5-9

sy Ea 5-9

0-4

Ea 0-4

Ha rd

Figure 3: Iteration-2 capture rates. Bar height = % of a subgroup’s population routed to each leaf (mean ± std, five seeds). For CMNIST, each leaf shows the own-class matching-colour subgroup (left bar) and the four own-class mismatched-colour subgroups stacked (right bar). subgroup’s population routed to each leaf, averaged across five seeds. Across all datasets, majority subgroups concentrate in easy leaves while minority subgroups concentrate in hard leaves (Figure 3). Waterbirds. The hard landbird leaf captures 82.2% of the minority landbird-on-water, while the easy leaf retains 92.3% of the majority landbird-on-land. The waterbird branch mirrors this: hard captures 74.4% of waterbird-on-land, easy retains 89.5% of waterbird-on-water. CelebA. The blond-male concentrates in its hard leaf (73.3%), and the not-blond hard leaf catches 28.7% of not-blond-female. Easy leaves retain 85.1% of not-blond-male and 74.2% of blond-female. ISIC. The easy benign leaf captures 99.7% of benign-with-patch, exploiting the color-patch shortcut. The hard benign leaf isolates 47.3% of the benign-no-patch subgroup, forcing reliance on lesion-based features rather than the artifact. UMNIST. The hard 5–9 leaf captures 71.4% of the undersampled digit 8, separating it from the majority digit 5–9 (non-8) population which the easy leaf retains at 86.4%. CMNIST. The same pattern carries to the multi-class case. Within each class, samples whose color matches the class consolidate in its easy leaf (98.7%), while samples whose color does not match concentrate in the hard leaf of the true class (83.7%). The four off-class colors contribute roughly equally to each hard-leaf stack (Figure 3), so the routing isolates the color-mismatched samples regardless of which off-class color they carry. For datasets with only three inherent subgroups (ISIC, UMNIST), iteration 2 produces four leaves but one remains sparsely populated rather than artificially splitting a natural subgroup; the sparse leaf absorbs residual samples from the majority without disrupting the primary split. We further analyze routing in Iteration 3 (eight leaves) in Appendix J. The hierarchical pattern persists at finer granularity: minority subgroups continue to concentrate in their class’s hard branch (Waterbirds 75.1% landbird-on-water, CelebA 60.3% blond-male, UMNIST 44.7% digit-8), and on Waterbirds the sparse-node merging mechanism (Section 3.5) activates to drop unused leaves. 6.3

Qualitative Interpretability: Attribution Analysis

The capture-rate analysis shows that hard leaves isolate minority subgroups, but do they actually rely on different features? We use LayerGradCAM [33] to inspect what image regions each leaf attends to; if the structural separation is meaningful, easy leaves should rely on the spurious cue while hard leaves should attend to the semantic class feature. Figure 4 shows five randomly selected samples per leaf on Waterbirds, ISIC, and CelebA. 8

Easy Waterbird

Hard Waterbird

Easy Benign

Hard Benign

Easy Blonde

Hard Blonde

Figure 4: Iteration-2 LayerGradCAM attributions. Easy leaves (left) localize on the spurious cue—background on Waterbirds, the color-patch artifact on ISIC, the face on CelebA. Hard leaves (right) shift attention to the semantic class feature: the bird body, the lesion, and the hair respectively. Waterbirds. The easy waterbird leaf attends to the water and surrounding scene—the lake surface and the wake behind the bird—rather than the bird itself. The hard leaf localizes tightly on the bird’s body across both water and forest backgrounds. ISIC. The easy benign leaf locks onto the colored skin-marker ring that frames many benign lesions in training, while the central lesion receives little attention. The hard benign leaf inverts this: it centers on the lesion—its boundary and pigmented body—and ignores the ring artifact. CelebA. The easy blond leaf attends to the lower face—eyes, cheeks, mouth—using it as a gender proxy. The hard blond leaf shifts upward to the hair, the actual class signal; in the samples shown, this includes a cap occluding the hair and faces where the gender shortcut would mislead. We omit UMNIST and CMNIST from this figure: their spurious cues are global properties like digit identity/color rather than localizable regions, so attribution maps cannot meaningfully highlight them.

Conclusion

Existing methods for spurious correlations adjust parameters to improve worst-group accuracy but leave the classifier opaque about the latent groups it has learned. NCT takes a different approach: training difficulty becomes a routing signal that partitions samples into easy and hard branches over successive iterations, with specialized heads at every node. The resulting tree serves as both classifier and partition—each leaf encodes a predicted class together with the difficulty path that produced it, and hard branches consistently isolate minority subgroups across binary and multi-class spurious correlations. Across five benchmarks, this framework delivers worst-group accuracy competitive with strong baselines—leading on ISIC across all supervision tiers—while making the latent group structure visible at inference rather than hidden in the parameters. The structural guarantees come with two conditions. First, routing relies on simplicity bias (Theorem 4.2): if minority features are themselves easy to learn, difficulty-based routing fails to isolate them. Second, the depth-selection rule depends on pWGA2 tracking true WGA, and may stop early when the proxy collapses faster than the true metric. Severe class imbalance can potentially amplify both, since a small misrouting rate can let an oversized easy class outnumber the genuine minority in a hard branch. We identify two natural directions for future work. First, extending NCT to language settings— where spurious correlations arise from lexical artifacts or demographic markers—would test whether simplicity bias provides the same routing signal under different feature geometry. Second, the routing rule can admit richer signals beyond correctness, opening a natural research direction: gradientbased, attribution-based, or representation-disagreement criteria could partition data along axes that correctness alone may miss. We see NCT as an evidence that the latent group structure of a dataset can be recovered as the architecture itself, discovered without supervision and visible at inference. 9

References [1] Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard S Zemel, Wieland Brendel, Matthias Bethge, and Felix Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2:665–673, 2020. doi: 10.1038/s42256-020-00257-z. [2] Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations, 2020. [3] Shuxian Liu, Kai-Wei Chang, Andrew L. Maas, Ran Tao, and Tatsunori B. Hashimoto. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, 2021. [4] Elliot Creager, Jörn-Henrik Jacobsen, and Richard Zemel. Environment inference for invariant learning. In International Conference on Machine Learning, 2021. [5] Polina Kirichenko, Pavel Izmailov, and Andrew Gordon Wilson. Last layer re-training is sufficient for robustness to spurious correlations. In International Conference on Learning Representations (ICLR), 2022. [6] Nimit S. Sohoni, Jared A. Dunnmon, Geoffrey Angus, Albert Gu, and Christopher Ré. No subclass left behind: Fine-grained robustness in coarse-grained classification problems. In Advances in Neural Information Processing Systems, volume 33, pages 19339–19352, 2020. [7] Rwiddhi Chakraborty, Adrian Sletten, and Michael C Kampffmeyer. Exmap: Leveraging explainability heatmaps for unsupervised group robustness to spurious correlations. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2024. [8] Michael Zhang, Nimit S Sohoni, Hongyang R Zhang, Chelsea Finn, and Christopher Ré. Correct-n-contrast: A contrastive approach for improving robustness to spurious correlations. In International Conference on Machine Learning (ICML), pages 26484–26516. PMLR, 2022. [9] Mohammad Pezeshki, Diane Bouchacourt, Mark Ibrahim, Nicolas Ballas, Pascal Vincent, and David Lopez-Paz. Discovering environments with xrm. In International Conference on Machine Learning, 2024. [10] Junhyun Nam, Hyuntak Cha, Sungsoo Ahn, Jaeho Lee, and Jinwoo Shin. Learning from failure: Training debiased classifier from biased classifier. In Advances in Neural Information Processing Systems, 2020. [11] Inwoo Hwang, Sangjun Lee, Yunhyeok Kwak, Seong Joon Oh, Damien Teney, Jin-Hwa Kim, and Byoung-Tak Zhang. Selecmix: Debiased learning by contradicting-pair sampling. In Advances in Neural Information Processing Systems, 2022. [12] Yujin Han and Difan Zou. Improving group robustness on spurious correlation requires preciser group inference. In International Conference on Machine Learning (ICML), pages 17480– 17504, 2024. [13] Saeid Asgari Taghanaki, Kumar Abhishek, Kenji Kawaguchi, and Amir Azimi. Masktune: Mitigating spurious correlations by forcing to explore. In Advances in Neural Information Processing Systems, 2022. [14] Pang Wei Koh, Thao Nguyen, Yew Siang Tang, Stephen Mussmann, Emma Pierson, Been Kim, and Percy Liang. Concept bottleneck models. In International Conference on Machine Learning (ICML), pages 5338–5348. PMLR, 2020. [15] Md Rifat Arefin, Yan Zhang, Elnaz Barshan, Xiang Yue, Kenji Kawaguchi, and H. Hacihadibadi. Unsupervised concept discovery mitigates spurious correlations. In International Conference on Machine Learning (ICML). PMLR, 2024. [16] Meike Nauta, Ron van Bree, and Christin Seifert. Neural prototype trees for interpretable fine-grained image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 14933–14943, 2021. 10

[17] Alvin Wan, Lisa Dunlap, Daniel Ho, Jihan Yin, Scott Lee, Henry Jin, Suzanne Petryk, Sarah Adel Bargal, and Joseph E. Gonzalez. NBDT: Neural-backed decision tree. In International Conference on Learning Representations (ICLR), 2021. [18] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeffrey Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations (ICLR), 2017. [19] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural Computation, 3(1):79–87, 1991. [20] Ryutaro Tanno, Kai Arulkumaran, Daniel Alexander, Antonio Criminisi, and Aditya Nori. Adaptive neural trees. In International Conference on Machine Learning (ICML), pages 6166–6175. PMLR, 2019. [21] Michael McCloskey and Neal J Cohen. Catastrophic interference in connectionist networks: The sequential learning problem. Psychology of learning and motivation, 24:109–165, 1989. [22] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. In Advances in Neural Information Processing Systems, volume 33, pages 9573–9585, 2020. [23] Katherine Hermann, Ting Chen, and Simon Kornblith. The origins and prevalence of texture bias in convolutional neural networks. In Advances in Neural Information Processing Systems, volume 33, pages 19000–19015, 2020. [24] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), 2015. [25] Noel C. F. Codella, David Gutman, M. Emre Celebi, Brian Helba, Michael A. Marchetti, Stephen W. Dusza, Aadi Kalloo, Konstantinos Liopyris, Nabin Mishra, Harald Kittler, and Allan Halpern. Skin lesion analysis toward melanoma detection: A challenge at the 2017 international symposium on biomedical imaging (isbi), hosted by the international skin imaging collaboration (isic). In 2018 IEEE 15th International Symposium on Biomedical Imaging (ISBI 2018), pages 168–172. IEEE, 2018. [26] Philipp Tschandl, Cliff Rosendahl, and Harald Kittler. The ham10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions. Scientific Data, 5(1): 1–9, 2018. [27] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893, 2019. [28] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A largescale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition, pages 248–255. IEEE, 2009. [29] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770–778, 2016. [30] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998. [31] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019. [32] Takuya Akiba, Shotaro Sano, Toshihiko Yanase, Takeru Ohta, and Masanori Koyama. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 2623–2631, 2019. 11

[33] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In Proceedings of the IEEE international conference on computer vision, pages 618–626, 2017. [34] Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The caltechucsd birds-200-2011 dataset. Technical Report CNS-TR-2011-001, California Institute of Technology, 2011. [35] Bolei Zhou, Agata Lapedriza, Aditya Khosla, Aude Oliva, and Antonio Torralba. Places: A 10 million image database for scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 40(6):1452–1464, 2018.

Appendix Contents A Extended Experimental Setup

A.1 Dataset Details . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

A.2 Preprocessing . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

B Implementation Details

B.1 Architecture . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

B.2 Training Protocol . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

B.3 Hyperparameters . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

B.4 Hyperparameter Search . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

C Baseline Implementations

D Cross-Head Calibration Diagnostics

E Training-Time Comparison

F Depth-Selection Diagnostics

F.1

Per-Iteration Test Performance . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F.2

Validation Proxy Values . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

F.3

Proxy Quality: Pseudo-WGA vs. True-WGA . . . . . . . . . . . . . . . . . . . .

G Node Sampling Strategies

G.1 Definitions. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

G.2 Per-Dataset Assignment. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

H Component Ablations

H.1 Trainable Scope . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

H.2 Head Architecture . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

H.3 Sparse-Node Merging . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

H.4 Auxiliary-Loss Weight . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .

Theoretical Proofs

I.1

Proof of Theorem 4.2 (Minority Enrichment) . . . . . . . . . . . . . . . . . . . .

I.2

Constructive Lower Bound for ∆ . . . . . . . . . . . . . . . . . . . . . . . . . . .

Iteration-3 Routing Analysis

Extended Experimental Setup

A.1

Dataset Details

We evaluate on five benchmark datasets exhibiting spurious correlations. Table 2 summarizes the dataset characteristics and split sizes, and Table 3 provides the training set group distributions. Table 2: Dataset Summary. Benchmark datasets with task descriptions, spurious attributes, and split sizes. Dataset

Task

Spurious Attribute

Waterbirds CelebA ISIC UMNIST CMNIST

Bird type (binary) Hair color (binary) Lesion diagnosis (binary) Digit range (binary) Digit pair (5-class)

Background Gender Colored patches Digit identity Color

Train

Val

Test

4,795 162,770 19,124 43,542 54,000

1,199 19,867 2,390 12,000 6,000

5,794 19,962 2,392 10,000 10,000

Waterbirds. Constructed by superimposing bird images from CUB-200-2011 [34] onto backgrounds from Places [35]. The task is binary classification of bird type (waterbird vs. landbird), where background (water vs. land) is spuriously correlated with the label. In training, 95% of waterbirds appear on water and 95% of landbirds on land. The validation and test sets are balanced across groups. CelebA. Binary hair color classification (blond vs. non-blond) on the CelebA dataset [24], where gender serves as the spurious attribute. Blond males constitute only 1,387 samples (0.85%) in training, making this the critical minority group. ISIC. Dermoscopic images from ISIC 2019 [26, 25] for binary classification (benign vs. malignant). Following Sohoni et al. [6], colored patches (acquisition artifacts) that appear predominantly on benign lesions serve as the spurious attribute. UMNIST. A modified MNIST [6] with binary classification: digits 0–4 vs. 5–9. Digit ‘8’ is undersampled to 5% of its original frequency in training, creating a spurious association between digit identity and class membership. CMNIST. A 5-class colored MNIST variant following Zhang et al. [8]: digits are grouped into pairs {0, 1}, {2, 3}, {4, 5}, {6, 7}, {8, 9} and each class is dyed predominantly with a fixed color. The training set has 99.5% color–class correlation; the validation and test sets are uncorrelated. The full 25 = 5 × 5 class-by-color grid serves as the ground-truth group structure, giving a multi-class spurious-correlation setting absent from the other four benchmarks. Aligned subgroups average ≈ 10,746 samples each; conflicting subgroups average ≈ 13 samples each. A.2

Preprocessing

Per-dataset preprocessing is summarized in Table 4. We adopt the published preprocessing pipelines of the corresponding baselines verbatim and do not alter resolution, augmentation, or normalization. Waterbirds follows Sagawa et al. [2]; CelebA follows the ImageNet-generic recipe used by Kirichenko et al. [5]; ISIC and UMNIST follow Sohoni et al. [6]; CMNIST follows Zhang et al. [8].

Implementation Details

B.1

Architecture

For Waterbirds, CelebA, and ISIC, we use a ResNet-50 [29] backbone initialized with ImageNet [28] pretrained weights. For UMNIST we use a LeNet-4 [30] backbone, and for CMNIST a LeNet-5 backbone, both trained from scratch. At iteration 1 each class head is an MLP block (Waterbirds, CelebA, ISIC) or a Linear classifier (UMNIST, CMNIST). At iterations 2–3 every 14

Table 3: Group Distributions. Training set composition by subgroup. Minority groups are italicized. For CMNIST, the 5 aligned and 20 conflicting subgroups are aggregated for compactness. Dataset

Group

Count

Waterbirds

Landbird on Land Landbird on Water Waterbird on Land Waterbird on Water

3,498 184 56 1,057

72.9 3.8 1.2 22.0

CelebA

Not Blond Female Not Blond Male Blond Female Blond Male

71,629 66,874 22,880 1,387

44.0 41.1 14.1 0.9

ISIC

Benign No Patch Benign With Patch Malignant

9,861 7,420 1,843

51.6 38.8 9.6

UMNIST

Digits 0–4 Digits 5–9 (non-8) Digit 8

24,449 18,859 234

56.2 43.3 0.5

CMNIST

Class–color aligned (5 groups) Class–color conflicting (20 groups)

53,730 270

99.5 0.5

Table 4: Preprocessing pipelines. All settings are taken from the cited baselines without modification. “RRC” = RandomResizedCrop; “RC(s, p)” = RandomCrop of size s with padding p; “CC” = CenterCrop; “HFlip”/“VFlip” = horizontal/vertical flip with p = 0.5. Dataset

Res.

Train transform

Eval transform

Normalization

Waterbirds CelebA ISIC UMNIST CMNIST

224 224 224 32 32

RRC(224), scale [0.7, 1.0], ratio [0.75, 1.33]; HFlip RRC(224), scale [0.7, 1.0], ratio [0.75, 1.33]; HFlip RRC(224), scale [0.7, 1.0], ratio [0.75, 1.33]; HFlip, VFlip RC(28, 4) → Resize 32; HFlip Resize 40 → RC(32, 0); no flip

Resize 256 → CC 224 Resize 256 → CC 224 Resize 256 → CC 224 Resize 32 Resize 40 → CC 32

ImageNet ImageNet ISIC-specific MNIST (0.131, 0.308) (0.5, 0.5)

dataset uses the asymmetric-head configuration of Section 3.2: easy children are linear and hard children use a Linear→ReLU→Dropout block. The hidden block is wrapped by an input batch-norm and a hidden batch-norm; iteration-1 heads omit these batch-norm layers (they are replaced by identity), so the additional normalization only takes effect once the hierarchy starts to grow. B.2

Training Protocol

Each iteration consists of two phases (Section 3.4); the per-iteration Phase 1 fraction is dataset-specific (Table 5). Iteration 1 trains for only 1–3 epochs — its purpose is hard-sample identification, not convergence, mirroring the JTT-style identification stage [3]. Iterations 2+ train for 50–100 epochs with early stopping. Checkpoint selection uses validation loss at iteration 1 (the routing partition does not yet exist) and pWGA2 thereafter. B.3

Hyperparameters

Table 5 lists the dataset-specific hyperparameters we tuned per benchmark. Choices held constant across all five datasets are: one-vs-all BCE loss, asymmetric heads from iteration 2 onward (linear easy child, linear–ReLU–dropout hard child), sparse-merge threshold mmin = 20, and depth-selection tolerance zα = 1.96. The optimizer is AdamW for every dataset except CelebA, where SGD with momentum 0.9 is used. B.4

Hyperparameter Search

Hyperparameter values in Table 5 were selected with Optuna [32] using the Tree-structured Parzen Estimator (TPE) sampler over 50–100 trials per dataset, with 20 warm-up trials of random sampling 15

Table 5: Dataset-specific hyperparameters. Constants held across all datasets (loss, asymmetric head structure, mmin , zα ) are listed in the surrounding text; backbone and head architecture choices are described in Appendix B.1. Parameter

Waterbirds

CelebA

ISIC

UMNIST

CMNIST

Head hidden dim Head dropout

32 0.3

128 0.3

96 0

8 0

64 0.2

Backbone LR Head LR LR decay factor Weight decay

2.7 × 10−5 1.3 × 10−5 2.5 2.4 × 10−3

1 × 10−5 1 × 10−4 1.5 0.1

8 × 10−5 1 × 10−4 2.5 0.1

5 × 10−3 5 × 10−4 1.5 0.1

2 × 10−3 1 × 10−5 1.0 5 × 10−5

Batch size Epochs (iter 1/2/3) Phase 1 ratio (iter 1/2/3) Scheduler Early stop patience

128 1/100/100 0.0/0.3/0.3 Plateau 15

64 1/50/50 1.0/0.3/0.3 Plateau 10

256 1/50/50 1.0/0.5/0.5 Step 15

64 64 3/50/50 2/50/50 0.0/0.5/0.5 0.0/0.2/0.7 Plateau Plateau 15 5

Sampling strategy Aux loss weight (λaux ) Class weight cap

Class weights Downsample Class weights 1.0 2.0 1.0 40 40 40

Geomean 0.5 40

Geomean 1.0 40

before switching to TPE. The optimization objective was pWGA2 on the validation set (Section 3.5), computed using routing-derived pseudo-labels rather than ground-truth subgroup annotations. The search ranged over backbone and head learning rates (10−6 to 10−2 ), LR decay factor (1.0 to 3.0), weight decay (10−5 to 2.0), head hidden dimension (16 to 128), head dropout (0.0 to 0.5), Phase 1 training ratios, auxiliary loss weight (0.0 to 2.0) and sampling strategy ({class weights, downsample, geomean}).

Baseline Implementations

Group DRO [2]. All Group DRO numbers in Table 1 are taken directly from prior baseline papers: Waterbirds and CelebA from Sagawa et al. [2], ISIC and UMNIST from Sohoni et al. [6], and CMNIST from Zhang et al. [8]. We do not retrain Group DRO ourselves on any dataset. GEORGE [6]. GEORGE clusters the ERM feature space into pseudo-subgroups, trains Group DRO on those clusters, and selects the best checkpoint by worst-cluster validation accuracy— no ground-truth group labels are accessed at any stage. We do not retrain GEORGE ourselves: Waterbirds, CelebA, ISIC, and UMNIST numbers are taken from Sohoni et al. [6] and CMNIST numbers from Zhang et al. [8]. JTT [3]. Waterbirds and CelebA numbers are taken from Liu et al. [3]; CMNIST numbers are taken from Zhang et al. [8]. For ISIC and UMNIST we extend the official JTT codebase. The two JTT-specific knobs—identification epochs Tup and upweight factor λup —are tuned via Optuna against validation WGA, yielding (Tup , λup ) = (1, 50) on ISIC and (1, 20) on UMNIST. Model selection uses validation worst-group accuracy. CnC [8]. Waterbirds, CelebA, and CMNIST numbers are taken from Zhang et al. [8]. For ISIC and UMNIST we use the official release and tune the stage-2 contrastive parameters—number of anchors / positives / negatives / easy-negatives, contrastive weight, temperature, and contrastive batch factor—via Optuna. The tuned configuration converges to 17 samples per role, contrastive weight 0.75, temperature 0.1, and batch factor 32. DFR [5]. Waterbirds and CelebA numbers are taken from Kirichenko et al. [5]. For ISIC and UMNIST we use the paper’s main variant (DFRVal Tr ): the ERM-trained backbone is frozen and a logistic head is retrained on a group-balanced validation subset, with 20 retrains at evaluation. The inverse regularisation strength C is tuned via Optuna; class weights are not tuned, following the paper’s main-variant recipe. We do not report CMNIST: the released variant is binary, and the multi-class 16

adaptation requires non-trivial changes to the group-balanced retraining objective that we leave to future work. EIIL [4]. Waterbirds and CelebA numbers are taken from Creager et al. [4]; CMNIST numbers are taken from Zhang et al. [8]. For ISIC and UMNIST we use the official two-stage pipeline. Stage-1 environment-inference uses the reference defaults (2 environments, environment lr 10−3 , 10 000 inference steps). Stage-2 trains a Group DRO learner on the inferred environments (η = 0.01, patience 10); the remaining stage-2 hyperparameters are tuned via Optuna against validation worst-group accuracy. ExMap [7]. Waterbirds and CelebA numbers are taken from Chakraborty et al. [7]. For ISIC and UMNIST we use the global G-ExMap variant of the released pipeline: spectral clustering of Layer-wise Relevance Propagation (LRP) attribution maps with the cluster count chosen automatically by the eigengap heuristic and LRP γ = 4, followed by the authors’ DFR-style last-layer retraining with 20 retrains at evaluation. The LRP per-batch size is tuned via Optuna and converges to 8 on ISIC and 512 on UMNIST. We do not report CMNIST: ExMap inherits DFR’s binary group-balancing for the retraining stage, and a faithful multi-class extension is non-trivial and left to future work.

Cross-Head Calibration Diagnostics

Because every NCT head is trained with independent BCE, the arg max inference rule (Equation (1)) implicitly assumes that head logits are on comparable scales. Three implicit factors push them in that direction: the uniform binary BCE target, the shared parent-layer representation, and the per-iteration batch-norm modules introduced in Appendix B.1. We audit the resulting calibration with two diagnostics: the per-head logit mean/std across all test samples, and the decision margin (winning logit minus second-best logit) at each predicted leaf. Per-head logit statistics. Table 6 reports the mean and standard deviation of each head’s logit distribution on the test set, averaged across the five seeds. At iteration 1 the two class-head means are near-antisymmetric (a consequence of BCE with flipped binary targets). At iteration 2 the means diverge, but in every case the per-head standard deviations are comparable to or larger than the typical mean separation, so the logit distributions overlap and argmax can still route samples to the head that fits them best. Table 6: Per-head logit mean ± standard deviation on the test set, averaged across five seeds. Heads that were merged by the sparse-node rule and never selected are omitted. CMNIST is summarized in prose because its 5 iter-1 heads and 10 iter-2 heads do not fit the column layout.

Iter 1 Dataset Waterbirds CelebA ISIC UMNIST

Iter 2

+0.27 ±1.01 +0.39 ±0.60 +0.80 ±1.44 +0.89 ±6.21

−0.16 ±0.90 −0.39 ±0.62 −0.81 ±1.42 −0.91 ±6.22

−1.09 ±2.06 +1.17 ±2.02 +0.47 ±5.82 −2.50 ±6.24

+0.46 ±1.37 −0.39 ±0.94 −2.75 ±2.99 −3.81 ±2.83

−0.70 ±1.55 −1.27 ±1.39 −7.26 ±4.64 −4.57 ±6.29

−0.39 ±1.13 −0.65 ±0.81 −5.01 ±1.96 −4.45 ±3.99

For CMNIST (5-class, 5 heads at iter 1 and 10 at iter 2), the iter-1 head means span [−5.49, −4.05] with an average within-head std of 5.59, and the iter-2 head means span [−1.21, +0.63] with an average within-head std of 4.00. The per-head std is several times larger than the spread of means at every iteration, so the same overlap argument applies. Decision margins. Table 7 reports the decision margin (winning logit minus runner-up) per predicted leaf, averaged across five seeds. All margins are strictly positive, confirming that argmax is decisive even when per-head means differ. The smallest binary-task margin is 0.82 (Waterbirds iter-2 H0); for CMNIST the minimum mean margin across the 5 iter-1 heads is 10.95, and the minimum across the 10 iter-2 heads is 7.71. These two diagnostics together indicate that, despite the absence of explicit calibration, the implicit alignment is sufficient for argmax routing. 17

Table 7: Mean decision margin ± std per predicted leaf (averaged across five seeds). CMNIST margins are summarized in prose for the same column-layout reason.

Iter 1 Dataset Waterbirds CelebA ISIC UMNIST

Iter 2

1.94 ±0.84 1.33 ±0.71 3.27 ±1.83 11.61 ±6.03

1.46 ±0.85 0.99 ±0.72 1.70 ±1.11 10.34 ±5.80

0.82 ±0.26 2.81 ±1.79 8.93 ±3.27 7.26 ±4.85

3.26 ±1.71 1.34 ±0.97 3.55 ±2.20 1.62 ±1.25

3.04 ±1.37 2.52 ±1.81 4.89 ±3.22 5.85 ±3.59

2.39 ±1.79 0.90 ±0.66 2.26 ±1.85 2.45 ±1.74

Training-Time Comparison

Table 8 reports wall-clock training time per method across the five benchmarks, measured on a single H100 GPU under identical hardware. The two-stage methods that retrain a full robust model on top of ERM sit at the heavy end. GEORGE roughly doubles ERM on every dataset because its GDRO stage runs the same epoch budget. JTT on Waterbirds takes 567 minutes because its Stage 2 upweights the error set by λup =50, inflating the effective training set. ExMap and EIIL on ISIC are dominated by their ResNet-50 ERM stage rather than by the LRP or environment-inference modules. DFR’s last-layer retrain is effectively free, so its reported total is dominated by the ERM stage. NCT’s training cost stays low for three reasons. Iteration 1 runs for only 1–3 epochs on every dataset except CelebA, since its purpose is hard-sample identification rather than convergence [3]. Phase 1 of iterations 2–3 trains only the new head layers (≈ 1% of parameters) and skips the backbone backward pass. The node sampler further shrinks each iteration-2/3 epoch on CelebA (downsampling to the smallest node size) and on UMNIST/CMNIST (geomean sampling); Waterbirds and ISIC use class-weighted BCE on full data, so they benefit only from the first two factors. The combined effect is visible in the table: on CelebA, NCT costs 41.6 minutes against DFR’s 148, and on ISIC it costs 106 minutes against ExMap’s 177.7 and EIIL’s 236. Table 8: Wall-clock training time (minutes, single H100 GPU). Numbers are total end-to-end wall-clock time on the published per-dataset configuration. ERM and Group DRO are paper-sourced and not retrained in our pipeline. “–” indicates baselines not adapted to that benchmark. Method

Waterbirds

CelebA

ISIC

UMNIST

CMNIST

GEORGE JTT CnC DFR EIIL ExMap

112.9 567.4 29.2 19.7 14.3 20.5

337.4 – – 148.2 47.7 70.4

130.3 – 251.8 64.6 236.2 177.7

16.1 23.5 13.5 3.1 6.5 3.5

19.1 5.4 61.7 – 5.4 –

NCT (Ours)

28.1

41.6

105.8

2.6

3.2

Depth-Selection Diagnostics

F.1

Per-Iteration Test Performance

Table 9 gives the per-iteration test WGA and average accuracy across the five benchmarks. The optimal hierarchy depth is dataset-dependent and, on Waterbirds and CMNIST, even seeddependent. The depth-selection rule of Section 3.5 adapts per seed: it splits 4-1 in favor of iteration 2 on CelebA, picks iteration 3 for four of five seeds on ISIC, and splits 3-2 in favor of iteration 2 on Waterbirds, 3-2 in favor of iteration 3 on UMNIST, and 4-1 in favor of iteration 2 on CMNIST, all without consulting ground-truth groups. 18

Table 9: Impact of hierarchy depth. Worst-group accuracy (WGA) and average accuracy at each iteration of the hierarchy. Best WGA and average accuracy are highlighted in bold.

Dataset Waterbirds CelebA

Iteration 1

Iteration 2

Iteration 3

WGA

Avg

WGA

Avg

WGA

Avg

18.6 ±4.3 23.7 ±2.2

59.9 ±1.2 85.3 ±1.2

87.8 ±1.0 86.1 ±0.9

92.7 ±0.6 88.2 ±0.9

86.8 ±1.9 82.0 ±5.4

92.7 ±0.8 90.0 ±1.2

ISIC (Non-Patch) 0.812 ±0.001 0.899 ±0.001 0.916 ±0.004 0.953 ±0.002 0.925 ±0.003 0.960 ±0.002 ISIC (Histopath) 0.702 ±0.001 0.899 ±0.001 0.871 ±0.007 0.953 ±0.002 0.881 ±0.004 0.960 ±0.002 UMNIST CMNIST F.2

80.8 ±8.5 0.0 ±0.0

96.6 ±1.0 20.7 ±1.0

94.4 ±1.1 72.3 ±3.0

97.2 ±0.6 87.6 ±1.5

93.0 ±2.6 71.9 ±2.8

97.5 ±0.4 87.7 ±2.1

Validation Proxy Values

Table 10 reports the validation proxy values (pseudo-WGA, or pseudo-AUROC for ISIC) consumed by the depth-selection rule (Equation (4)). The iter-2-to-iter-3 drop in the validation proxy is largest on CelebA, small on Waterbirds and CMNIST, and reversed on ISIC, mirroring the per-seed selection counts above. Table 10: Validation pseudo-WGA per iteration (mean over five seeds). ISIC reports pseudoAUROC instead of pseudo-WGA, matching its main-text p metric. The depth-selection rule stops at the first iteration whose proxy value falls more than z p⋆ (1 − p⋆ )/nworst below the running best (z = 1.96). Dataset Waterbirds CelebA ISIC (pseudo-AUROC) UMNIST CMNIST

F.3

Iter 2

Iter 3

0.821 ±0.015 0.702 ±0.047 0.937 ±0.011 0.781 ±0.028 0.792 ±0.029

0.793 ±0.043 0.657 ±0.038 0.945 ±0.005 0.778 ±0.024 0.779 ±0.047

Proxy Quality: Pseudo-WGA vs. True-WGA

The depth-selection rule and the within-iteration early-stopping criterion both rely on pWGA2 (Section 3.5), computed from the model’s own routing labels rather than ground-truth groups. Its usefulness therefore rests on whether it tracks true WGA. Table 11 reports the Spearman rank correlation ρ between the validation pWGA2 trajectory and the held-out true WGA across training epochs (iterations 2 and 3 combined). The correlation is high on Waterbirds, CelebA, and CMNIST (ρ = 0.87, 0.92, and 0.90), moderate on ISIC (ρ = 0.71), and weakest on UMNIST (ρ = 0.43). On UMNIST’s lower proxy correlation. The weaker Spearman on UMNIST (ρ = 0.43) reflects mild simplicity bias on this benchmark: ERM alone reaches 93.9% WGA (Table 1), above JTT and CnC, so iteration 1 already classifies many minority digit-8 samples correctly. The hard 5–9 leaf still captures most digit-8 samples at test time (71.4%, Section 6.2), but a fraction is routed to the easy branch instead. This mixing means the hard leaf no longer cleanly corresponds to the minority subgroup, lowering pWGA2 ’s rank correlation with oracle WGA. The method itself remains effective: UMNIST achieves 93.7% test WGA, suggesting pWGA2 is robust to moderate misrouting—it still selects a depth that delivers strong worst-group performance even when its rank fidelity drops.

Node Sampling Strategies

From iteration 2 onwards the leaves are imbalanced by construction (hard children always smaller than their easy siblings), and the natural class distribution of Waterbirds, CelebA, and CMNIST is also imbalanced at iteration 1. We support three rebalancing strategies; this appendix gives the precise definition of each, the per-dataset assignment, and a sensitivity comparison. 19

Table 11: pWGA2 proxy quality. Spearman correlation between the validation pWGA2 trajectory and held-out true WGA across training (mean ± std over five seeds). Values close to 1 indicate the proxy ranks checkpoints the same way the oracle would.

G.1

Dataset

Spearman ρ

Waterbirds CelebA ISIC UMNIST CMNIST

0.87 ±0.07 0.92 ±0.04 0.71 ±0.04 0.43 ±0.21 0.90 ±0.08

Definitions.

Let K (t) be the number of leaves at iteration t and nj the count of training samples assigned to node j. • Class weights: Sample-level loss is multiplied by wj = N/(K (t) · nj ), capped at a per-dataset maximum. Iteration 1 uses uniform weights; from iteration 2 onwards wj is recomputed per epoch from the current pseudo-label assignment. • Downsample: At every epoch we draw a fresh subsample of size minj nj from each node. The epoch length is K (t) · minj nj . Class weighting is disabled because every leaf already contributes equally. Q (t) • Geomean: At every epoch we resample each node to size n̄ = ( j nj )1/K . Majority nodes are downsampled; minority nodes are oversampled with replacement. The geometric mean minimizes the maximum stretch factor across nodes and bounds the variance introduced by oversampling. G.2

Per-Dataset Assignment.

Waterbirds and ISIC. Training sets are small (4,795 and 19,124 samples), and after iteration 2 the rarest hard child of Waterbirds contains on the order of 102 samples. Downsample would collapse the epoch to that count per leaf and starve every node, and geomean still trims the overall sample count. Class-weighted BCE keeps every sample available and produces the most stable WGA on both. CelebA. With 162,770 training samples, even the rarest hard child after iteration 2 retains enough examples for stable optimisation under downsample. Class weighting destabilises training because the imbalance ratio is very large (the largest weight cap is hit on every minority leaf), and geomean introduces variance through the large oversampling factor needed on hard children. Downsample balances batches without inflating any single sample. UMNIST and CMNIST. Totals are large (43,542 and 60,000 samples) but hard children are extremely sparse—digit ‘8’ has 234 training samples on UMNIST, and CMNIST’s bias-conflicting color groups are even smaller. Class weighting produces extreme per-sample weights, while downsample starves all nodes; geomean balances the two, moving every node toward n̄ with bounded oversampling on the small leaves.

Component Ablations

We isolate four components: trainable scope, head architecture, sparse-node merging and the auxiliaryloss weight. For every cell we pick the better of iteration 2 and 3 per seed (test WGA, or overall AUROC for ISIC) and average across seeds. The default-configuration column reuses the paper’s main-table seeds; non-default columns are run on a separate three-seed sweep. 20

H.1

Trainable Scope

Each iteration is normally split into two phases: Phase 1 trains the new heads only with the backbone and parent layers frozen, then Phase 2 fine-tunes the entire network. Table 12 compares this standard schedule against two extremes: head-only, which keeps the backbone and parent layers frozen throughout, and full fine-tuning, which skips Phase 1 entirely. Head-only training collapses on every dataset—the pretrained backbone alone cannot resolve the hard children. Full fine-tuning is competitive with the standard schedule on Waterbirds, CelebA, UMNIST, and ISIC, where the dataset’s iter-1 Phase 1 ratio is already near zero, but loses 10 pp on CMNIST. Table 12: Trainable-scope ablation. Standard two-phase training vs. head-only and full fine-tuning.

H.2

Dataset

Standard

Head-Only

Full FT

Waterbirds (WGA) CelebA (WGA) CMNIST (WGA) UMNIST (WGA) ISIC (Overall AUROC)

88.0 ±0.9 86.1 ±0.9 72.9 ±2.8 93.7 ±2.5 .959 ±.001

42.3 ±4.8 45.2 ±3.6 0.8 ±1.1 57.6 ±12.3 .900 ±.004

87.2 ±2.9 85.2 ±0.9 62.8 ±18.1 90.1 ±8.0 .949 ±.012

Head Architecture

At iteration 2 and beyond, easy and hard children of a split can use different head architectures. Table 13 compares the asymmetric default (linear easy, MLP hard) with three alternatives: both children linear, both children MLP, and the reversed asymmetry (MLP easy, linear hard). The asymmetric configuration is the strongest choice on CelebA and CMNIST, and within seed noise of the best alternative on Waterbirds, UMNIST, and ISIC. The reversed asymmetry consistently underperforms or matches the standard order, confirming that the additional capacity belongs on the hard branch where the conflicting features live, not on the easy branch where the parent has already done most of the work. We retain the default for its principled motivation: easy children classify samples the parent already gets right and need little capacity, while hard children carry the conflicting features that motivate a richer block. Table 13: Head-architecture ablation. Linear = both children linear; MLP = both children MLP; Asymmetric = linear easy / MLP hard (default); Reversed = MLP easy / linear hard.

H.3

Dataset

Linear

MLP

Asymmetric

Reversed

Waterbirds (WGA) CelebA (WGA) CMNIST (WGA) UMNIST (WGA) ISIC (Overall AUROC)

86.0 ±0.7 83.8 ±2.3 59.0 ±20.4 92.0 ±4.8 .955 ±.003

88.6 ±0.4 85.3 ±1.4 68.2 ±3.8 92.9 ±3.6 .959 ±.002

88.0 ±0.9 86.1 ±0.9 72.9 ±2.8 93.7 ±2.5 .959 ±.001

87.8 ±1.4 84.2 ±2.7 66.2 ±8.5 83.4 ±4.4 .957 ±.005

Sparse-Node Merging

Hard children of an iteration-2 split can inherit very few samples. The sparse-node merging rule folds any hard child with fewer than mmin training samples into its easy sibling. Table 14 compares disabling the rule (mmin = 0) against the default mmin , separately at iteration 2 and iteration 3. The threshold only ever fires on Waterbirds at iteration 3, where it converts a 0.7 pp regression into a small recovery and reduces seed variance. On every other dataset the iteration-2 and iteration-3 hard children stay above the cutoff (the per-dataset mmin = 20 is set conservatively so the rule fires only when a hard branch would otherwise be optimised on a handful of samples), so the two columns are identical by construction. We retain the rule as a safety net for deeper hierarchies. 21

Table 14: Sparse-node merging ablation. “Disabled” = mmin = 0. “Default” uses the per-dataset mmin from the main results. Values are reported separately at iteration 2 and iteration 3 (3-seed mean ± std). Iteration 2

Iteration 3

Dataset

Disabled

Default

Disabled

Default

Waterbirds (WGA) CelebA (WGA) CMNIST (WGA) UMNIST (WGA) ISIC (Overall AUROC)

88.4 ±0.4 86.1 ±0.9 54.5 ±30.9 94.3 ±1.5 .950 ±.005

85.5 ±2.7 82.0 ±5.4 58.3 ±19.0 93.5 ±3.4 .958 ±.004

86.2 ±2.1 82.0 ±5.4 58.3 ±19.0 93.5 ±3.4 .958 ±.004

Table 15: Auxiliary-loss coefficient sweep. WGA (or Overall AUROC for ISIC) as λaux varies.

λaux

0.0

0.3

0.5

0.7

1.0

2.0

Waterbirds (WGA) 86.4 ±0.5 88.6 ±1.5 87.7 ±2.1 88.7 ±1.1 88.0 ±0.9 88.4 ±1.6 CelebA (WGA) 68.5 ±3.4 74.3 ±4.0 76.8 ±3.9 78.7 ±4.2 81.5 ±2.8 86.1 ±0.9 CMNIST (WGA) 50.6 ±27.9 53.3 ±27.2 61.2 ±16.0 53.0 ±28.5 72.9 ±2.8 65.7 ±10.5 UMNIST (WGA) 89.1 ±7.2 93.5 ±2.3 93.7 ±2.5 92.9 ±4.0 95.4 ±1.8 92.5 ±5.0 ISIC (Overall AUROC) .955 ±.004 .955 ±.002 .959 ±.003 .958 ±.001 .959 ±.001 .957 ±.001 H.4

Auxiliary-Loss Weight

The auxiliary term λaux L(t−1) balances child specialization against parent stability. Table 15 sweeps λaux ∈ {0, 0.3, 0.5, 0.7, 1.0, 2.0} on all five datasets. Removing the auxiliary loss costs 1.6 pp WGA on Waterbirds, 17.6 pp on CelebA, and 22 pp on CMNIST; UMNIST and ISIC are within seed noise. The empirical optimum is dataset-dependent and does not always match the chosen default, but every dataset benefits from at least one nonzero coefficient.

Theoretical Proofs

I.1

Proof of Theorem 4.2 (Minority Enrichment)

We give the full derivation of Theorem 4.2 deferred from Section 4. (1)

Proof. Since ℓi = yi , we write y for the true label throughout. Under simplicity bias (Assumption 4.1), the model predicts ℓ̂(1) ≈ a. We derive the probability of error for each group: • Majority groups (a = y): The model predicts ℓ̂(1) ≈ a = y, which is correct. An error occurs only with probability ϵ: P (ℓ̂(1) ̸= y | a = y) = P (ℓ̂(1) ̸= a) = ϵ • Minority groups (a ̸= y): The model predicts ℓ̂(1) ≈ a ̸= y, which is incorrect. An error occurs with high probability 1 − ϵ: P (ℓ̂(1) ̸= y | a ̸= y) = P (ℓ̂(1) = a) = 1 − ϵ We apply Bayes’ theorem to find the posterior probability that a misclassified sample belongs to a minority group: P (a ̸= y | ℓ̂(1) ̸= y) =

P (ℓ̂(1) ̸= y | a ̸= y)P (a ̸= y)

P (ℓ̂(1) ̸= y) (1 − ϵ)(1 − ρ) = ϵρ + (1 − ϵ)(1 − ρ) 22

I.2

Constructive Lower Bound for ∆

We instantiate the approximation gap of Theorem 4.4 on a tractable Gaussian feature model and derive a closed-form lower bound that grows with the spurious correlation strength. I.2.1

Setup

Let y, a ∈ {0, 1} denote the class label and spurious attribute, respectively, and let Φ denote the standard normal CDF. Each input has features z = (zs , zc ) drawn according to zs ∼ N (2a − 1)µs , 1 , (9) zc ∼ N (2y − 1)µc , 1 , (10) where µs , µc > 0 control the magnitudes of the spurious and core features, and the map (2a − 1) sends {0, 1} → {−1, +1}. We assume equal class priors and that zs , zc are conditionally independent given (y, a), so the class-conditional covariance is the identity. A linear classifier w = (w1 , w2 ) predicts the label via sign(w1 zs + w2 zc ). I.2.2

Specialist Classifiers

Easy subgroup (a = y). Both features align with the label. For two Gaussian classes with equal priors and shared identity covariance, the Bayes-optimal linear direction is the difference of class means (the standard LDA result, with the shared covariance eliminating the quadratic term). The two class-conditional means on the easy subgroup are (µs , µc ) for y = 1 and (−µs , −µc ) for y = 0, so (µs , µc ) − (−µs , −µc ) = 2(µs , µc ) ∝ (µs , µc ), giving the optimal weight vector wE ∝ (µs , µc ). Hard subgroup (a ̸= y). The spurious feature flips sign relative to the label: when y = 1 we have a = 0, so zs ∼ N (−µs , 1) instead of N (µs , 1). The class-conditional means become (−µs , µc ) for y = 1 and (µs , −µc ) for y = 0, so by the same LDA argument the optimal weight vector is wH ∝ (−µs , µc ). Specialist risk. Take wE = (µs , µc ). On the easy subgroup with y = 1, the score wE · z = µs zs + µc zc is a linear combination of independent Gaussians, hence Gaussian itself with E[wE · z] = µ2s + µ2c ,

Var(wE · z) = µ2s + µ2c .

A misclassification occurs when this score is negative, so ⋆ RE = P wE · z < 0 y = 1

0 − (µ2s + µ2c ) p = Φ µ2s + µ2c

p = Φ − µ2s + µ2c .

(11)

The case y = 0 is symmetric and yields the same value. For the hard subgroup, wH = (−µs , µc ) pairs with zs whose mean is negated relative to the easy case, so wH · z has the same Gaussian ⋆ ⋆ distribution as wE · z on the easy subgroup, and RH = RE . I.2.3

Single Classifier

A single linear classifier shares (w1 , w2 ) across both subgroups, giving w1 µs + w2 µc RE (w) = Φ − , ∥w∥ −w1 µs + w2 µc RH (w) = Φ − . ∥w∥ The worst-group risk is max(RE (w), RH (w)), and we seek the w that minimises it. 23

(12) (13)

The minimax is attained at w1 = 0. The risks (12) and (13) are related by the symmetry w1 7→ −w1 , which swaps RE and RH , so the worst-group risk is symmetric in w1 . This motivates checking w1 = 0, but symmetry alone does not force the optimum onto the axis (a symmetric function may have off-axis minima); we therefore verify the bound directly. Since Φ is monotone increasing, max(RE (w), RH (w)) = Φ(− min(aE , aH )) ,

aE =

w1 µs + w2 µc −w1 µs + w2 µc , aH = . ∥w∥ ∥w∥

Using the identity min(b + c, −b + c) = c − |b| with b = w1 µs and c = w2 µc , and bounding in three steps: min(aE , aH ) =

(ii) w2 µc − |w1 |µs (i) w µ |w2 |µc p p 2 c ≤ ≤ p 2 2 2 2 2 w1 + w2 w1 + w2 w1 + w22

(iii)

≤ µc ,

(14)

where (i) uses |w1 |µs ≥ 0, (ii) uses w2 ≤ |w2 |, and (iii) uses |w2 | ≤ ∥w∥. Equality holds in (i) iff w1 = 0, in (ii) iff w2 ≥ 0, and in (iii) iff w1 = 0; jointly, equality throughout requires w1 = 0 and w2 > 0 (the strict sign on w2 is needed for ∥w∥ > 0). Resulting risk. From (14) we obtain max(RE (w), RH (w)) ≥ Φ(−µc ) for every w ̸= 0, with equality iff w1 = 0. The single classifier therefore attains worst-group risk Φ(−µc ) and does so only by discarding the spurious feature entirely. I.2.4

Closed-Form Gap

Combining the two cases yields p ∆ = Φ(−µc ) − Φ − µ2s + µ2c > 0

whenever µs > 0,

(15)

p since µ2s + µ2c > µc and Φ is monotone. The gap grows with the spurious magnitude µs : stronger spurious correlations make structural separation more beneficial. Numerical example.

For a Waterbirds-like regime with µs = 2 and µc = 1, √ Φ(−1) ≈ 0.159, Φ(− 5) ≈ 0.013,

so ∆ ≈ 0.146, i.e. a 14.6 percentage-point gap. I.2.5

Interpretation

JTT and similar reweighting methods upweight hard samples within a single classifier; that classifier still faces the gap ∆ in (15). NCT’s separate heads remove this gap. Routing errors do not compound at inference because all leaf heads run in parallel and the prediction is taken via arg max. Hierarchy is particularly helpful because child heads receive parent representations (Section 3.2), enabling progressive specialisation that is more sample-efficient than a flat mixture discovering both classes and subgroups simultaneously.

Iteration-3 Routing Analysis

Iteration 3 splits each iter-2 leaf into easy and hard children, yielding eight leaves per binary task and twenty leaves on CMNIST. Figure 5 reports the per-subgroup capture rate at each iter-3 leaf. The pattern is consistent across datasets: bias-aligned subgroups consolidate in the easy-easy (EE) leaf of their own class, and minority subgroups concentrate in their class’s hard-easy (HE) leaf—the iter-2 hard child further refined by the additional split. Waterbirds. The minority landbird-on-water subgroup lands 75.1 ± 4.5% in LB:HE; the symmetric minority waterbird-on-land lands 70.5 ± 6.5% in WB:HE. Majority subgroups stay in their EE leaves (LB-Land 79.0 ± 16.5%, WB-Water 91.0 ± 1.7%). Two leaves (LB:HH, WB:EH) are unused on every seed because the sparse-node merging rule (Section 3.5) folded their iter-2 children into the easy siblings. 24

Waterbirds

CelebA

E = Easy split H = Hard split Ben = Benign Mal = Malignant

lE Ma

H nH

nE Be

UMNIST

HH NB

HE WB

EE WB

HH LB

EE LB

E = Easy split H = Hard split NB = Not Blond B = Blond

CMNIST Subgroup

100

Subgroup

100

Digit 0-4 Digit 5-9 Non-8 (Easy) Digit 5-9 Digit-8 (Hard)

% of Subgroup Population Captured

E = Easy split H = Hard split LB = Landbird WB = Waterbird

Benign With Patch (Easy) Benign No Patch (Hard) Malignant No Patch

Subgroup

100

Not Blond Male (Easy) Not Blond Female (Hard) Blond Female (Easy) Blond Male (Hard)

% of Subgroup Population Captured

ISIC Subgroup

100

LB on Land (Easy) LB on Water (Hard) WB on Water (Easy) WB on Land (Hard)

Subgroup

100

Class 0-1 Class 2-3 Class 4-5 Class 6-7 Class 8-9

E = Easy split H = Hard split Left bar: aligned color Right stack: 4 mismatched colors

E = Easy split H = Hard split

HH 8-9

HE 8-9

8-9

EE 8-9

HH 6-7

HE 6-7

6-7

EE 6-7

HE 4-5

4-5

EE 4-5

4-5

HH 2-3

HE 2-3

EH 2-3

EE 2-3

0-1

HE 0-1

EE 0-1

0-1

HE 5-9

5-9

EE 5-9

5-9

HH 0-4

HE 0-4

0-4

Figure 5: Iteration-3 capture rates. Bar height = % of a subgroup’s population routed to each leaf (mean ± std, five seeds). Iter-3 leaf labels: EE = easy → easy, EH = easy → hard, HE = hard → easy, HH = hard → hard. For CMNIST, each leaf shows the own-class matching-color subgroup (left bar) and the four own-class mismatched-color subgroups stacked (right bar). CelebA. The minority blond-male subgroup concentrates in B:HE (60.3 ± 2.6%), with another 12.9% in NB:EE. Majority not-blond-male retains 88.5% in NB:EE. Blond-female consolidates in B:EE (72.1%) with 14.1% in B:HE. ISIC. The patch shortcut is locked in: benign-with-patch captures 99.5% of its population in Ben:EE. The conflicting benign-no-patch subgroup splits roughly evenly between Ben:EE (33.3%) and Ben:HE (32.7%). Malignant samples reach a malignant-class leaf at ≈ 79% (Mal:EE 59.7%, Mal:EH 9.3%, Mal:HE 10.1%). UMNIST. The undersampled digit-8 minority concentrates in 5-9:HE (44.7 ± 13.0%); combined with the other 5-9 hard descendants, the 5-9 hard branch holds around 85% of digit-8 samples. Majority subgroups stay in their EE leaves (Digit 0-4 at 80.3%, Digit 5-9 Non-8 at 75.6%). CMNIST. Across all five classes, bias-aligned subgroups consolidate in their own EE leaf at 97.1% on average and stay within their true-class branch at 99.2%. Bias-conflicting subgroups (color ̸= digit class) reach a hard leaf of their true class at 83.0%, averaged across the twenty conflicting subgroups. The iter-2 to iter-3 transition therefore refines the within-class routing without leaking conflicting samples to the wrong digit branch.

Record · ID 299931 · SHA-256 827c5e8fa5622724
Retrieved via Conceptio — every document is proof-bundled with source, license, and retrieval metadata.