torch_geometric.nn.aggr.EquilibriumAggregation

class EquilibriumAggregation(in_channels: int, out_channels: int, num_layers: List[int], grad_iter: int = 5, lamb: float = 0.1)[source]

Bases: Aggregation

The equilibrium aggregation layer from the “Equilibrium Aggregation: Encoding Sets via Optimization” paper.

The output of this layer \(\mathbf{y}\) is defined implicitly via a potential function \(F(\mathbf{x}, \mathbf{y})\), a regularization term \(R(\mathbf{y})\), and the condition

\[\mathbf{y} = \min_\mathbf{y} R(\mathbf{y}) + \sum_{i} F(\mathbf{x}_i, \mathbf{y}).\]

The given implementation uses a ResNet-like model for the potential function and a simple \(L_2\) norm \(R(\mathbf{y}) = \textrm{softplus}(\lambda) \cdot {\| \mathbf{y} \|}^2_2\) for the regularizer with learnable weight \(\lambda\).

Parameters:

  • in_channels (int) – Size of each input sample.

  • out_channels (int) – Size of each output sample.

  • num_layers (List[int) – List of hidden channels in the potential function.

  • grad_iter (int) – The number of steps to take in the internal gradient descent. (default: 5)

  • lamb (float) – The initial regularization constant. (default: 0.1)

reset_parameters()[source]

Resets all learnable parameters of the module.

forward(x: Tensor, index: Optional[Tensor] = None, ptr: Optional[Tensor] = None, dim_size: Optional[int] = None, dim: int = -2) → Tensor[source]

Forward pass.

Parameters:

  • x (torch.Tensor) – The source tensor.

  • index (torch.Tensor, optional) – The indices of elements for applying the aggregation. One of index or ptr must be defined. (default: None)

  • ptr (torch.Tensor, optional) – If given, computes the aggregation based on sorted inputs in CSR representation. One of index or ptr must be defined. (default: None)

  • dim_size (int, optional) – The size of the output tensor at dimension dim after aggregation. (default: None)

  • dim (int, optional) – The dimension in which to aggregate. (default: -2)

  • max_num_elements – (int, optional): The maximum number of elements within a single aggregation group. (default: None)