Support projection_lcl_residual_blocks and lcl_residual+mlp projection type in adversarial disentanglement module.

Author: arnor-sigurdsson

src/eir/setup/schema_modules/adversarial_schemas.py

@@ -49,6 +49,12 @@ class AdversarialConfig:
     :param projection_type:
         Type of projection layer to use before the adversarial discriminator.
 
+    :param projection_lcl_residual_blocks:
+        When ``projection_type='lcl+mlp_residual'``, use progressive LCL residual
+        blocks to reduce dimensionality before the final MLP residual block,
+        instead of a single LCL projection. Useful for very high-dimensional
+        embeddings where a dense MLP would blow up in parameters.
+
     :param embedding_cache_target:
         Whether to cache 'input' or 'output' of the embedding layer.
         Default 'output' caches the layer's output activations.
@@ -71,8 +77,14 @@ class AdversarialConfig:
     dropout_p: float = 0.1
     stochastic_depth_p: float = 0.0
     projection_type: Literal[
-        "linear", "lcl", "lcl_residual", "mlp_residual", "grouped_linear"
+        "linear",
+        "lcl",
+        "lcl_residual",
+        "mlp_residual",
+        "lcl+mlp_residual",
+        "grouped_linear",
     ] = "linear"
+    projection_lcl_residual_blocks: bool = True
     embedding_cache_target: Literal["input", "output"] = "output"
     target_cache_target: Literal["input", "output"] = "output"
 

src/eir/train_utils/adversarial.py

@@ -43,8 +43,14 @@ def __init__(
         dropout_p: float = 0.1,
         stochastic_depth_p: float = 0.0,
         projection_type: Literal[
-            "linear", "lcl", "lcl_residual", "mlp_residual", "grouped_linear"
+            "linear",
+            "lcl",
+            "lcl_residual",
+            "mlp_residual",
+            "lcl+mlp_residual",
+            "grouped_linear",
         ] = "linear",
+        projection_lcl_residual_blocks: bool = True,
     ):
         super().__init__()
         self.embedding_dim = embedding_dim
@@ -53,6 +59,7 @@ def __init__(
         self.layers = layers if layers is not None else [2]
         self.dropout_p = dropout_p
         self.projection_type = projection_type
+        self.projection_lcl_residual_blocks = projection_lcl_residual_blocks
 
         from_shape = torch.Size([embedding_dim])
         to_shape = torch.Size([fc_dim])
@@ -62,6 +69,7 @@ def __init__(
             to_shape_no_batch=to_shape,
             cache_fusion_type="sum",
             projection_type=projection_type,
+            projection_lcl_residual_blocks=projection_lcl_residual_blocks,
         )
 
         mlp_blocks = []
@@ -143,6 +151,9 @@ def hook_add_adversarial_losses(
                 dropout_p=adv_config.dropout_p,
                 stochastic_depth_p=adv_config.stochastic_depth_p,
                 projection_type=adv_config.projection_type,
+                projection_lcl_residual_blocks=(
+                    adv_config.projection_lcl_residual_blocks
+                ),
             )
             module = module.to(device=device)
             adversarial_modules[adv_config.name] = module