[Feature] Add bias to edge decoder (#1125)

*Issue #, if available:*
An extension of issue #335 

*Description of changes:*
Adding bias to edge decoders.

Regression test results after my changes:
```
Dataset		Test Performance	Validation Performance	Total Time	Epoch Time/Seconds
OGB-MAG-NC	Acc: 0.4070		Acc: 0.4151		0:25:08		69.963
OGB-PROD-NC	Acc: 0.7254		Acc: 0.9137		0:46:33		108.104
```

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

python/graphstorm/config/argument.py

@@ -1752,15 +1752,15 @@ def log_report_frequency(self):
     ###################### Model training related ######################
     @property
     def decoder_bias(self):
-        """ Node decoder bias. decoder_bias must be a boolean. Default is False.
+        """ Decoder bias. decoder_bias must be a boolean. Default is True.
         """
         # pylint: disable=no-member
         if hasattr(self, "_decoder_bias"):
             assert self._decoder_bias in [True, False], \
                 "decoder_bias should be in [True, False]"
             return self._decoder_bias
-        # By default, node decoder bias is False
-        return False
+        # By default, decoder bias is True
+        return True
 
     @property
     def dropout(self):

python/graphstorm/gsf.py

@@ -439,7 +439,8 @@ def create_builtin_reconstruct_efeat_decoder(g, decoder_input_dim, config, train
     decoder = EdgeRegression(decoder_input_dim,
                              target_etype=target_etype,
                              out_dim=feat_dim,
-                             dropout=dropout)
+                             dropout=dropout,
+                             use_bias=config.decoder_bias)
 
     loss_func = RegressionLossFunc()
     return decoder, loss_func
@@ -617,14 +618,16 @@ def create_builtin_edge_decoder(g, decoder_input_dim, config, train_task):
                                      dropout_rate=dropout,
                                      regression=False,
                                      target_etype=target_etype,
-                                     norm=config.decoder_norm)
+                                     norm=config.decoder_norm,
+                                     use_bias=config.decoder_bias)
         elif decoder_type == "MLPDecoder":
             decoder = MLPEdgeDecoder(decoder_input_dim,
                                      num_classes,
                                      multilabel=config.multilabel,
                                      target_etype=target_etype,
                                      num_ffn_layers=config.num_ffn_layers_in_decoder,
-                                     norm=config.decoder_norm)
+                                     norm=config.decoder_norm,
+                                     use_bias=config.decoder_bias)
         elif decoder_type == "MLPEFeatEdgeDecoder":
             decoder_edge_feat = config.decoder_edge_feat
             assert decoder_edge_feat is not None, \
@@ -648,7 +651,8 @@ def create_builtin_edge_decoder(g, decoder_input_dim, config, train_task):
                 target_etype=target_etype,
                 dropout=config.dropout,
                 num_ffn_layers=config.num_ffn_layers_in_decoder,
-                norm=config.decoder_norm)
+                norm=config.decoder_norm,
+                use_bias=config.decoder_bias)
         else:
             assert False, f"decoder {decoder_type} is not supported."
 
@@ -680,15 +684,17 @@ def create_builtin_edge_decoder(g, decoder_input_dim, config, train_task):
                                      target_etype=target_etype,
                                      dropout_rate=dropout,
                                      regression=True,
-                                     norm=config.decoder_norm)
+                                     norm=config.decoder_norm,
+                                     use_bias=config.decoder_bias)
         elif decoder_type == "MLPDecoder":
             decoder = MLPEdgeDecoder(decoder_input_dim,
                                      1,
                                      multilabel=False,
                                      target_etype=target_etype,
                                      regression=True,
                                      num_ffn_layers=config.num_ffn_layers_in_decoder,
-                                     norm=config.decoder_norm)
+                                     norm=config.decoder_norm,
+                                     use_bias=config.decoder_bias)
         elif decoder_type == "MLPEFeatEdgeDecoder":
             decoder_edge_feat = config.decoder_edge_feat
             assert decoder_edge_feat is not None, \
@@ -713,7 +719,8 @@ def create_builtin_edge_decoder(g, decoder_input_dim, config, train_task):
                 dropout=config.dropout,
                 regression=True,
                 num_ffn_layers=config.num_ffn_layers_in_decoder,
-                norm=config.decoder_norm)
+                norm=config.decoder_norm,
+                use_bias=config.decoder_bias)
         else:
             assert False, "decoder not supported"
         loss_func = RegressionLossFunc()

python/graphstorm/model/__init__.py

@@ -61,7 +61,8 @@
                            LinkPredictWeightedRotatEDecoder,
                            LinkPredictTransEDecoder,
                            LinkPredictContrastiveTransEDecoder,
-                           LinkPredictWeightedTransEDecoder)
+                           LinkPredictWeightedTransEDecoder,
+                           EdgeRegression)
 
 from .gnn_encoder_base import GraphConvEncoder
 

python/graphstorm/model/edge_decoder.py

@@ -132,6 +132,11 @@ class DenseBiDecoder(GSEdgeDecoder):
     norm: str
         Normalization methods. Not used, but reserved for complex DenseBiDecoder child class
         implementation. Default: None.
+    use_bias: bool
+        Whether the edge decoder uses a bias parameter. Default: True.
+
+    .. versionchanged:: 0.4.0
+        Add a new argument "use_bias" so users can control whether decoders have bias.
     """
     def __init__(self,
                  in_units,
@@ -141,7 +146,8 @@ def __init__(self,
                  num_basis=2,
                  dropout_rate=0.0,
                  regression=False,
-                 norm=None):
+                 norm=None,
+                 use_bias=True):
         super().__init__()
 
         self.in_units = in_units
@@ -157,6 +163,7 @@ def __init__(self,
         assert isinstance(target_etype, tuple) and len(target_etype) == 3, \
             "Target etype must be a tuple of a canonical etype."
         self.target_etype = target_etype
+        self.use_bias = use_bias
 
         self._init_model()
 
@@ -171,7 +178,7 @@ def _init_model(self):
         self.dropout = nn.Dropout(self.dropout)
         self.basis_para = nn.Parameter(
             th.randn(self.num_basis, self.in_units, self.in_units))
-        self.combine_basis = nn.Linear(self.num_basis, basis_out, bias=False)
+        self.combine_basis = nn.Linear(self.num_basis, basis_out, bias=self.use_bias)
         self.reset_parameters()
 
         if self.regression:
@@ -331,13 +338,16 @@ class EdgeRegression(GSEdgeDecoder):
     norm: str, optional
         Normalization methods. Not used, but reserved for complex edge regression.
         implementation. Default: None.
+    use_bias: bool
+        Whether the edge decoder uses a bias parameter. Default: True.
     """
     def __init__(self,
                  h_dim,
                  target_etype,
                  out_dim=1,
                  dropout=0,
-                 norm=None):
+                 norm=None,
+                 use_bias=True):
         super(EdgeRegression, self).__init__()
         self._h_dim = h_dim
         self._out_dim = out_dim
@@ -349,6 +359,7 @@ def __init__(self,
             "Target etype must be a tuple of a canonical etype," \
             f"e.g., (src_ntype, etype, dst_ntype), but got {target_etype}."
         self._target_etype = target_etype
+        self._use_bias = use_bias
 
         self._init_model()
 
@@ -360,7 +371,7 @@ def _init_model(self):
         if self._norm is not None:
             logging.warning("Embedding normalization (batch norm or layer norm) "
                             "is not supported in EdgeRegression")
-        self.linear = nn.Linear(h_dim * 2, h_dim, bias=True)
+        self.linear = nn.Linear(h_dim * 2, h_dim, bias=self._use_bias)
         self.relu = nn.ReLU()
         self.dropout = nn.Dropout(self._dropout)
         self.regression_head = nn.Linear(h_dim, out_dim, bias=True)
@@ -500,6 +511,11 @@ class MLPEdgeDecoder(GSEdgeDecoder):
     norm: str
         Normalization methods. Not used, but reserved for complex MLPEdgeDecoder child class
         implementation. Default: None.
+    use_bias: bool
+        Whether the edge decoder uses a bias parameter. Default: True.
+
+    .. versionchanged:: 0.4.0
+        Add a new argument "use_bias" so users can control whether decoders have bias.
     """
     def __init__(self,
                  h_dim,
@@ -510,7 +526,8 @@ def __init__(self,
                  dropout=0,
                  regression=False,
                  num_ffn_layers=0,
-                 norm=None):
+                 norm=None,
+                 use_bias=True):
         super(MLPEdgeDecoder, self).__init__()
         self.h_dim = h_dim
         self.multilabel = multilabel
@@ -526,6 +543,7 @@ def __init__(self,
         assert isinstance(target_etype, tuple) and len(target_etype) == 3, \
             "Target etype must be a tuple of a canonical etype."
         self.target_etype = target_etype
+        self.use_bias = use_bias
 
         self._init_model()
 
@@ -543,6 +561,8 @@ def _init_model(self):
 
         # Here we assume the source and destination nodes have the same dimension.
         self.decoder = nn.Parameter(th.randn(self.h_dim * 2, self.out_dim))
+        if self.use_bias:
+            self.bias = nn.Parameter(th.zeros(self.out_dim))
         assert self.num_hidden_layers == 1, "More than one layers not supported"
         nn.init.xavier_uniform_(self.decoder,
                                 gain=nn.init.calculate_gain('relu'))
@@ -561,8 +581,8 @@ def _compute_logits(self, g, h):
                 The dictionary containing the embeddings
             Returns
             -------
-            th.Tensor
-                Output of forward
+            out: th.Tensor
+                Output of forward.
         """
         with g.local_scope():
             u, v = g.edges(etype=self.target_etype)
@@ -574,6 +594,8 @@ def _compute_logits(self, g, h):
             if self.num_ffn_layers > 0:
                 h = self.ngnn_mlp(h)
             out = th.matmul(h, self.decoder)
+            if self.use_bias:
+                out = out + self.bias
         return out
 
     # pylint: disable=unused-argument
@@ -706,6 +728,11 @@ class MLPEFeatEdgeDecoder(MLPEdgeDecoder):
     norm: str
         Normalization methods. Not used, but reserved for complex MLPEFeatEdgeDecoder child
         class implementation. Default: None.
+    use_bias: bool
+        Whether the edge decoder uses a bias parameter. Default: True.
+
+    .. versionchanged:: 0.4.0
+        Add a new argument "use_bias" so users can control whether decoders have bias.
     """
     def __init__(self,
                  h_dim,
@@ -716,7 +743,8 @@ def __init__(self,
                  dropout=0,
                  regression=False,
                  num_ffn_layers=0,
-                 norm=None):
+                 norm=None,
+                 use_bias=True):
         self.feat_dim = feat_dim
         super(MLPEFeatEdgeDecoder, self).__init__(h_dim=h_dim,
                                                   out_dim=out_dim,
@@ -725,7 +753,8 @@ def __init__(self,
                                                   dropout=dropout,
                                                   regression=regression,
                                                   num_ffn_layers=num_ffn_layers,
-                                                  norm=norm)
+                                                  norm=norm,
+                                                  use_bias=use_bias)
 
     def _init_model(self):
         """ Init decoder model
@@ -747,6 +776,8 @@ def _init_model(self):
         # combine output of nn_decoder and feat_decoder
         self.combine_decoder = nn.Parameter(th.randn(self.h_dim * 2, self.h_dim))
         self.decoder = nn.Parameter(th.randn(self.h_dim, self.out_dim))
+        if self.use_bias:
+            self.bias = nn.Parameter(th.zeros(self.out_dim))
         self.dropout = nn.Dropout(self.dropout)
 
         self.nn_decoder_norm = None
@@ -784,10 +815,12 @@ def _compute_logits(self, g, h, e_h):
                 The minibatch graph
             h: dict of Tensors
                 The dictionary containing the embeddings
+            e_h: dict of Tensor
+                The input edge embeddings in the format of {(src_ntype, etype, dst_ntype): emb}.
             Returns
             -------
-            th.Tensor
-                Output of forward
+            out: Tensor
+                Output of forward.
         """
         assert e_h is not None, "edge feature is required"
         with g.local_scope():
@@ -819,6 +852,8 @@ def _compute_logits(self, g, h, e_h):
                 combine_h = self.combine_norm(combine_h)
             combine_h = self.relu(combine_h)
             out = th.matmul(combine_h, self.decoder)
+            if self.use_bias:
+                out = out + self.bias
 
         return out
 

python/graphstorm/model/node_decoder.py

@@ -39,15 +39,18 @@ class EntityClassifier(GSLayer):
         Normalization methods. Not used, but reserved for complex node classifier
         implementation. Default: None.
     use_bias: bool
-        Whether the node decoder uses a bias parameter. Default: False.
+        Whether the node decoder uses a bias parameter. Default: True.
+
+    .. versionchanged:: 0.4.0
+        Add a new argument "use_bias" so users can control whether decoders have bias.
     """
     def __init__(self,
                  in_dim,
                  num_classes,
                  multilabel,
                  dropout=0,
                  norm=None,
-                 use_bias=False):
+                 use_bias=True):
         super(EntityClassifier, self).__init__()
         self._in_dim = in_dim
         self._num_classes = num_classes
@@ -170,14 +173,17 @@ class EntityRegression(GSLayer):
         Normalization methods. Not used, but reserved for complex node regression
         implementation. Default: None.
     use_bias: bool
-        Whether the node decoder uses a bias parameter. Default: False.
+        Whether the node decoder uses a bias parameter. Default: True.
+
+    .. versionchanged:: 0.4.0
+        Add a new argument "use_bias" so users can control whether decoders have bias.
     """
     def __init__(self,
                  h_dim,
                  dropout=0,
                  out_dim=1,
                  norm=None,
-                 use_bias=False):
+                 use_bias=True):
         super(EntityRegression, self).__init__()
         self._h_dim = h_dim
         self._out_dim = out_dim

tests/unit-tests/data_utils.py

@@ -46,6 +46,32 @@ def generate_mask(idx, length):
     th_mask = th.tensor(mask, dtype=th.bool)
     return th_mask
 
+def generate_dummy_constant_graph(in_units):
+    """
+    Generate a dummy heterogeneous graph to test edge decoder.
+
+    Return
+    -------
+    g: a heterogeneous graph.
+
+    h: node embeddings.
+
+    edge_type: graph schema ("n0", "r0", "n1")
+    """
+    u = th.tensor([0, 0])
+    v = th.tensor([1, 2])
+    edge_type = ("n0", "r0", "n1")
+    g = dgl.heterograph({
+        edge_type: (u, v)
+    })
+
+    h = {
+        "n0": th.ones(g.num_nodes("n0"), in_units),
+        "n1": th.ones(g.num_nodes("n1"), in_units)
+    }
+
+    return g, h, edge_type
+
 def generate_dummy_hetero_graph_for_efeat_gnn(is_random=True):
     """
     generate a dummy heterogeneous graph to test the get_edge_feat_size() method.