PaperlayerView exact revision on arXiv

Source-linked AI summary

Implicit Feature Reweighting for Label-Free Graph Generalization

Tianjun Yao, Chen Xi, Xinyang Li, Lixian Zhong

arXiv:2610.29571v1cs.LG

TL;DR

Graph OOD methods often require class labels to identify predictive substructures, limiting label-free representation learning. SCARCE pairs graph reconstruction with adversarial environment confusion using a given or inferred environment partition, and reports strong label-free OOD results on GOOD-Motif and DrugOOD.

  • Problem

    Many graph OOD approaches require class labels to decide which substructures are predictive, while label-free graph methods remain few and do not establish what their representations retain.

    Source: §1 · pp. 1–2

  • Method

    SCARCE combines VGAE graph reconstruction with adversarial confusion across a given or inferred training-environment partition, without class labels during representation learning.

    Source: §1 · p. 2, §4 · p. 4, §3 · p. 4

  • Results

    17 of 18 DrugOOD settings rank SCARCE best among the compared methods, and it attains state-of-the-art label-free OOD accuracy on GOOD-Motif.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Reconstruction retains core and environment factors, while confusion contracts environment-discriminative directions into an anisotropic ridge penalty for an ℓ2-regularized downstream probe.

    Source: §1 · p. 2

    The representation-learning setting excludes class labels and uses the environment index as its only side information; labels enter only when fitting the evaluation-time linear probe.

    Source: §3 · p. 4

Abstract

from arXiv · show

Graph-level out-of-distribution (OOD) generalization is critical for deploying graph neural networks in real-world applications where distribution shifts are inevitable. Existing approaches typically require class labels during representation learning to decide which substructures are predictive, limiting their applicability. We propose SCARCE, a label-free representation-learning framework that yields distribution-shift-robust graph representations from an environment partition of the training graphs alone, whether that partition comes from dataset metadata or, when none is available, from clustering a self-supervised embedding. Our method operates in two stages: (1) a variational graph autoencoder is warm-started for graph reconstruction; and (2) an adversarial confusion mechanism forces the encoder to produce representations from which a discriminator cannot distinguish environments, while preserving graph-reconstruction fidelity. We provide theoretical analysis showing that environment confusion contracts environment-discriminative directions in the representation space, which, under regularized downstream classifiers, is equivalent to imposing a stronger directional penalty on environment-associated features. This feature-reweighting mechanism reduces reliance on spurious correlations without erasing information, offering a label-free, continuous relaxation of last-layer reweighting methods. Experiments on a controlled synthetic datasets and the DrugOOD datasets show that SCARCE achieves state-of-the-art OOD generalization among label-free methods, using no class labels during representation learning.

1 INTRODUCTION

Graph OOD methods often rely on class labels to identify predictive substructures, while label-free graph methods remain underexplored and do not establish what their representations retain. SCARCE uses an environment partition without class labels to retain predictive factors through reconstruction and contract environment-associated directions, with reported gains on graph OOD benchmarks.

1 INTRODUCTION

  • Graph OOD approaches commonly need class labels to identify predictive substructures, environment labels to identify unstable correlations, or both.Label-free alternatives on graphs are few, use invariance-inspired contrastive objectives, and do not establish what their learned representations retain.Source: §1 · pp. 1–2
  • SCARCE uses an environment partition alone, which may come from metadata or inferred clusters, rather than class labels during representation learning.Class labels are described as costly in drug discovery, where assays may have tens of labels while unlabeled molecular graphs number in the tens of millions.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • VGAE reconstruction retains both environment-specific and core features, and the paper verifies that both remain linearly accessible.The authors argue that likelihood-based reconstruction retains the whole input, unlike contrastive objectives that may suppress features not needed to distinguish instances.Source: §1 · p. 2
  • Adversarial confusion contracts environment-discriminative directions rather than erasing them, equivalent to an anisotropic ridge penalty on a downstream probe.The paper reports state-of-the-art label-free OOD accuracy on GOOD-Motif and the DrugOOD comparison in the first point.Source: §1 · p. 2
  • SCARCE attains state-of-the-art label-free OOD accuracy on GOOD-Motif and is best in 17 of 18 DrugOOD settings.The introduction also reports that the method works when no environment annotation exists and that a controlled synthetic study confirms the theory's preconditions.Source: §1 · pp. 2–3

2 RELATED WORK

Prior graph OOD methods typically use class labels to select predictive substructures or enforce invariance across environments. SCARCE instead learns by reconstruction and contracts environment-associated directions, requiring an environment partition but no class labels during representation learning.

OOD generalization with labeled supervision

Graph OOD generalization

Unsupervised graph representation learning

  • RGCL [Li et al., 2022b] and OOD-GCL use contrastive objectives, while SCARCE uses reconstruction to retain core and environment factors and contracts rather than erases environment directions.OOD-GCL also penalizes variation across inferred environments; SCARCE uses clustered pseudo-environments in a label-free objective.Source: §2 · p. 3

Feature reweighting and concept erasure

  • Last-layer retraining [Kirichenko et al., 2023] reweights core over spurious features using group-balanced data, whereas concept erasure [Ravfogel et al., 2022; Belrose et al., 2023] removes linearly decodable protected-attribute information.The paper positions SCARCE between them: representation-space reweighting without group-balanced labeled data or a protected attribute.Source: §2 · p. 3

3 PROBLEM SETTING

The paper studies graph-level unsupervised domain generalization: learning representations robust to shifts without class labels during representation learning. It models graphs with core, environment-specific, and residual factors, and allows environments to be supplied or inferred.

3 PROBLEM SETTING

  • Graph-level unsupervised domain generalization seeks representations robust to distribution shifts without class labels during representation learning.The goal is to minimize worst-case risk across training and unseen environments using an encoder and downstream classifier.Source: §3 · pp. 3–4
  • Each training graph has an environment index, supplied by dataset metadata or inferred without labels by clustering a self-supervised embedding.The formulation does not depend on which source provides the partition.Source: §3 · p. 3
  • The generative model writes G = g(C, S, ξ), where C determines the label across environments, S has environment-varying label correlation, and ξ is independent residual variation.Because S–Y correlation may weaken or reverse at test time, C is the only reliable predictor of Y in this model.Source: §3 · p. 4
  • The encoder uses message-passing GNN layers and sum pooling to produce a graph-level latent; GIN is the default backbone, with GCN also reported.The latent representation is denoted Z = fθ(G).Source: §3 · p. 4
  • Class labels are unavailable during representation learning, the environment index is the only side information and may be inferred, and labels are used only to fit an evaluation-time linear probe.This constraint rules out invariant-learning methods that require class labels to identify unstable correlations.Source: §3 · p. 4

4 METHOD

SCARCE combines VGAE reconstruction with adversarial environment confusion to retain graph information while contracting environment-discriminative directions. Its analysis connects this contraction to directional regularization of downstream probes and specifies when it can help.

4.1 WHAT FEATURES DOES A VGAE RETAIN?

  • 90.7% OOD accuracy for unsupervised VGAE DFR at epoch 10 accompanies theorem-backed retention of both factors and early linear access to invariant features.Both information bounds are strictly positive when Lrec < min{H(S | C, ξ), H(C | S, ξ)}.Source: §4.1 · pp. 4–5, Figure 2 · p. 5

4.2 ADVERSARIAL ENVIRONMENT CONFUSION

  • Alternating adversary and encoder updates weaken environment distinctions while reconstruction constrains the encoder from discarding all variation.The adversary is trained on detached representations; the encoder then minimizes confusion toward uniform environment predictions while retaining the reconstruction objective.Source: §4.2 · pp. 5–6

4.3 THEORETICAL ANALYSIS: WHY CONFUSION HELPS OOD GENERALIZATION

  • Confusion contracts adversary-discriminative directions, equivalent under an ℓ2-regularized linear probe to increasing the ridge penalty along the contracted direction while leaving orthogonal directions unchanged.The effective penalty along uE is λ/ε^2, versus λ on orthogonal directions; the feature remains available to the probe rather than being erased.Source: §4.3 · pp. 6–7
  • Contraction strictly reduces OOD risk when a training-only spurious correlation lies apart from the core signal and the downstream probe is finitely regularized; it can hurt when both directions align.As regularization approaches zero, the probe can undo contraction and its benefit vanishes.Source: §4.3 · p. 7, Figure 3 · p. 7

5 EXPERIMENTS

SCARCE’s confusion-based feature contraction helps across synthetic and graph benchmarks, with the strongest molecular results spanning all tested label budgets. The synthetic and controlled studies also show that its benefit depends on representation geometry and probe regularization.

5.1 SYNTHETIC EVALUATION

  • In the synthetic study, contraction reduces OOD error for orthogonal directions, but increases error when shortcut and stable-signal directions are fully aligned.The benefit shrinks as overlap increases, showing that contraction can suppress stable signal alongside the shortcut.Source: §5.1 · p. 7
  • In the synthetic study, contraction has negligible effect at λ=10^-6, peaks in benefit at intermediate regularization, and diminishes under excessive regularization.The near-zero-regularization probe can rescale its coefficient to undo contraction.Source: §5.1 · pp. 7–8

5.2 CONTROLLED GRAPH BENCHMARK

  • On GOOD-Motif, SCARCE outperforms GraphCL, MVGRL, OOD-GCL, and the plain VGAE warm-start, while RGCL performs on par.The benchmark uses domain indices as environments, with class labels withheld from the encoder.Source: §5.2 · p. 8, Table 1 · p. 8
  • On GOOD-Motif, confusion’s accuracy gain is non-monotonic: it is +3.8% at C=10^-3, vanishes at C=1, and returns as +1.4 and +2.4% at C ≥ 10^2.At C=1 the probe can partially undo contraction; at high C, the baseline probe can exploit shortcuts already contracted by confusion.Source: §5.2 · p. 8, Figure 4 · p. 8

5.3 REAL-WORLD MOLECULAR EVALUATION

  • DrugOOD experiments use clustered pseudo-environments because assay, scaffold, and size attributes do not provide usable within-training partitions.The method clusters frozen MVGRL embeddings without labels; cluster imbalance is handled through α, and K is selected on validation.Source: §5.3 · pp. 8–9
  • SCARCE leads the strongest baseline by 4.1, 3.7, and 2.7 points at 5%, 10%, and 20% labels, respectively, averaged over six DrugOOD splits.It is best in 17 of 18 dataset–ratio cells; the sole exception is IC50-size at 5% labels.Source: §5.3 · p. 9, Table 2 · p. 9
  • Across top-ranked DrugOOD configurations, β=10 is preferred in nearly every setting, while the best K varies across splits.The authors connect the strong confusion weight to the strongly regularized probe regime of the 5–20% label settings.Source: §5.3 · pp. 8–9

6 CONCLUSION

SCARCE combines reconstruction-based feature retention with adversarial contraction of environment-discriminative directions. Its experiments support the mechanism and its preconditions, and the authors suggest the recipe also extends to other likelihood-based encoders.

6 CONCLUSION

  • SCARCE pairs a VGAE with adversarial confusion, retaining core and environment factors while contracting environment-discriminative directions rather than erasing them.For any ℓ2-regularized probe, this contraction is equivalent to an anisotropic ridge penalty; experiments support the mechanism and its preconditions.Source: §6 · p. 9

ETHICS STATEMENT

The authors report that SCARCE uses publicly available benchmarks without personal data or human subjects. They identify no direct negative societal impact beyond generic impacts of representation learning.

ETHICS STATEMENT

  • The authors report no direct negative societal impact beyond those generic to representation learning.They note that removing label requirements and tolerating coarse or inferred environments may lower predictor-building costs where labeled and annotated data are scarce.Source: ETHICS STATEMENT · p. 10

A.1 PROOF OF THEOREM 4.1 (RECONSTRUCTION FORCES RETENTION)

The proof bounds how much graph information reconstruction leaves unexplained by the latent, then uses factor recoverability to bound residual uncertainty in each generative factor. These bounds are positive under the stated entropy condition.

A.1 PROOF OF THEOREM 4.1 (RECONSTRUCTION FORCES RETENTION)

  • Reconstruction loss is at least H(G | Z), so any encoder–decoder pair with loss Lrec leaves at most Lrec nats of G unexplained by Z.The proof obtains this bound from conditional entropy and the decoder distribution.Source: §A.1 · p. 13
  • Under injectivity in each factor argument, S and C are recoverable from the graph when the other factor and ξ are fixed, yielding separate information bounds for both factors.For S, I(Z; S | C, ξ) ≥ H(S | C, ξ) − Lrec; the corresponding C bound follows symmetrically.Source: §A.1 · p. 13
  • The per-factor bounds are strictly positive when Lrec < min{H(S | C, ξ), H(C | S, ξ)}, and assumption A2 makes the right-hand sides finite.The proof contrasts this guarantee with a discriminative objective, which can discard C when S predicts Y in training environments.Source: §A.1 · p. 13

A.2 PROOF OF LEMMA 4.2 (LOCAL CONTRACTION)

The proof shows that confusion selectively contracts representation directions used by the adversary to distinguish environments, while leaving orthogonal directions unchanged. Under a regularized linear probe, this contraction is equivalent to penalizing use of the environment direction more strongly.

A.2 PROOF OF LEMMA 4.2 (LOCAL CONTRACTION)

  • Directions with larger adversarial singular values contract more strongly, while directions orthogonal to the adversary’s span remain unchanged.The result describes a single-step, first-order contraction for a fixed linear adversary.Source: §A.2 · p. 14
  • Under the rank-one approximation, repeated confusion steps shrink the environment component while preserving the orthogonal component.The cumulative model assumes a stable dominant adversarial direction; the implemented encoder and adversary depart from this simplified setup.Source: §A.2 · p. 15

A.3 PROOF OF PROPOSITION 4.3 (RIDGE IDENTITY)

  • For ε > 0, contraction preserves information but makes use of the environment direction more costly under an ℓ2-regularized linear probe.The change of variables preserves prediction loss while adding a directional penalty to the regularizer.Source: §A.3 · pp. 15–16

A.4 EXACT WORKED EXAMPLE: STATEMENT AND PROOF

The worked example analyzes when contracting an environment-associated feature helps or hurts OOD risk under ridge regularization. It shows benefit under a training-to-test correlation shift when the contracted feature is separated from the stable signal, but harm when the directions align.

A.4 EXACT WORKED EXAMPLE: STATEMENT AND PROOF

  • When a spurious training correlation vanishes at test time and λ > 0, stronger contraction strictly reduces OOD risk.The analysis also requires the contracted environment direction to remain separated from invariant task-relevant information.Source: §A.4 · pp. 16, 18
  • The OOD-risk decomposition separates squared core-feature bias, noisy-core-channel penalty, and variance from loading on the shifting environment feature.The derivation evaluates risk at test-time environment correlation.Source: §A.4 · p. 17
  • Without downstream regularization, the fitted predictor can rescale the contracted direction and undo the contraction.The fitted coefficients become independent of ε as regularization approaches zero.Source: §A.4 · pp. 16, 18
  • When the environment direction fully aligns with the core signal, contraction hurts because it compresses task-relevant information.The worked example reports R(1) ≈ 0.667 and R(0.1) ≈ 0.864.Source: §A.4 · pp. 16, 18

B DFR PROTOCOL

The appendix specifies a DFR protocol that reweights frozen encoder features using a group-balanced labeled set, then broadens the real-world results and reports evaluated hyperparameter configurations across DrugOOD settings and label ratios.

B DFR PROTOCOL

  • DFR freezes the encoder, forms a group-balanced set using label–environment groups, standardizes its representations, and fits a sparse linear head.The head’s regularization is selected by worst-group validation accuracy, and the balancing and fitting steps are repeated across ten resamples.Source: B DFR PROTOCOL · p. 19
  • DFR changes only the weighting of frozen features, so its OOD accuracy measures what the encoder already makes linearly accessible.The protocol evaluates checkpoints from supervised ERM or unsupervised VGAE training on GOOD-Motif concept shift.Source: B DFR PROTOCOL · p. 19

C ADDITIONAL REAL-WORLD RESULTS

  • Real-world experiments fit linear probes using 5%, 10%, or 20% of labeled training data, selecting models on validation and reserving test data for final evaluation.The expanded analyses cover the available molecular shifts.Source: C ADDITIONAL REAL-WORLD RESULTS · p. 19

C.1 HYPERPARAMETER SENSITIVITY ACROSS DRUGOOD

  • Tables 3–5 list six evaluated (α, β, K) configurations per DrugOOD dataset–domain block, ordered by test ROC-AUC for each label ratio.Untested configurations are neither included nor interpolated.Source: §C.1 · p. 19

C.2 COMPLETE TRAINING DYNAMICS

The training-dynamics figures document bounded optimization losses across the DrugOOD settings, while the tables summarize selected hyperparameter configurations. The loss plots are diagnostics, not evidence by themselves of invariance or causal removal of environment information.

C.2 COMPLETE TRAINING DYNAMICS

  • Across the runs, reconstruction, confusion, and adversary objectives remain bounded, though their scales and transient behavior vary by dataset and split.Stable losses rule out persistent numerical divergence but do not establish invariance or causal removal of environment information.Source: §C.2 · p. 19
  • Tables 3–5 order the six displayed evaluated configurations within each dataset–domain setting by test ROC-AUC at 5%, 10%, and 20% label ratios.Untested configurations are not interpolated.Source: Table 3 · p. 20, Table 4 · p. 21, Table 5 · p. 22
  • Figures 5–10 show training dynamics for six DrugOOD settings, with each caption reporting its validation-selected (α, β, K) configuration.Figure 5 identifies reconstruction, confusion, and adversary losses by line style and marker; lines average three seeds and shading denotes mean ± sample standard deviation.Source: Figure 5 · p. 23, Figure 6 · p. 23, Figure 7 · p. 24

Related papers

Loading 2610.29571v1…