DeepSet: Adapt output dimension of invariant module inside the equivariant module (#557) (#561)

* adapt output dim of invariant module in equivariant module

See #557. The DeepSet showed bad performance and was not able to learn
diverse summary statistics. Reducing the dimension of the output of the
invariant module inside the equivariant module improves this, probably
because the invidividual information of each set member gains importance
compared to the shared information provided by the invariant module.

There might be better settings for this, so we might update the default
later on. However, this is already an improvement over the previous
setting.

* DeepSet: adapt docstring to reflect code

Author: vpratz

bayesflow/networks/deep_set/deep_set.py

@@ -28,7 +28,7 @@ def __init__(
         output_pooling: str = "mean",
         mlp_widths_equivariant: Sequence[int] = (64, 64),
         mlp_widths_invariant_inner: Sequence[int] = (64, 64),
-        mlp_widths_invariant_outer: Sequence[int] = (64, 64),
+        mlp_widths_invariant_outer: Sequence[int] = (64, 4),
         mlp_widths_invariant_last: Sequence[int] = (64, 64),
         activation: str = "silu",
         kernel_initializer: str = "he_normal",
@@ -68,7 +68,7 @@ def __init__(
         mlp_widths_invariant_inner : Sequence[int], optional
             Widths of the inner MLP layers within the invariant module. Default is (64, 64).
         mlp_widths_invariant_outer : Sequence[int], optional
-            Widths of the outer MLP layers within the invariant module. Default is (64, 64).
+            Widths of the outer MLP layers within the invariant module. Default is (64, 4).
         mlp_widths_invariant_last : Sequence[int], optional
             Widths of the MLP layers in the final invariant transformation. Default is (64, 64).
         activation : str, optional
@@ -80,7 +80,7 @@ def __init__(
         spectral_normalization : bool, optional
             Whether to apply spectral normalization to stabilize training. Default is False.
         **kwargs
-            Additional keyword arguments passed to the equivariant and invariant modules.
+            Additional keyword arguments passed to the base class.
         """
 
         super().__init__(**kwargs)