Add projectors to DeepSet

bayesflow/networks/deep_set/deep_set.py

@@ -30,7 +30,7 @@ def __init__(
         mlp_widths_invariant_inner: Sequence[int] = (64, 64),
         mlp_widths_invariant_outer: Sequence[int] = (64, 64),
         mlp_widths_invariant_last: Sequence[int] = (64, 64),
-        activation: str = "gelu",
+        activation: str = "silu",
         kernel_initializer: str = "he_normal",
         dropout: int | float | None = 0.05,
         spectral_normalization: bool = False,
@@ -72,7 +72,7 @@ def __init__(
         mlp_widths_invariant_last : Sequence[int], optional
             Widths of the MLP layers in the final invariant transformation. Default is (64, 64).
         activation : str, optional
-            Activation function used throughout the network, such as "gelu". Default is "gelu".
+            Activation function used throughout the network, such as "gelu". Default is "silu".
         kernel_initializer : str, optional
             Initialization strategy for kernel weights, such as "he_normal". Default is "he_normal".
         dropout : int, float, or None, optional

bayesflow/networks/deep_set/equivariant_layer.py

@@ -94,6 +94,7 @@ def __init__(
             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
 
@@ -137,7 +138,10 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         output_set = ops.concatenate([input_set, invariant_summary], axis=-1)
 
         # Pass through final equivariant transform + residual
-        output_set = input_set + self.equivariant_fc(output_set, training=training)
+        out_fc = self.equivariant_fc(output_set, training=training)
+        out_projected = self.out_fc_projector(out_fc)
+        output_set = input_set + out_projected
+
         if self.layer_norm is not None:
             output_set = self.layer_norm(output_set, training=training)
 

bayesflow/networks/deep_set/invariant_layer.py

@@ -74,6 +74,7 @@ def __init__(
             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,
@@ -82,6 +83,7 @@ def __init__(
             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:
@@ -106,8 +108,10 @@ def call(self, input_set: Tensor, training: bool = False, **kwargs) -> Tensor:
         """
 
         set_summary = self.inner_fc(input_set, training=training)
+        set_summary = self.inner_projector(set_summary)
         set_summary = self.pooling_layer(set_summary, training=training)
         set_summary = self.outer_fc(set_summary, training=training)
+        set_summary = self.outer_projector(set_summary)
         return set_summary
 
     @sanitize_input_shape