PairAlign(in_dim, hid_dim, num_classes, num_layers=2, cls_dim=128, cls_layers=2, dropout=0.0, backbone='GS', pooling='mean', ew_type='pseudobeta', rw_lmda=1.0, ls_lambda=1.0, lw_lambda=0.005, label_rw=False, edge_rw=False, ew_start=0, ew_freq=10, lw_start=0, lw_freq=10, gamma_reg=0.0001, weight_CE_src=False, bn=False, act=F.relu, weight_decay=0.0001, lr=0.001, epoch=200, device='cuda:0', batch_size=0, num_neigh=-1, verbose=2, **kwargs)

Bases: BaseGDA

Pairwise Alignment Improves Graph Domain Adaptation (ICML-24).

Parameters:
  • in_dim (int) –

    Input feature dimension.

  • hid_dim (int) –

    Hidden dimension of model.

  • num_classes (int) –

    Total number of classes.

  • num_layers (int, default: 2 ) –

    Total number of gnn layers in model. Default: 2.

  • cls_dim (int, default: 128 ) –

    Hidden dimension for classification layer. Default: 128.

  • cls_layers (int, default: 2 ) –

    Total number of cls layers in model. Default: 2.

  • ew_start (int, default: 0 ) –

    Starting epoch for edge reweighting. Default: 0.

  • ew_freq (int, default: 10 ) –

    Frequency for edge reweighting. Default: 10.

  • lw_start (int, default: 0 ) –

    Starting epoch for label reweighting. Default: 0.

  • lw_freq (int, default: 10 ) –

    Frequency for label reweighting. Default: 10.

  • dropout (float, default: 0.0 ) –

    Dropout rate. Default: 0..

  • pooling (string, default: 'mean' ) –

    Aggregation in gnn. Default: mean.

  • ew_type (string, default: 'pseudobeta' ) –

    Use the true edge weight or not. Default: pseudobeta.

  • rw_lmda (float, default: 1.0 ) –

    Trade-off parameter for edge reweight. Default: 1.0.

  • ls_lambda (float, default: 1.0 ) –

    Regularize the distance to 1 in w optimization. Default: 1.0.

  • lw_lambda (float, default: 0.005 ) –

    Regularize the distance to 1 in beta optimization. Default: 0.005.

  • label_rw (bool, default: False ) –

    Reweight the label or not. Default: False.

  • edge_rw (bool, default: False ) –

    Reweight the edge in source graph or not. Default: False.

  • gamma_reg (float, default: 0.0001 ) –

    Mimic the variance of the edges to normalize the weight. Default: 1e-4.

  • weight_CE_src (bool, default: False ) –

    Reweight the loss by src class or not. Default: False.

  • backbone (string, default: 'GS' ) –

    The backbone of PairAlign. Default: GS.

  • weight_decay (float, default: 0.0001 ) –

    Weight decay (L2 penalty). Default: 0.0001.

  • act (callable activation function or None, default: relu ) –

    Activation function if not None. Default: torch.nn.functional.relu.

  • lr (float, default: 0.001 ) –

    Learning rate. Default: 0.001.

  • bn (bool, default: False ) –

    Batch normalization or not. Default: False.

  • epoch (int, default: 200 ) –

    Maximum number of training epoch. Default: 200.

  • device (str, default: 'cuda:0' ) –

    GPU or CPU. Default: cuda:0.

  • batch_size (int, default: 0 ) –

    Minibatch size, 0 for full batch training. Default: 0.

  • num_neigh (int, default: -1 ) –

    Number of neighbors in sampling, -1 for all neighbors. Default: -1.

  • verbose (int, default: 2 ) –

    Verbosity mode. Range in [0, 3]. Larger value for printing out more log information. Default: 2.

  • **kwargs –

    Other parameters for the model.
CE_loss(pred, label)

Standard cross-entropy loss computation.

Parameters:
  • pred (Tensor) –

    Model predictions.

  • label (Tensor) –

    Ground truth labels.
Returns:
  • Tensor –

    Computed cross-entropy loss.
CE_loss_weight(pred, label, label_weight, weight_src)

Weighted cross-entropy loss with label-specific weights.

Parameters:
  • pred (Tensor) –

    Model predictions.

  • label (Tensor) –

    Ground truth labels.

  • label_weight (Tensor) –

    Weights for each label.

  • weight_src (bool) –

    Whether to use source domain class weights.
Returns:
  • Tensor –

    Weighted loss value.
Notes
  • Combines class weights and label weights
  • Supports optional source domain weighting
  • Normalizes by number of classes
  • Handles both balanced and imbalanced scenarios
CE_srcweight_loss(pred, label)

Weighted cross-entropy loss for source domain.

Parameters:
  • pred (Tensor) –

    Model predictions.

  • label (Tensor) –

    Ground truth labels.
Returns:
  • Tensor –

    Weighted cross-entropy loss.
Notes
  • Computes inverse frequency weights for classes
  • Applies class-specific weights to loss
  • Handles class imbalance in source domain
LS_optimization(cov, muhat_tgt, mu_src, lambda_reg)

Perform least squares optimization for weight computation.

Parameters:
  • cov (Tensor) –

    Covariance matrix between source and target.

  • muhat_tgt (Tensor) –

    Target domain mean features.

  • mu_src (Tensor) –

    Source domain mean features.

  • lambda_reg (float) –

    Regularization parameter.
Returns:
  • ndarray –

    Optimal transformation weights.
Notes

  • Solves constrained optimization problem:

    • Minimizes L2 distance between transformed source and target
    • Includes regularization towards uniform weights
    • Maintains probability constraints

  • Prints covariance matrix rank for debugging

  • Uses CVXPY with SCS solver
cal_edge_prob_sep(src_graph, tgt_graph)

Calculate separate edge probabilities for source and target domains.

Parameters:
  • src_graph (Data) –

    Source domain graph.

  • tgt_graph (Data) –

    Target domain graph.
Returns:
  • tuple –

    Contains: - src_edge_prob : torch.Tensor Edge probabilities in source domain. - tgt_edge_prob : torch.Tensor Edge probabilities in target domain.
Notes

Implementation steps:

  1. Converts sparse adjacency to dense format
  2. Creates one-hot label encodings using sparse matrices
  3. Computes class-conditional edge probabilities
  4. Normalizes by node counts per class
calc_edge_rw_pseudo(src_graph, tgt_graph, yhat_src, yhat_tgt)

Calculate edge reweighting using pseudo-labels.

Parameters:
  • src_graph (Data) –

    Source domain graph.

  • tgt_graph (Data) –

    Target domain graph.

  • yhat_src (Tensor) –

    Source domain predictions.

  • yhat_tgt (Tensor) –

    Target domain predictions.
Returns:
  • tuple –

    Contains: - diff : float L1 difference between computed and true edge weights. - beta_diff : float L1 difference between computed and true beta values.
Notes

Implementation steps:

Edge Feature Computation

  • Constructs sparse edge features for both domains
  • Uses einsum for efficient tensor operations
  • Handles class-conditional edge probabilities

Distribution Alignment

  • Computes covariance and mean statistics
  • Performs least squares optimization
  • Applies ratio-based weight calculation

Weight Processing

  • Handles numerical edge cases (inf, nan)
  • Updates graph edge weights
  • Computes alignment metrics
calc_label_rw(y_src, y_hat_tgt, cov, lambda_reg)

Calculate label reweighting coefficients using constrained optimization.

Parameters:
  • y_src (Tensor) –

    Source domain label distribution.

  • y_hat_tgt (Tensor) –

    Target domain predicted label distribution.

  • cov (Tensor) –

    Covariance matrix between predictions and true labels.

  • lambda_reg (float) –

    Regularization parameter for weight deviation.
Returns:
  • ndarray –

    Optimal label weights that align source and target distributions.
Notes

  • Solves a quadratic program with:

    • L2 loss between transformed and target distributions
    • L2 regularization towards uniform weights
    • Sum-to-one and non-negativity constraints

  • Uses CVXPY with SCS solver for optimization

calc_ratio_weight(src_graph, kmm_weight, gamma_reg)

Calculate edge ratio weights using kernel mean matching.

Parameters:
  • src_graph (Data) –

    Source domain graph data.

  • kmm_weight (ndarray) –

    Kernel mean matching weights.

  • gamma_reg (float) –

    Regularization parameter for numerical stability.
Returns:
  • tuple –

    Contains: - gamma : numpy.ndarray Edge ratio weights. - p_ei_tgt_reg : numpy.ndarray Regularized target edge probabilities.
Notes
  • Computes class-conditional edge probabilities
  • Applies regularization for stability
  • Handles edge probability ratios between domains
  • Maintains probability constraints
calc_true_beta(src_graph, tgt_graph)

Calculate true beta values for edge distribution alignment.

Parameters:
  • src_graph (Data) –

    Source domain graph.

  • tgt_graph (Data) –

    Target domain graph.
Notes
  • Creates edge class matrices for both domains
  • Computes edge class probabilities
  • Calculates ratio between target and source probabilities
  • Handles numerical edge cases (inf, nan)
  • Stores results in source graph for reference
  • Maintains class-pair relationships in edge classification
calc_true_ew(src_graph, tgt_graph)

Calculate true edge weights between source and target domains.

Parameters:
  • src_graph (Data) –

    Source domain graph.

  • tgt_graph (Data) –

    Target domain graph.
Notes
  • Computes edge probability ratios
  • Handles numerical stability with gamma regularization
  • Stores true edge weights for reference
calc_true_lw(src_graph, tgt_graph)

Calculate true label weights between domains.

Parameters:
  • src_graph (Data) –

    Source domain graph.

  • tgt_graph (Data) –

    Target domain graph.
Notes
  • Computes class distribution in both domains
  • Calculates ratio of target to source distributions
  • Stores true weights and initializes current weights
  • Uses one-hot encoding for label processing
fit(source_data, target_data)

Train the PairAlign model on source and target domain data.

Parameters:
  • source_data (Data) –

    Source domain graph data.

  • target_data (Data) –

    Target domain graph data.
Notes

Training process includes:

Initial Setup

  • Edge weight initialization
  • Data loader configuration
  • True weight calculations (edge, label, beta)

Training Loop

  • Forward pass with reweighting
  • Loss computation and optimization
  • Periodic weight updates
  • Performance monitoring
forward_model(source_data, target_data, epoch)

Forward pass of the PairAlign model.

Parameters:
  • source_data (Data) –

    Source domain graph data.

  • target_data (Data) –

    Target domain graph data.

  • epoch (int) –

    Current training epoch.
Returns:
  • tuple –

    Contains: - loss : torch.Tensor Combined loss from classification. - pred_src : torch.Tensor Source domain predictions. - pred_tgt : torch.Tensor Target domain predictions.
Notes
  • Updates edge weights periodically if edge_rw is enabled
  • Updates label weights periodically if label_rw is enabled
  • Handles both weighted and unweighted classification losses
init_model(**kwargs)

Initialize the PairAlign base model.

Parameters:
  • **kwargs –

    Additional parameters for model initialization.
Returns:
  • ReweightGNN –

    Initialized model with specified architecture parameters.
Notes

Configures model with:

  • GNN backbone with specified dimensions
  • Classification layers
  • Reweighting mechanisms for edges and labels
predict(data)

Make predictions on input data.

Parameters:
  • data (Data) –

    Input graph data.
Returns:
  • tuple –

    Contains: - logits : torch.Tensor Model predictions. - labels : torch.Tensor True labels.
Notes
  • Evaluates model in inference mode
  • Returns both predictions and ground truth
process_graph(data)

Process input graph data.

Parameters:
  • data (Data) –

    Input graph to be processed.
Notes

Placeholder method as preprocessing is handled by:

  • Edge weight calculation
  • Label weight computation
  • Beta value estimation
update_ew(src_data, tgt_data, pred_src, pred_tgt)

Update edge weights based on domain alignment.

Parameters:
  • src_data (Data) –

    Source domain graph data.

  • tgt_data (Data) –

    Target domain graph data.

  • pred_src (Tensor) –

    Source domain predictions.

  • pred_tgt (Tensor) –

    Target domain predictions.
Notes

  • Supports multiple edge weight calculation methods:

    • pseudobeta: Dynamic weight calculation
    • truth: Uses true edge weights
    • other: Ratio-based weight calculation

  • Updates edge weights in source graph

update_lw(src_data, tgt_data, pred_src, pred_tgt)

Update label weights for domain alignment.

Parameters:
  • src_data (Data) –

    Source domain graph data.

  • tgt_data (Data) –

    Target domain graph data.

  • pred_src (Tensor) –

    Source domain predictions.

  • pred_tgt (Tensor) –

    Target domain predictions.
Returns:
  • ndarray –

    Updated label weights.
Notes
  • Computes label distribution alignment
  • Uses least squares optimization with regularization
  • Maintains probability constraints