AdversarialLearner(data_src, data_tar, dim_hidden=64, num_layer=2, source_clf=True, norm_mode='PN', norm_scale=1.0)

Bases: Module

Adversarial learning framework combining source, target, and discriminator networks.

Parameters:
  • data_src (Data) –

    Source domain graph data

  • data_tar (Data) –

    Target domain graph data

  • dim_hidden (int, default: 64 ) –

    Hidden dimension size. Default: 64

  • num_layer (int, default: 2 ) –

    Number of layers in networks. Default: 2

  • source_clf (bool, default: True ) –

    Whether to use source classification. Default: True

  • norm_mode (str, default: 'PN' ) –

    Normalization mode ('PN', etc.). Default: 'PN'

  • norm_scale (float, default: 1.0 ) –

    Normalization scale factor. Default: 1.
Notes

Architecture components:

  • SourceLearner for source domain processing
  • TargetLearnerAE for target domain processing with reconstruction
  • Discriminator for adversarial training
get_probs_cross_domain(data_src, data_tar, idx1, idx2, return_representation=False)

Compute similarity probabilities between source and target domains.

Parameters:
  • data_src (Data) –

    Source domain graph data

  • data_tar (Data) –

    Target domain graph data

  • idx1 (Tensor) –

    Source domain indices

  • idx2 (Tensor) –

    Target domain indices

  • return_representation (bool, default: False ) –

    Whether to return encoded features. Default: False
Returns:
  • tuple –

    Contains:

    • torch.Tensor: Cross-domain pair probabilities
    • torch.Tensor: Source classification probabilities
    • torch.Tensor: Target classification probabilities
    • torch.Tensor: Source encoded features (if return_representation=True)
    • torch.Tensor: Target encoded features (if return_representation=True)
Notes
  • Encodes both domains separately
  • Computes cross-domain similarities
  • Returns detached representations if requested
  • Returns zero probabilities if source_clf is False
get_probs_within_domain(data, idx1, idx2, domain='target')

Compute similarity probabilities within a single domain.

Parameters:
  • data (Data) –

    Input graph data

  • idx1 (Tensor) –

    First set of node indices

  • idx2 (Tensor) –

    Second set of node indices

  • domain (str, default: 'target' ) –

    Domain type ('source' or 'target'). Default: 'target'
Returns:
  • tuple –

    Contains:

    • torch.Tensor: Pair similarity probabilities
    • torch.Tensor: Classification probabilities (exp of log_probs)
Notes
  • For source domain: Uses source learner directly
  • For target domain: Encodes first, then uses source similarity network
  • Returns zero probabilities if source_clf is False
BridgedGraph(data_src, data_tar, k_cross=20, k_within=6, epsilon=0.5, dim_hidden=64, batch_size=1000, num_layer=2, num_epoch=200, lr=0.001, weight_decay=0.005, source_clf=True, norm_mode='PN', norm_scale=1.0, device=None)

Bases: Module

pair_enumeration(x1, x2)

input: [B,D] return: [B*B,D] input [[a], [b]] return [[a,a], [b,a], [a,b], [b,b]]

Decoder(dim_in, dim_hidden, dim_out, num_layer=2, use_norm=False, dropout=0.5, act_fn='relu', norm_mode='PN', norm_scale=1.0)

Bases: Module

Multi-layer MLP decoder with optional normalization and various activation functions.

Parameters:
  • dim_in (int) –

    Input feature dimension

  • dim_hidden (int) –

    Hidden layer dimension

  • dim_out (int) –

    Output feature dimension

  • num_layer (int, default: 2 ) –

    Number of linear layers. Default: 2

  • use_norm (bool, default: False ) –

    Whether to use PairNorm normalization. Default: False

  • dropout (float, default: 0.5 ) –

    Dropout rate. Default: 0.5

  • act_fn (str, default: 'relu' ) –

    Activation function type ('relu', 'leakyrelu', 'tanh', 'sigmoid'). Default: 'relu'

  • norm_mode (str, default: 'PN' ) –

    PairNorm mode if use_norm is True. Default: 'PN'

  • norm_scale (float, default: 1.0 ) –

    Normalization scaling factor. Default: 1.
Notes

Architecture:

  • Multiple linear layers with configurable width
  • Optional PairNorm after each intermediate layer
  • Configurable activation functions
  • Bias terms included in linear layers
forward(z)

Forward pass of the decoder.

Parameters:
  • z (Tensor) –

    Input features of shape [batch_size, dim_in]
Returns:
  • Tensor –

    Reconstructed features of shape [batch_size, dim_out]
Notes

Process for each intermediate layer:

  • Linear transformation
  • Optional PairNorm (if use_norm=True)
  • Activation function

Final layer:

  • Only linear transformation
Discriminator(dim_in, dim_hidden, num_layer=2, use_bn=False, use_pair_norm=False, dropout=0.5, act_fn='leakyrelu', sigmoid_output=True, norm_mode='PN', norm_scale=1.0)

Bases: Module

Multi-layer discriminator network with configurable normalization and activation.

Parameters:
  • dim_in (int) –

    Input feature dimension

  • dim_hidden (int) –

    Hidden layer dimension

  • num_layer (int, default: 2 ) –

    Number of linear layers. Default: 2

  • use_bn (bool, default: False ) –

    Whether to use batch normalization. Default: False

  • use_pair_norm (bool, default: False ) –

    Whether to use PairNorm normalization. Default: False

  • dropout (float, default: 0.5 ) –

    Dropout rate. Default: 0.5

  • act_fn (str, default: 'leakyrelu' ) –

    Activation function type ('relu', 'leakyrelu', 'tanh', 'sigmoid'). Default: 'leakyrelu'

  • sigmoid_output (bool, default: True ) –

    Whether to apply sigmoid to final output. Default: True

  • norm_mode (str, default: 'PN' ) –

    PairNorm mode if use_pair_norm is True. Default: 'PN'

  • norm_scale (float, default: 1.0 ) –

    Normalization scaling factor. Default: 1.
Notes

Architecture:

  • Multiple linear layers with configurable width
  • Optional BatchNorm or PairNorm after each intermediate layer
  • Configurable activation functions
  • Optional sigmoid output for binary classification
forward(z)

Forward pass of the discriminator.

Parameters:
  • z (Tensor) –

    Input features of shape [batch_size, dim_in]
Returns:
  • Tensor –

    Discrimination probabilities of shape [batch_size, 1]

    • If sigmoid_output=True: Range [0,1]
    • If sigmoid_output=False: Unbounded logits
Notes

Process for each intermediate layer:

  • Linear transformation

  • Normalization (if enabled):

    • BatchNorm (if use_bn=True)
    • PairNorm (if use_pair_norm=True)

  • Activation function

Final layer:

  • Linear transformation
  • Optional sigmoid activation
GraphEncoder(dim_in, dim_out, dim_hidden=64, layer_num=2, root_weight=True, norm_mode='PN-SCS', norm_scale=1, log_softmax=False)

Bases: Module

Multi-layer GraphSAGE encoder with normalization and regularization.

Parameters:
  • dim_in (int) –

    Input feature dimension

  • dim_out (int) –

    Output feature dimension

  • dim_hidden (int, default: 64 ) –

    Hidden layer dimension. Default: 64

  • layer_num (int, default: 2 ) –

    Number of GraphSAGE layers. Default: 2

  • root_weight (bool, default: True ) –

    Whether to add self-loops in GraphSAGE. Default: True

  • norm_mode (str, default: 'PN-SCS' ) –

    PairNorm mode ('None', 'PN', 'PN-SI', 'PN-SCS'). Default: 'PN-SCS'

  • norm_scale (float, default: 1 ) –

    Normalization scaling factor. Default: 1

  • log_softmax (bool, default: False ) –

    Whether to apply log_softmax to output. Default: False
Notes

Architecture:

  • Multiple GraphSAGE convolution layers
  • PairNorm after each intermediate layer
  • ReLU activation and dropout
  • Optional log_softmax output
forward(x, edge_index)

Forward pass of the GraphEncoder.

Parameters:
  • x (Tensor) –

    Node feature matrix of shape [num_nodes, dim_in]

  • edge_index (Tensor) –

    Graph connectivity in COO format of shape [2, num_edges]
Returns:
  • Tensor –

    Encoded node features of shape [num_nodes, dim_out]
Notes

For intermediate layers:

  • GraphSAGE convolution
  • PairNorm
  • ReLU activation
  • Dropout (p=0.5)

Final layer:

  • GraphSAGE convolution only

Optional log_softmax

KBLBase(data_src, data_tar, device, k_cross=20, k_within=6, epsilon=0.5, bridge_batch_size=1000, dim_hidden=64, num_layer=2, num_epoch=200, lr=0.001, weight_decay=0.005, source_clf=True, norm_mode='PN', norm_scale=1.0)

Bases: Module

Base class for KBL.

Parameters:
  • data_src (Data) –

    Source domain graph data

  • data_tar (Data) –

    Target domain graph data

  • device (device) –

    Device to use for computation

  • k_cross (int, default: 20 ) –

    Number of cross-domain edges per node. Default: 20

  • k_within (int, default: 6 ) –

    Number of within-domain edges per node. Default: 6

  • epsilon (float, default: 0.5 ) –

    Similarity threshold. Default: 0.5

  • bridge_batch_size (int, default: 1000 ) –

    Batch size for bridge construction. Default: 1000

  • dim_hidden (int, default: 64 ) –

    Hidden dimension size. Default: 64

  • num_layer (int, default: 2 ) –

    Number of GNN layers. Default: 2

  • num_epoch (int, default: 200 ) –

    Number of training epochs. Default: 200

  • lr (float, default: 0.001 ) –

    Learning rate. Default: 0.001

  • weight_decay (float, default: 0.005 ) –

    Weight decay for optimization. Default: 5e-3

  • source_clf (bool, default: True ) –

    Whether to use source classifier. Default: True

  • norm_mode (str, default: 'PN' ) –

    Normalization mode. Default: 'PN'

  • norm_scale (float, default: 1.0 ) –

    Normalization scale factor. Default: 1.
Notes

Architecture components:

  • BridgedGraph for domain connection
  • GNNBase for node classification
  • Memory-efficient batch processing
forward(data)

Forward pass through the KBL model.

Parameters:
  • data (Data) –

    Input graph data containing:

    • x: Node features
    • edge_index: Graph connectivity
Returns:
  • Tensor –

    Log probabilities of node classifications Shape: [num_nodes, num_classes]
Notes

Processes input through GNN layers for node classification using the enhanced graph structure

get_bridged_graph()

Generate and return a bridged graph connecting source and target domains.

Returns:
  • Data –

    Merged graph containing:

    • Combined node features
    • Original edges from both domains
    • Cross-domain bridge edges
    • Within-domain enhancement edges
    • Updated masks for splits
Notes
  • Trains the bridge model using adversarial learning
  • Generates bridged graph with specified batch size
PairEnumerator(data, mode='train')

Node pair sampling and enumeration utility for graph data.

Parameters:
  • data (Data) –

    Input graph data

  • mode (str, default: 'train' ) –

    Sampling mode ('train', 'val', 'test', 'all'). Default: 'train'
Notes
  • Organizes nodes by class for stratified sampling
  • Supports different data splits (train/val/test)
  • Handles missing labels (-1)
  • Creates class-wise node buckets for efficient sampling
balanced_sampling(max_class_num=2, sample_size=40000, shuffle=True)

Generate balanced node pairs across classes.

Parameters:
  • max_class_num (int, default: 2 ) –

    Maximum number of classes to sample from. Default: 2

  • sample_size (int, default: 40000 ) –

    Total number of pairs to generate. Default: 40000

  • shuffle (bool, default: True ) –

    Whether to shuffle the pairs. Default: True
Returns:
  • tuple –

    Contains:

    • torch.Tensor: First nodes in pairs
    • torch.Tensor: Second nodes in pairs
Notes
  • Generates both same-class and cross-class pairs
  • Same-class pairs: 50% of sample_size
  • Cross-class pairs: 50% distributed across class combinations
  • Sampling with replacement within classes
pair_enumeration(x1, x2)

Generate all possible pairs between two sets of nodes.

Parameters:
  • x1 (Tensor) –

    First set of nodes, shape [B,D]

  • x2 (Tensor) –

    Second set of nodes, shape [B,D]
Returns:
  • Tensor –

    Concatenated pairs, shape [B*B,2D]

Examples:

Input: x1 = [[a], [b]] Output: [[a,a], [b,a], [a,b], [b,b]]

Notes
  • Ensures inputs are 2-dimensional tensors
  • Creates cartesian product of input sets
  • Memory-efficient implementation using repeat operations
sampling(max_class_num=2, sample_size=40000, shuffle=True)

Sample node pairs using square root sampling strategy.

Parameters:
  • max_class_num (int, default: 2 ) –

    Maximum number of classes to sample from. Default: 2

  • sample_size (int, default: 40000 ) –

    Total number of pairs to generate. Default: 40000

  • shuffle (bool, default: True ) –

    Whether to shuffle the pairs. Default: True
Returns:
  • tuple –

    Contains:

    • torch.Tensor: First nodes in pairs
    • torch.Tensor: Second nodes in pairs
Notes
  • Uses sqrt(sample_size) nodes per class
  • Generates pairs through enumeration
  • Sampling with replacement within classes
  • Optional shuffling of final pairs
PairNorm(mode='PN', scale=10)

Bases: Module

Implementation of PairNorm for graph neural networks.

PairNorm is a normalization layer that helps prevent over-smoothing in deep GNNs by normalizing node features both individually and across the graph.

Parameters:
  • mode (str, default: 'PN' ) –

    Normalization mode. Options:

    • 'None': No normalization
    • 'PN': Original PairNorm
    • 'PN-SI': Scale-Individually version
    • 'PN-SCS': Scale-and-Center-Simultaneously version Default: 'PN'
  • scale (float, default: 10 ) –

    Scaling factor for normalized features. Default: 10
Notes

The 'PN-SCS' mode is an extension of the original paper that works well in practice, especially for GCN and GAT architectures.

mode:

  • 'None' : No normalization
  • 'PN' : Original version
  • 'PN-SI' : Scale-Individually version
  • 'PN-SCS' : Scale-and-Center-Simultaneously version ('SCS'-mode is not in the paper but we found it works well in practice, especially for GCN and GAT.)

PairNorm is typically used after each graph convolution operation.

forward(x)

Apply PairNorm to input features.

Parameters:
  • x (Tensor) –

    Input node features of shape [num_nodes, num_features]
Returns:
  • Tensor –

    Normalized node features of shape [num_nodes, num_features]
Notes

Implements four normalization strategies:

  • 'None': Returns input unchanged
  • 'PN': Centers and scales features using mean row norm
  • 'PN-SI': Centers and scales features individually
  • 'PN-SCS': Scales individually then centers features

A small epsilon (1e-6) is added for numerical stability during normalization.

Similar(in_channels, num_clf_classes, dropout=0.6, use_clf=True)

Bases: Module

Similarity computation module with optional classification capability.

Parameters:
  • in_channels (int) –

    Input feature dimension

  • num_clf_classes (int) –

    Number of classification classes

  • dropout (float, default: 0.6 ) –

    Dropout rate. Default: 0.6

  • use_clf (bool, default: True ) –

    Whether to include classification branch. Default: True
Notes

Architecture components:

  • Bias attention network (128->64->128)
  • Feature transformation network with BatchNorm
  • Optional classification layer
  • Cosine similarity computation
forward(x, idx1, idx2)

Forward pass for within-domain similarity computation.

Parameters:
  • x (Tensor) –

    Node features

  • idx1 (Tensor) –

    First set of node indices

  • idx2 (Tensor) –

    Second set of node indices
Returns:
  • tuple –

    Contains:

    • torch.Tensor: Similarity scores [batch_size, 1]
    • torch.Tensor: Classification log probabilities (if use_clf=True)
forward_cross_domain(x_src, x_tar, idx1, idx2)

Forward pass for cross-domain similarity computation.

Parameters:
  • x_src (Tensor) –

    Source domain features

  • x_tar (Tensor) –

    Target domain features

  • idx1 (Tensor) –

    Source node indices

  • idx2 (Tensor) –

    Target node indices
Returns:
  • tuple –

    Contains:

    • torch.Tensor: Similarity scores
    • torch.Tensor: Source classification log probabilities (if use_clf=True)
    • torch.Tensor: Target classification log probabilities (if use_clf=True)
similarity(x, idx1, idx2)

Compute similarity between nodes within same domain.

Parameters:
  • x (Tensor) –

    Node features

  • idx1 (Tensor) –

    First set of node indices

  • idx2 (Tensor) –

    Second set of node indices
Returns:
  • Tensor –

    Similarity scores between selected node pairs
similarity_cross_domain(x_src, x_tar, idx1, idx2)

Compute similarity between source and target domain nodes.

Parameters:
  • x_src (Tensor) –

    Source domain features

  • x_tar (Tensor) –

    Target domain features

  • idx1 (Tensor) –

    Source node indices

  • idx2 (Tensor) –

    Target node indices
Returns:
  • Tensor –

    Similarity scores between selected node pairs
similarity_cross_domain_batch(x_src, x_tar, idx1, idx2)

Compute cross-domain similarity scores in batches.

Parameters:
  • x_src (Tensor) –

    Source domain features

  • x_tar (Tensor) –

    Target domain features

  • idx1 (Tensor) –

    Source node indices

  • idx2 (Tensor) –

    Target node indices
Returns:
  • Tensor –

    Batched similarity scores between node pairs
Notes
  • Processes in batches of 100 for memory efficiency
  • Uses attention-enhanced cosine similarity
SourceLearner(data, dim_hidden=64, norm_mode='None', norm_scale=1, use_clf=True)

Bases: Module

Source domain learning module combining graph encoding and similarity matching.

Parameters:
  • data (Data) –

    Source domain graph data

  • dim_hidden (int, default: 64 ) –

    Hidden dimension size. Default: 64

  • norm_mode (str, default: 'None' ) –

    Normalization mode for GraphEncoder. Default: 'None'

  • norm_scale (float, default: 1 ) –

    Normalization scale factor. Default: 1

  • use_clf (bool, default: True ) –

    Whether to use classification. Default: True
Notes

Architecture components:

  • GraphEncoder backbone for feature extraction
  • Similar module for similarity computation
  • Optional classification capability
forward(data, idx1, idx2, return_representation=False)

Forward pass of the source learner.

Parameters:
  • data (Data) –

    Input graph data

  • idx1 (Tensor) –

    First set of node indices

  • idx2 (Tensor) –

    Second set of node indices

  • return_representation (bool, default: False ) –

    Whether to return encoded features. Default: False
Returns:
  • tuple –

    Contains:

    • torch.Tensor: Pair similarity probabilities
    • torch.Tensor: Classification logits
    • torch.Tensor: Encoded features (if return_representation=True)
TargetLearnerAE(data, dim_eq_trans=128, dim_hidden=64, norm_mode='None', norm_scale=1)

Bases: Module

Target domain learner with autoencoder architecture.

Parameters:
  • data (Data) –

    Target domain graph data

  • dim_eq_trans (int, default: 128 ) –

    Dimension for equivalent transformation. Default: 128

  • dim_hidden (int, default: 64 ) –

    Hidden dimension size. Default: 64

  • norm_mode (str, default: 'None' ) –

    Normalization mode. Default: 'None'

  • norm_scale (float, default: 1 ) –

    Normalization scale factor. Default: 1
Notes

Architecture components:

  • Equivalent transformation layer
  • GraphEncoder for feature encoding
  • Decoder for feature reconstruction
decode(z)

Decode latent features.

Parameters:
  • z (Tensor) –

    Encoded features
Returns:
  • Tensor –

    Reconstructed features in tanh-normalized space
encode(data)

Encode input graph data.

Parameters:
  • data (Data) –

    Input graph data
Returns:
  • tuple –

    Contains: - torch.Tensor: Encoded features - torch.Tensor: Transformed input features
forward(data)

Forward pass through the autoencoder.

Parameters:
  • data (Data) –

    Input graph data
Returns:
  • tuple –

    Contains:

    • torch.Tensor: Transformed input features
    • torch.Tensor: Encoded features
    • torch.Tensor: Reconstructed features