Boundary Embedding Shaping (BES)

• BES is a label-informed adaptive contrastive learning module for GNNs that selectively refines node embeddings near decision boundaries to mitigate structural entanglement.
• It functions as a post-encoding refinement stage by applying an attention-based shaping layer to boundary nodes while preserving the overall topology of embeddings.
• BES employs supervised contrastive objectives with perturbation-aware optimization to improve node classification and link prediction performance across diverse graph datasets.

Boundary Embedding Shaping (BES) is a label-informed, adaptive contrastive learning plug-in module for graph neural networks (GNNs) that addresses graph structural disentanglement by selectively reshaping node embeddings near class decision boundaries, where spurious correlations from semantically irrelevant neighbors are most likely to destabilize prediction (Chen et al., 18 Jun 2026). Rather than robustifying all nodes uniformly, BES identifies a boundary region in embedding space, mines hard boundary nodes, refines their representations with an attention-based shaping layer, and scales updates by their global embedding perturbation so that boundary discrimination improves with minimal disturbance to non-boundary embeddings (Chen et al., 18 Jun 2026).

1. Problem formulation and conceptual basis

BES is motivated by a specific failure mode of message-passing GNNs: structural entanglement, in which learned node embeddings mix class-relevant information with misleading local topology. The paper argues that this contamination is not uniformly harmful across the graph. It is most damaging for nodes near class decision boundaries, because small amounts of structural noise can move an embedding across a class separator and blur the effective margin. BES therefore treats boundary-region entanglement as the primary bottleneck, in contrast to robust GNN methods that denoise or regularize all nodes in the same way (Chen et al., 18 Jun 2026).

The conceptual model is a latent-variable factorization in which each node observation $x$ is generated from latent factors $z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$, with $z_{\mathrm{inv}}$ carrying stable, class-determining information and $z_{\mathrm{var}}$ carrying changing local structural effects that act as noise. The paper writes the generative process as

$$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$

Within this view, structural disentanglement aims to preserve invariant class structure while suppressing variant topological effects. Corollary 3.1 expresses the ideal target as

$$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$

The boundary emphasis is justified theoretically by Proposition 3.1: $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$ The paper interprets this as showing that disentanglement error scales inversely with boundary margin. This suggests that the smallest-margin region, rather than the average node, is the dominant source of robustness loss.

2. Architectural position within a GNN pipeline

BES is not a standalone GNN architecture. It is a post-encoding refinement stage that improves representations produced by one or more base encoders and then passes the refined embeddings to a downstream decoder for node classification or link prediction (Chen et al., 18 Jun 2026).

For a base encoder $\Phi$, BES introduces an attention-based shaping layer $A^{(B)}$ with output

$Z^{(B)} = A^{(B)}(\Phi(X);\theta^{(B)}).$

The appendix generalizes the encoder side to multiple encoders,

$z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$0

The underlying encoder computation remains standard message passing. For node $z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$1,

$z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$2

$z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$3

$z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$4

BES leaves this encoder stage conceptually intact and operates on the resulting embeddings rather than modifying message passing or graph connectivity directly.

The shaping module itself is defined as a standard multi-head self-attention block: $z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$5 with $z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$6, $z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$7, and $z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$8. Head outputs are concatenated and projected with a residual connection,

$z=(z_{\mathrm{inv}}, z_{\mathrm{var}})$9

In the reported implementation, two attention layers $z_{\mathrm{inv}}$0 and $z_{\mathrm{inv}}$1 are trained separately with gravity loss, followed by a Chebyshev-convolution decoder. BES is therefore best understood as a modular representation-shaping layer that can be attached to existing backbones rather than as a replacement for them.

3. Boundary-region estimation and boundary-node mining

BES begins by estimating a coarse boundary slab in embedding space rather than scanning the full graph uniformly. The boundary region between two classes or clusters $z_{\mathrm{inv}}$2 and $z_{\mathrm{inv}}$3 is defined by

$z_{\mathrm{inv}}$4

where $z_{\mathrm{inv}}$5 and $z_{\mathrm{inv}}$6 are class or cluster means, $z_{\mathrm{inv}}$7 is the covariance term used in the Mahalanobis-like projection, and $z_{\mathrm{inv}}$8 is the slab width. The paper models each cluster with a Gaussian $z_{\mathrm{inv}}$9 (Chen et al., 18 Jun 2026).

Within this slab, BES computes a structural variability, or shift score, over embedding-space neighbors. In prose, the score is the fraction of local embedding-space neighbors having a different label than the candidate node. Nodes with high shift score, for example $z_{\mathrm{var}}$0, are designated boundary nodes. The method is explicitly label-informed only on the training split: labels used for boundary identification come from training nodes, and test nodes and test labels are excluded.

The use of embedding-space neighbors rather than graph-topology neighbors is deliberate. The paper argues that these neighbors contribute more directly to the objective and empirically preserve the graph’s heterophilic characteristics while producing a more stable signal. Table 6 is described as showing that the two definitions are close, with embedding-space neighborhoods slightly better on some datasets.

This mining strategy makes BES selective in two senses. First, it restricts intervention to the decision-boundary slab. Second, within that slab it focuses on nodes whose local representation neighborhoods are most label-misaligned. A plausible implication is that BES converts graph robustness from a global smoothing problem into a targeted local geometry problem.

4. Contrastive shaping and perturbation-aware optimization

For each mined boundary node $z_{\mathrm{var}}$1, BES applies supervised contrastive shaping. Positive examples are same-class nodes, and negatives are nodes from other classes. The exact pairwise similarities are defined as

$z_{\mathrm{var}}$2

$z_{\mathrm{var}}$3

The paper interprets these terms as “positive gravity” and “negative gravity”: the boundary node is pulled toward its class manifold and pushed away from competing classes (Chen et al., 18 Jun 2026).

To avoid the $z_{\mathrm{var}}$4 cost of exact pairwise attraction to all positives, BES introduces a center-based approximation and states that it is gradient-equivalent. The key identity is

$z_{\mathrm{var}}$5

Since the second term is constant with respect to $z_{\mathrm{var}}$6, attraction to all same-class positives can be replaced by attraction to the class centroid. The surrogate similarities become

$z_{\mathrm{var}}$7

$z_{\mathrm{var}}$8

The hinge-like $z_{\mathrm{var}}$9 prevents over-penalizing nodes already sufficiently close to their own class center relative to the nearest boundary exemplar.

The paper refers to the resulting objective as the “gravity loss for structural disentanglement,” but does not provide a single final batched summation formula in the provided text. It is, however, explicit that BES uses a boundary-focused supervised contrastive objective instantiated with centroid-based gravity terms.

A second adaptive mechanism controls how strongly refinement can perturb the overall representation geometry. At each iteration, BES samples a batch of $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$0 boundary nodes, computes a loss gradient, and performs a virtual update

$$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$1

It then measures the induced global embedding perturbation: $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$2 The actual update is scaled inversely by this perturbation, with learning rate $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$3 and scaling hyperparameter $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$4. The intended effect is to prevent overshooting and preserve non-boundary embeddings while boundary nodes are being reshaped.

5. Theoretical interpretation and geometric effect on embedding space

The theoretical justification for BES is organized around the relation between contrastive learning, invariant factor recovery, and class-boundary margin. Theorem 3.1 states that minimizing a contrastive loss with suitable positive-pair construction can block-identify the invariant latent factors $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$5. The positive-pair assumption is label-informed: positive pairs share the same invariant latent block while allowing sufficient variation in the variant block. In graph terms, two nodes from different local subgraphs but the same class can serve as positives because they preserve semantics while varying in structural context (Chen et al., 18 Jun 2026).

This construction makes the intended geometry explicit. Same-class positives encourage embeddings to align around invariant class structure, while negatives strip away variant, semantically unstable structural effects. The proposition-level margin bound then motivates concentrating this contrastive pressure where the boundary margin is smallest.

The paper’s qualitative claims about geometry are correspondingly local. BES aims to increase inter-class separation in boundary areas and improve local margin without unnecessarily altering the global embedding arrangement. The visualizations in Figures 1, 3, 5, 6, and 7 are described as showing ambiguous nodes leaving the boundary slab and class clusters becoming more separable. The same mechanism also improves intra-class compactness around hard nodes, because attraction is to same-class positives or class centroids.

The method is explicitly not based on stochastic graph augmentations. One of its conceptual premises is that structural variation already exists in graphs, especially heterophilic ones, and can therefore be mined directly from the observed data. This suggests that BES treats graphs with noisy neighborhoods by operating on mined representation geometry rather than by synthesizing perturbations or pruning edges.

6. Empirical performance and ablation structure

BES is evaluated on node classification and link prediction across homogeneous datasets (WikiCS, Cora, CiteSeer, Pubmed), heterogeneous datasets (Chameleon, Cornell, Texas), a heterophilic dataset (Roman-empire), and OGB benchmarks (ogbn-arxiv, ogbl-collab). Non-OGB datasets use random $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$6 train/validation/test splits; OGB uses official protocols; results are averaged over 5 runs (Chen et al., 18 Jun 2026).

The most prominent reported evidence is that BES improves common backbones as a plug-in. The headline claim is that it boosts GCN performance by an average of $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$7 in node classification, with the best gain of $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$8 on WikiCS.

Setting	BES result	Comparison
Node classification, Cora	89.46	GCN 87.74
Node classification, CiteSeer	78.40	GCN 75.74
Node classification, Pubmed	91.03	GCN 87.36
Node classification, WikiCS	85.52	GCN 80.48
Link prediction, Pubmed	97.89	IGCL 96.88
Link prediction, ogbl-collab	65.84	beats reported baselines

On Pubmed node classification, BES reaches $$z=(z_{\mathrm{inv}},z_{\mathrm{var}})\sim p_z,\qquad y=g_y(z_{\mathrm{inv}}),\qquad x=g(z).$$9, outperforming IGCL $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$0, H2GCN $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$1, and MPNN-SAGE $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$2. On WikiCS, it improves over IGCL $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$3 and slightly exceeds MPNN-GAT $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$4 with $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$5. For link prediction on homogeneous datasets, BES reports Cora $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$6, CiteSeer $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$7, Pubmed $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$8, and WikiCS $$\max_f I(z;Z_{\mathrm{inv}})\quad \text{s.t.}\quad I(z;Z_{\mathrm{var}})\to 0.$$9. On WikiCS link prediction it slightly surpasses H2GCN $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$0 and GCN $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$1.

The ablations are consistent with the boundary-focused claim. On WikiCS node classification, encoder combination without BES gives $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$2, one BES layer gives $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$3, and two layers give $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$4. Table 16 reports gains from using mined boundary nodes: WikiCS improves from $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$5 to $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$6, Pubmed from $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$7 to $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$8, Roman-empire from $$\inf_A \mathbb E[\|z-A Z_{\mathrm{inv}}\|^2] \le O(\sqrt{\mathcal E}) + O\!\left(\frac{1}{\operatorname{margin}(d Z_{\mathrm{inv}})}\right).$$9 to $\Phi$0, and ogbn-arxiv from $\Phi$1 to $\Phi$2. The reported best settings include temperature $\Phi$3 and scaling factor $\Phi$4. Efficiency experiments comparing the exact pairwise loss with the center-based approximation show similar or slightly better accuracy for the center-based version with substantially lower convergence time.

The method is not computationally free. The overall time complexity is

$\Phi$5

where $\Phi$6 is the number of BES layers, $\Phi$7 the training iterations per layer, $\Phi$8 the number of nodes, $\Phi$9 the average degree or neighborhood size, and $A^{(B)}$0 the embedding dimension. Space complexity is

$A^{(B)}$1

with $A^{(B)}$2 classes. The dominant overhead comes from the adaptive perturbation calculation across all nodes during each update.

7. Scope, limitations, and terminological boundaries

Several features of BES are easiest to understand by stating what it does not do. It does not replace the encoder, does not alter message passing inside the base GNN directly, does not edit original node features, does not delete graph edges, and does not require a special inference-time boundary-mining stage. Its intervention is concentrated at the representation and loss level, during training, through label-informed boundary-node mining and attention-based shaping (Chen et al., 18 Jun 2026).

The method is also supervised, or at least explicitly label-informed. Training labels are required for boundary scores and for the construction of positive and negative pairs. The paper identifies low-label and self-supervised settings as future work. Another limitation is dependence on the quality of the initial embedding geometry: if embeddings are extremely poor, the slab-based boundary approximation and mined boundary nodes may be noisy. The paper partly addresses this through perturbation-aware updates, but this remains an implicit limitation. It also notes that the exact final batch loss formula combining class loss and gravity loss is not fully specified in the provided text.

The term “boundary” in Boundary Embedding Shaping refers to class decision boundaries in embedding space, not to physical or geometric boundaries in the sense of free-boundary surface theory, and not to diagnostic instrumentation. The acronym BES is used in unrelated literatures for Beam Emission Spectroscopy (Barbisan et al., 2022), while geometric boundary-embedding problems include free boundary isometric embedding into the unit ball (Koerber, 2019). This terminological overlap is superficial. In the graph-learning sense established by “Boundary Embedding Shaping with Adaptive Contrastive Learning for Graph Structural Disentanglement,” BES denotes a boundary-local representation refinement strategy for GNNs (Chen et al., 18 Jun 2026).