Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Zhang, Simon, DeMilt, Ryan P., Jin, Kun, Xia, Cathy H.
Formato: Preprint
Publicado: 2026
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866908981032648704
author Zhang, Simon
DeMilt, Ryan P.
Jin, Kun
Xia, Cathy H.
author_facet Zhang, Simon
DeMilt, Ryan P.
Jin, Kun
Xia, Cathy H.
contents Out-of-distribution (OoD) generalization occurs when representation learning encounters a distribution shift. This occurs frequently in practice when training and testing data come from different environments. Covariate shift is a type of distribution shift that occurs only in the input data, while the concept distribution stays invariant. We propose RIA - Regularization for Invariance with Adversarial training, a new method for OoD generalization under convariate shift. Motivated by an analogy to $Q$-learning, it performs an adversarial exploration for counterfactual data environments. These new environments are induced by adversarial label invariant data augmentations that prevent a collapse to an in-distribution trained learner. It works with many existing OoD generalization methods for covariate shift that can be formulated as constrained optimization problems. We develop an alternating gradient descent-ascent algorithm to solve the problem in the context of causally generated graph data, and perform extensive experiments on OoD graph classification for various kinds of synthetic and natural distribution shifts. We demonstrate that our method can achieve high accuracy compared with OoD baselines.
format Preprint
id arxiv_https___arxiv_org_abs_2604_08404
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization
Zhang, Simon
DeMilt, Ryan P.
Jin, Kun
Xia, Cathy H.
Machine Learning
Out-of-distribution (OoD) generalization occurs when representation learning encounters a distribution shift. This occurs frequently in practice when training and testing data come from different environments. Covariate shift is a type of distribution shift that occurs only in the input data, while the concept distribution stays invariant. We propose RIA - Regularization for Invariance with Adversarial training, a new method for OoD generalization under convariate shift. Motivated by an analogy to $Q$-learning, it performs an adversarial exploration for counterfactual data environments. These new environments are induced by adversarial label invariant data augmentations that prevent a collapse to an in-distribution trained learner. It works with many existing OoD generalization methods for covariate shift that can be formulated as constrained optimization problems. We develop an alternating gradient descent-ascent algorithm to solve the problem in the context of causally generated graph data, and perform extensive experiments on OoD graph classification for various kinds of synthetic and natural distribution shifts. We demonstrate that our method can achieve high accuracy compared with OoD baselines.
title Adversarial Label Invariant Graph Data Augmentations for Out-of-Distribution Generalization
topic Machine Learning
url https://arxiv.org/abs/2604.08404