ScoreAug: Score-Based Graph Augmentation

Updated 8 July 2026

ScoreAug is a score-based out-of-distribution augmentation method for graphs that leverages guided reverse diffusion to simulate unseen environmental features without explicit S/E disentanglement.
It integrates label and exploration guidance to balance OOD diversity and sample validity during the reverse SDE process.
Incorporating ScoreAug with GNN training boosts performance under distribution shifts, achieving significant gains on benchmarks like GOOD-Motif and GOOD-CMNIST.

Searching arXiv for the specified paper and closely related "ScoreAug"/"Score Augmentation" usages to ground the article in the cited literature. I’ll look up the target arXiv record and adjacent papers using the arXiv search interface. ScoreAug, also termed OODA in the graph-learning paper “Mitigating Graph Covariate Shift via Score-based Out-of-distribution Augmentation,” is a score-based out-of-distribution augmentation method for graph learning under covariate shift. It is motivated by the setting in which distribution shifts between training and testing datasets significantly impair model performance, while a causal view of graph invariant learning treats stable predictive features as label-causal and environmental features as the source of shift. In this formulation, ScoreAug avoids explicit stable/environmental disentanglement in input space and instead uses score-based graph generation to synthesize unseen environmental features while preserving the validity and stable features of overall graph patterns (Wang et al., 2024).

1. Problem setting and invariant-learning formulation

In the graph setting, a graph is written as $G=(A,X,E)$ , with $A\in\mathbb{R}^{n\times n}$ denoting adjacency, $X\in\mathbb{R}^{n\times a}$ node features, and $E\in\mathbb{R}^{n\times n\times b}$ edge features, often folded into $A$ . Each graph has a label $Y\in\mathcal{Y}$ determined by some unknown rule. The causal decomposition assumed in the method separates each graph into $S$ , the stable substructure or features, and $E$ , the environmental substructure or features. The invariant, covariate-shift assumption is

$P_{tr}(Y\mid G)=P_{te}(Y\mid G)$

for all environments, while

$P_{tr}(G)\neq P_{te}(G),$

or equivalently $A\in\mathbb{R}^{n\times n}$ 0 while $A\in\mathbb{R}^{n\times n}$ 1 stays the same (Wang et al., 2024).

The training environments $A\in\mathbb{R}^{n\times n}$ 2 are assumed to be only a subset of all possible $A\in\mathbb{R}^{n\times n}$ 3, whereas testing exposes the model to new environments $A\in\mathbb{R}^{n\times n}$ 4. The out-of-distribution generalization objective is to learn $A\in\mathbb{R}^{n\times n}$ 5 minimizing

$A\in\mathbb{R}^{n\times n}$ 6

Because only data from $A\in\mathbb{R}^{n\times n}$ 7 are available in practice, augmentation is treated as necessary rather than optional (Wang et al., 2024).

This formulation situates ScoreAug against perturbation-based graph augmentation methods that selectively perturb or mixup environmental features after attempting to disentangle stable and environmental components. The paper states two limitations of that family: heavy reliance on an accurate separation of stable and environmental features, and exploration confined to existing environmental features in the training distribution. ScoreAug is proposed precisely to address those two limits.

2. Score-based graph generation backbone

The generative component jointly diffuses node features $A\in\mathbb{R}^{n\times n}$ 8 and adjacency $A\in\mathbb{R}^{n\times n}$ 9 over continuous time. The forward Itô SDEs are

$X\in\mathbb{R}^{n\times a}$ 0

and the reverse SDEs are

$X\in\mathbb{R}^{n\times a}$ 1

$X\in\mathbb{R}^{n\times a}$ 2

Here $X\in\mathbb{R}^{n\times a}$ 3 is the marginal at time $X\in\mathbb{R}^{n\times a}$ 4, and the drifts and noise scales are chosen so that $X\in\mathbb{R}^{n\times a}$ 5 yields a simple prior, for example a standard normal (Wang et al., 2024).

The score model is a single graph transformer-based network

$X\in\mathbb{R}^{n\times a}$ 6

where each branch is implemented by $X\in\mathbb{R}^{n\times a}$ 7 layers of a graph transformer with attention and MLP blocks. Its input is the noisy pair $X\in\mathbb{R}^{n\times a}$ 8 together with a learned $X\in\mathbb{R}^{n\times a}$ 9-embedding, specified as sinusoidal or MLP-based, and its output estimates $E\in\mathbb{R}^{n\times n\times b}$ 0 and $E\in\mathbb{R}^{n\times n\times b}$ 1 (Wang et al., 2024).

Training uses a denoising score matching objective. Given the corruption kernel $E\in\mathbb{R}^{n\times n\times b}$ 2, the loss is

$E\in\mathbb{R}^{n\times n\times b}$ 3

The implementation expands this into the sum of two $E\in\mathbb{R}^{n\times n\times b}$ 4 losses, one for $E\in\mathbb{R}^{n\times n\times b}$ 5 and one for $E\in\mathbb{R}^{n\times n\times b}$ 6. The variance schedule $E\in\mathbb{R}^{n\times n\times b}$ 7 or $E\in\mathbb{R}^{n\times n\times b}$ 8 is typically linear from $E\in\mathbb{R}^{n\times n\times b}$ 9 to $A$ 0, and $A$ 1 is normalized into $A$ 2 and embedded by an MLP (Wang et al., 2024).

3. Environmental feature synthesis by guided reverse sampling

ScoreAug’s central move is to replace explicit $A$ 3 disentanglement with guided sampling of the trained diffusion model. The sampling conditions on two quantities: the desired class label $A$ 4, which is used to preserve stable features, and an exploration parameter $A$ 5, which pushes samples into lower-density, out-of-distribution regions (Wang et al., 2024).

The full conditional score is written as $A$ 6 and is decomposed by Bayes’ rule into three terms:

$A$ 7,
$A$ 8, called stable-pattern guidance,
$A$ 9, called OOD exploration guidance.

Following Lee et al. (2023), the paper sets

$Y\in\mathcal{Y}$ 0

so that

$Y\in\mathcal{Y}$ 1

This yields the guided score

$Y\in\mathcal{Y}$ 2

The class-conditional term $Y\in\mathcal{Y}$ 3 is approximated by a time-dependent classifier $Y\in\mathcal{Y}$ 4 trained to predict $Y\in\mathcal{Y}$ 5 from $Y\in\mathcal{Y}$ 6 (Wang et al., 2024).

Algorithm 1 in the paper instantiates this idea as a reverse-diffusion sampler. For each labeled training graph $Y\in\mathcal{Y}$ 7, the procedure samples $Y\in\mathcal{Y}$ 8, iterates from $Y\in\mathcal{Y}$ 9 down to $S$ 0, computes the score-network output $S$ 1, evaluates $S$ 2, obtains $S$ 3 by backpropagation through $S$ 4, and forms

$S$ 5

The state is then updated by Euler–Maruyama steps for both $S$ 6 and $S$ 7. The final output is an augmented graph $S$ 8 carrying the original label $S$ 9 (Wang et al., 2024).

This design makes the environmental synthesis implicit in the generative trajectory. A plausible implication is that the method trades explicit factor decomposition for control through conditional scores and density exploration.

4. Integration with downstream GNN training

Once ScoreAug generates out-of-distribution graphs, they are assembled into an augmentation set

$E$ 0

The downstream classifier $E$ 1, such as GIN or GCN, is trained on both the original training set and the augmented set with

$E$ 2

with $E$ 3 often set to $E$ 4 (Wang et al., 2024).

The paper also states an optional consistency regularizer,

$E$ 5

Under this training recipe, the classifier is exposed not only to the original observed environments but also to unseen-environment graphs synthesized by the score model (Wang et al., 2024).

Within the paper’s logic, this is the mechanism by which OOD generalization is improved: the augmentation stage broadens the environmental support seen during training without changing the graph label. This suggests that ScoreAug is best viewed as a data-space intervention attached to a standard supervised GNN pipeline rather than as a replacement for the predictive model.

5. Theoretical remarks, diagnostics, and empirical results

The theoretical content described in the paper is deliberately limited. The only formal proof shows that

$E$ 6

so that explicit modeling of $E$ 7 is unnecessary. No classical OOD generalization bound is given. Instead, control over divergence is diagnosed through MMD-RBF (Wang et al., 2024).

The empirical evaluation is reported on the GOOD benchmark. The stated highlights are:

GOOD-Motif, for structure-covariate shifts: +6–9% over ERM and best baselines.
GOOD-CMNIST, for node-feature shifts: +18% over AIA, +21% over DIR.
GOOD-HIV, for molecular scaffold/size shifts: +2–6%.
GOOD-SST2, for grammar tree length shifts: +0.7% over MaskFeature (Wang et al., 2024).

The ablation results isolate the two guidance terms. Without exploration guidance, that is, with $E$ 8, the paper reports no OOD diversity. Without stable-pattern guidance, meaning the $E$ 9 term is dropped, the method produces invalid OOD samples. Using both $P_{tr}(Y\mid G)=P_{te}(Y\mid G)$ 0 and $P_{tr}(Y\mid G)=P_{te}(Y\mid G)$ 1 guidance yields the best performance (Wang et al., 2024).

These findings are tightly aligned with the method’s design. The exploration term is responsible for moving samples into lower-density regions, whereas the classifier-guidance term is responsible for retaining class-consistent graph structure. This suggests a division of labor between “OOD diversity” and “sample validity” inside the guided reverse SDE.

The label “ScoreAug” is not unique to graph OOD generalization. The data block contains multiple unrelated procedures that use the same or a closely related name. This suggests that “ScoreAug” functions as a cross-domain label for methods that augment training, inference, or evaluation with score information, but not as a single standardized algorithm (Kim et al., 2023).

arXiv id	Domain	Core mechanism
(Wang et al., 2024)	Graph OOD generalization	Score-based graph generation with label and exploration guidance
(Kim et al., 2023)	Contrastive learning	Pairwise weighting by differences of score vectors
(Eddine et al., 2022)	Machine translation evaluation	Data augmented translations and score aggregation
(Ulloa-Pérez et al., 2024)	Causal effect estimation	Propensity-score model augmented with a prognostic score term
(Shen et al., 12 May 2026)	Neural likelihood surrogates	BCE augmented with exact score information and adaptive weighting
(Hou et al., 11 Aug 2025)	Diffusion models	Augment noisy inputs and predict augmented targets via an equivariant loss

In “ScoreCL: Augmentation-Adaptive Contrastive Learning via Score-Matching Function,” the method called ScoreAug re-weights InfoNCE pairs with

$P_{tr}(Y\mid G)=P_{te}(Y\mid G)$ 2

using denoising-score estimates of augmented views (Kim et al., 2023). In “DATScore,” score augmentation refers to adding augmented translations of the source and reference and aggregating eight hypothesis-centered generation directions (Eddine et al., 2022). In “Propensity Score Augmentation in Matching-based Estimation of Causal Effects,” augmentation means adding a one-dimensional prognostic covariate $P_{tr}(Y\mid G)=P_{te}(Y\mid G)$ 3 to the propensity-score logit model to reduce variance in $P_{tr}(Y\mid G)=P_{te}(Y\mid G)$ 4 matching (Ulloa-Pérez et al., 2024). In “Keeping Score,” the Score-Augmented Loss adds exact simulator score information $P_{tr}(Y\mid G)=P_{te}(Y\mid G)$ 5 to binary-classification training of neural likelihood surrogates (Shen et al., 12 May 2026). In “Score Augmentation for Diffusion Models,” augmentation is performed on noisy samples, and the denoiser is trained to predict the augmentation of the original target, yielding an equivariant learning objective over denoising spaces (Hou et al., 11 Aug 2025).

Because these usages differ in objective, mathematical object, and training protocol, disambiguation by domain is essential. In current arXiv usage, “ScoreAug” may refer to score-based graph OOD augmentation, score-guided pair weighting, data-augmented evaluation scoring, propensity-score model augmentation, score-augmented surrogate loss design, or equivariant score learning in diffusion models.