Message Passing#
Message passing module.
- class topomodelx.base.message_passing.MessagePassing(aggr_func: Literal['sum', 'mean', 'add'] = 'sum', att: bool = False, initialization: Literal['uniform', 'xavier_uniform', 'xavier_normal'] = 'xavier_uniform', initialization_gain: float = 1.414)[source]#
Define message passing.
This class defines message passing through a single neighborhood N, by decomposing it into 2 steps:
🟥 Create messages going from source cells to target cells through N.
🟧 Aggregate messages coming from different sources cells onto each target cell.
This class should not be instantiated directly, but rather inherited through subclasses that effectively define a message passing function.
This class does not have trainable weights, but its subclasses should define these weights.
- Parameters:
- aggr_funcLiteral[“sum”, “mean”, “add”], default=”sum”
Aggregation function to use.
- attbool, default=False
Whether to use attention.
- initializationLiteral[“uniform”, “xavier_uniform”, “xavier_normal”], default=”xavier_uniform”
Initialization method for the weights of the layer.
- initialization_gainfloat, default=1.414
Gain for the weight initialization.
References
[1]Hajij, Zamzmi, Papamarkou, Miolane, Guzmán-Sáenz, Ramamurthy, Birdal, Dey, Mukherjee, Samaga, Livesay, Walters, Rosen, Schaub. Topological deep learning: going beyond graph data (2023). https://arxiv.org/abs/2206.00606.
[2]Papillon, Sanborn, Hajij, Miolane. Architectures of topological deep learning: a survey on topological neural networks (2023). https://arxiv.org/abs/2304.10031.
- aggregate(x_message)[source]#
Aggregate messages on each target cell.
A target cell receives messages from several source cells. This function aggregates these messages into a single output feature per target cell.
🟧 This function corresponds to the within-neighborhood aggregation defined in [1]_ and [2]_.
- Parameters:
- x_messagetorch.Tensor, shape = (…, n_messages, out_channels)
Features associated with each message. One message is sent from a source cell to a target cell.
- Returns:
- Tensor, shape = (…, n_target_cells, out_channels)
Output features on target cells. Each target cell aggregates messages from several source cells. Assumes that all target cells have the same rank s.
- attention(x_source, x_target=None)[source]#
Compute attention weights for messages.
This provides a default attention function to the message-passing scheme.
Alternatively, users can subclass MessagePassing and overwrite the attention method in order to replace it with their own attention mechanism.
The implementation follows [1]_.
- Parameters:
- x_sourcetorch.Tensor, shape = (n_source_cells, in_channels)
Input features on source cells. Assumes that all source cells have the same rank r.
- x_targettorch.Tensor, shape = (n_target_cells, in_channels)
Input features on source cells. Assumes that all source cells have the same rank r.
- Returns:
- torch.Tensor, shape = (n_messages, 1)
Attention weights: one scalar per message between a source and a target cell.
- forward(x_source, neighborhood, x_target=None)[source]#
Forward pass.
This implements message passing for a given neighborhood:
from source cells with input features x_source,
via neighborhood defining where messages can pass,
to target cells with input features x_target.
In practice, this will update the features on the target cells.
If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.
The message passing is decomposed into two steps:
1. 🟥 Message: A message \(m_{y \rightarrow x}^{\left(r \rightarrow s\right)}\) travels from a source cell \(y\) of rank r to a target cell \(x\) of rank s through a neighborhood of \(x\), denoted \(\mathcal{N} (x)\), via the message function \(M_\mathcal{N}\):
\[m_{y \rightarrow x}^{\left(r \rightarrow s\right)} = M_{\mathcal{N}}\left(\mathbf{h}_x^{(s)}, \mathbf{h}_y^{(r)}, \Theta \right),\]where:
\(\mathbf{h}_y^{(r)}\) are input features on the source cells, called x_source,
\(\mathbf{h}_x^{(s)}\) are input features on the target cells, called x_target,
\(\Theta\) are optional parameters (weights) of the message passing function.
Optionally, attention can be applied to the message, such that:
\[m_{y \rightarrow x}^{\left(r \rightarrow s\right)} \leftarrow att(\mathbf{h}_y^{(r)}, \mathbf{h}_x^{(s)}) . m_{y \rightarrow x}^{\left(r \rightarrow s\right)}\]2. 🟧 Aggregation: Messages are aggregated across source cells \(y\) belonging to the neighborhood \(\mathcal{N}(x)\):
\[m_x^{\left(r \rightarrow s\right)} = \text{AGG}_{y \in \mathcal{N}(x)} m_{y \rightarrow x}^{\left(r\rightarrow s\right)},\]resulting in the within-neighborhood aggregated message \(m_x^{\left(r \rightarrow s\right)}\).
Details can be found in [1]_ and [2]_.
- Parameters:
- x_sourceTensor, shape = (…, n_source_cells, in_channels)
Input features on source cells. Assumes that all source cells have the same rank r.
- neighborhoodtorch.sparse, shape = (n_target_cells, n_source_cells)
Neighborhood matrix.
- x_targetTensor, shape = (…, n_target_cells, in_channels)
Input features on target cells. Assumes that all target cells have the same rank s. Optional. If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.
- Returns:
- torch.Tensor, shape = (…, n_target_cells, out_channels)
Output features on target cells. Assumes that all target cells have the same rank s.
- message(x_source, x_target=None)[source]#
Construct message from source cells to target cells.
🟥 This provides a default message function to the message passing scheme.
Alternatively, users can subclass MessagePassing and overwrite the message method in order to replace it with their own message mechanism.
- Parameters:
- x_sourceTensor, shape = (…, n_source_cells, in_channels)
Input features on source cells. Assumes that all source cells have the same rank r.
- x_targetTensor, shape = (…, n_target_cells, in_channels)
Input features on target cells. Assumes that all target cells have the same rank s. Optional. If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.
- Returns:
- torch.Tensor, shape = (…, n_source_cells, in_channels)
Messages on source cells.