Initialize projectors for invariant and equivariant DeepSet layers

Author: han-ol

bayesflow/networks/deep_set/equivariant_layer.py

@@ -88,12 +88,13 @@ def __init__(
         # Fully connected net + residual connection for an equivariant transform applied to each set member
         self.input_projector = layers.Dense(mlp_widths_equivariant[-1])
         self.equivariant_fc = MLP(
-            mlp_widths_equivariant,
+            mlp_widths_equivariant[:-1],
             dropout=dropout,
             activation=activation,
             kernel_initializer=kernel_initializer,
             spectral_normalization=spectral_normalization,
         )
+        self.out_fc_projector = keras.layers.Dense(mlp_widths_equivariant[-1], kernel_initializer=kernel_initializer)
 
         self.layer_norm = layers.LayerNormalization() if layer_norm else None
 
@@ -140,7 +141,7 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         out_fc = self.equivariant_fc(output_set, training=training)
         out_projected = self.out_fc_projector(out_fc)
         output_set = input_set + out_projected
-        # output_set = input_set + self.equivariant_fc(output_set, training=training)
+
         if self.layer_norm is not None:
             output_set = self.layer_norm(output_set, training=training)
 

bayesflow/networks/deep_set/invariant_layer.py

@@ -68,20 +68,22 @@ def __init__(
 
         # Inner fully connected net for sum decomposition: inner( pooling( inner(set) ) )
         self.inner_fc = MLP(
-            mlp_widths_inner,
+            mlp_widths_inner[:-1],
             dropout=dropout,
             activation=activation,
             kernel_initializer=kernel_initializer,
             spectral_normalization=spectral_normalization,
         )
+        self.inner_projector = keras.layers.Dense(mlp_widths_inner[-1], kernel_initializer=kernel_initializer)
 
         self.outer_fc = MLP(
-            mlp_widths_outer,
+            mlp_widths_outer[:-1],
             dropout=dropout,
             activation=activation,
             kernel_initializer=kernel_initializer,
             spectral_normalization=spectral_normalization,
         )
+        self.outer_projector = keras.layers.Dense(mlp_widths_outer[-1], kernel_initializer=kernel_initializer)
 
         # Pooling function as keras layer for sum decomposition: inner( pooling( inner(set) ) )
         if pooling_kwargs is None: