Merge pull request #160 from arnor-sigurdsson/residual-flow-improvements

Residual flow improvements

Author: arnor-sigurdsson

src/eir/models/fusion/fusion_default.py

@@ -100,7 +100,7 @@ def __init__(
             self.input_projections[name] = nn.Sequential(
                 nn.RMSNorm(normalized_shape=output_dim),
                 nn.Linear(in_features=output_dim, out_features=self.fusion_dim),
-                nn.GELU(),
+                nn.RMSNorm(normalized_shape=self.fusion_dim),
             )
 
         fusion_resblocks_kwargs = {

src/eir/models/input/array/models_locally_connected.py

@@ -241,7 +241,7 @@ def __init__(
             kernel_size=fc_0_kernel_size,
             bias=True,
         )
-        self.act_0 = nn.GELU()
+        self.norm_0 = nn.RMSNorm(normalized_shape=self.fc_0.out_features)
 
         cutoff = dynamic_cutoff or self.model_config.cutoff
         assert isinstance(cutoff, int)
@@ -287,7 +287,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         out = self.flatten_fn(x=input)
 
         out = self.fc_0(out)
-        out = self.act_0(out)
+        out = self.norm_0(out)
         out = self.lcl_blocks(out)
 
         return out
@@ -302,12 +302,10 @@ class ExpertBranch(nn.Module):
     def __init__(
         self,
         fc_0: LCL,
-        act_0: nn.Module,
         lcl_blocks: nn.Sequential,
     ):
         super().__init__()
         self.fc_0 = fc_0
-        self.act_0 = act_0
         self.lcl_blocks = lcl_blocks
 
     @property
@@ -318,7 +316,6 @@ def out_features(self) -> int:
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.fc_0(x)
-        x = self.act_0(x)
         x = self.lcl_blocks(x)
         return x
 
@@ -381,7 +378,6 @@ def __init__(
                 kernel_size=expert_fc_0_kernel,
                 bias=True,
             )
-            act_0 = nn.GELU()
 
             expert_kernel_width = _clamp_kernel_for_min_chunks(
                 kernel_size=kernel_width,
@@ -407,7 +403,6 @@ def __init__(
 
             branch = ExpertBranch(
                 fc_0=fc_0,
-                act_0=act_0,
                 lcl_blocks=lcl_blocks,
             )
             self.expert_branches[name] = branch

src/eir/models/input/tabular/tabular.py

@@ -185,9 +185,12 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         output = self.layer(input)
         output = self.mlp_blocks(output)
 
-        if self.training and self.drop_prob > 0.0:
-            mask = (torch.rand(1, device=output.device) >= self.drop_prob).float()
-            output = output * mask
+        if self.drop_prob > 0.0:
+            if self.training:
+                mask = (torch.rand(1, device=output.device) >= self.drop_prob).float()
+                output = output * mask
+            else:
+                output = output * (1.0 - self.drop_prob)
 
         return output
 

src/eir/models/layers/cnn_layers.py

@@ -194,6 +194,12 @@ def __init__(
         stochastic_depth_p: float = 0.0,
         conv_downsample_identity: bool = True,
     ):
+        """
+        TODO:   Look into adding a dynamic addition of a norm under downsample identity
+                (if dimensions are being changed in the downsample),
+                similar to what we have e.g. in MLPResidualBlock currently.
+
+        """
         super().__init__()
 
         self.in_channels = in_channels

src/eir/models/layers/lcl_layers.py

@@ -164,26 +164,26 @@ def __init__(
 
         self.out_features = self.fc_2.out_features
 
-        # generally seems to work better to always initialize to 1.0 in LCL blocks
-        # in contrast to what we do in the standard MLP blocks
-        self.ls = LayerScale(dim=self.out_features, init_values=1.0)
+        self.ls = LayerScale(dim=self.out_features, init_values=1e-05)
 
-        if in_features == self.out_features:
-            self.downsample_identity = lambda x: x
-        else:
+        self._norm_identity = full_preactivation or (in_features != self.out_features)
+        self.downsample_identity: nn.Module
+        if in_features != self.out_features:
             self.downsample_identity = LCL(
                 in_features=self.in_features,
                 out_feature_sets=1,
                 bias=True,
                 num_chunks=self.fc_2.out_features,
             )
+        else:
+            self.downsample_identity = nn.Identity()
 
         self.stochastic_depth = StochasticDepth(p=stochastic_depth_p, mode="batch")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         out = self.norm_1(x)
 
-        identity = out if self.full_preactivation else x
+        identity = out if self._norm_identity else x
         identity = self.downsample_identity(identity)
 
         out = self.fc_1(out)

src/eir/models/layers/mlp_layers.py

@@ -72,17 +72,18 @@ def __init__(
             bias=False,
         )
 
+        self._norm_identity = full_preactivation or (in_features != out_features)
         self.downsample_identity: nn.Module
-        if in_features == out_features:
-            self.downsample_identity = nn.Identity()
-            ls_init = 1e-05
-        else:
+        if in_features != out_features:
             self.downsample_identity = nn.Linear(
                 in_features=in_features,
                 out_features=out_features,
                 bias=True,
             )
             ls_init = 1.0
+        else:
+            self.downsample_identity = nn.Identity()
+            ls_init = 1e-05
 
         self.ls = LayerScale(
             dim=out_features,
@@ -97,7 +98,7 @@ def __init__(
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         out = self.norm_1(x)
 
-        identity = out if self.full_preactivation else x
+        identity = out if self._norm_identity else x
         identity = self.downsample_identity(identity)
 
         out = self.fc_1(out)

src/eir/models/layers/projection_layers.py

@@ -113,6 +113,8 @@ def get_lcl_projection_layer(
     out_feature_sets_candidates: Sequence[int] = tuple(range(1, 1024 + 1)),
     diff_tolerance: int = 0,
     kernel_width_divisible_by: int | None = None,
+    full_preactivation: bool = False,
+    dropout_p: float = 0.0,
 ) -> LCLResidualBlock | LCL | None:
     layer_class: type[LCLResidualBlock] | type[LCL]
     match layer_type:
@@ -140,11 +142,17 @@ def get_lcl_projection_layer(
         return None
 
     best_kernel_size, best_out_feature_sets = solution
-    best_layer = layer_class(
-        in_features=input_dimension,
-        kernel_size=best_kernel_size,
-        out_feature_sets=best_out_feature_sets,
-    )
+
+    kwargs: dict = {
+        "in_features": input_dimension,
+        "kernel_size": best_kernel_size,
+        "out_feature_sets": best_out_feature_sets,
+    }
+    if layer_type == "lcl_residual":
+        kwargs["full_preactivation"] = full_preactivation
+        kwargs["dropout_p"] = dropout_p
+
+    best_layer = layer_class(**kwargs)
 
     return best_layer
 

src/eir/models/output/tabular/shared_mlp_residual.py

@@ -98,13 +98,18 @@ def _build_shared(self, input_dimension: int) -> None:
 
         final_block = MLPResidualBlock(
             in_features=self.model_config.fc_task_dim,
-            out_features=self.total_outputs,
+            out_features=self.model_config.fc_task_dim,
             dropout_p=self.model_config.rb_do,
             stochastic_depth_p=self.model_config.stochastic_depth_p,
             full_preactivation=False,
         )
 
-        self.shared_branch = nn.Sequential(shared_branch_module, final_block)
+        final_norm = nn.RMSNorm(self.model_config.fc_task_dim)
+        final_proj = nn.Linear(self.model_config.fc_task_dim, self.total_outputs)
+
+        self.shared_branch = nn.Sequential(
+            shared_branch_module, final_block, final_norm, final_proj
+        )
 
     def _build_expert(self, input_dimension: int, num_experts: int) -> None:
         self.input_identity = nn.Identity()
@@ -140,6 +145,8 @@ def _build_expert(self, input_dimension: int, num_experts: int) -> None:
             torch.zeros(num_targets, num_experts),
         )
 
+        self.expert_norm = nn.RMSNorm(expert_dim)
+
         self.target_final_layers = nn.ModuleDict(
             {
                 name: nn.Linear(in_features=expert_dim, out_features=size)
@@ -176,6 +183,7 @@ def _forward_expert(self, inputs: torch.Tensor) -> dict[str, torch.Tensor]:
         # Per-target weighted average of expert outputs:
         # (T, E) @ (B, E, D) -> (B, T, D)
         all_mixed = torch.matmul(gate_weights, stacked)
+        all_mixed = self.expert_norm(all_mixed)
 
         per_target = []
         for i, name in enumerate(self.target_names):

src/eir/models/tensor_broker/tensor_broker_fusion_layers.py

@@ -89,7 +89,7 @@ def __init__(
         self,
         input_shape: torch.Size,
         context_shape: torch.Size,
-        gate_init_value: float = 0.01,
+        gate_init_value: float = 1e-05,
     ):
         super().__init__()
         self.input_shape = input_shape

src/eir/models/tensor_broker/tensor_broker_projection_layers.py

@@ -86,11 +86,6 @@ def get_projection_layer(
 
             mlp_input_target = target_dim
 
-            norm_layer = nn.RMSNorm(normalized_shape=input_dim)
-            act_layer = nn.GELU()
-            projection_layers.append(norm_layer)
-            projection_layers.append(act_layer)
-
             if projection_lcl_residual_blocks:
                 cur_dim = input_dim
                 while cur_dim // 4 > mlp_input_target:
@@ -101,6 +96,7 @@ def get_projection_layer(
                         layer_type="lcl_residual",
                         diff_tolerance=cur_dim // 100,
                         kernel_width_divisible_by=kernel_width_divisible_by,
+                        dropout_p=0.1,
                     )
                     if block is not None:
                         projection_layers.append(block)
@@ -116,6 +112,8 @@ def get_projection_layer(
                         projection_layers.append(fallback)
                         cur_dim = halve_target
             else:
+                projection_layers.append(nn.RMSNorm(normalized_shape=input_dim))
+
                 try:
                     lcl_projection_layer = get_1d_projection_layer(
                         input_dimension=input_dim,
@@ -125,9 +123,6 @@ def get_projection_layer(
                         kernel_width_divisible_by=kernel_width_divisible_by,
                     )
                 except ValueError:
-                    # Sometimes we cannot create are reasonable LCL projection
-                    # e.g. if target dim is much larger tha input dim so we have this
-                    # fallback
                     lcl_projection_layer = get_1d_projection_layer(
                         input_dimension=input_dim,
                         target_dimension=mlp_input_target,
@@ -137,16 +132,17 @@ def get_projection_layer(
                     )
 
                 projection_layers.append(lcl_projection_layer)
+                projection_layers.append(nn.RMSNorm(normalized_shape=mlp_input_target))
                 cur_dim = mlp_input_target
 
             mlp_residual_block = MLPResidualBlock(
                 in_features=cur_dim,
                 out_features=target_dim,
                 dropout_p=0.0,
-                full_preactivation=True,
                 stochastic_depth_p=0.0,
             )
             projection_layers.append(mlp_residual_block)
+
             projected_shape = to_shape_no_batch
 
         case "mlp_residual":
@@ -155,13 +151,18 @@ def get_projection_layer(
 
             projection_layer = MLPResidualBlock(
                 in_features=input_dim,
-                out_features=target_dim,
+                out_features=input_dim,
                 dropout_p=0.0,
-                full_preactivation=True,
                 stochastic_depth_p=0.0,
                 reduce_at_fc_1=False,
             )
             projection_layers.append(projection_layer)
+
+            projection_layers.append(nn.RMSNorm(normalized_shape=input_dim))
+            projection_layers.append(
+                nn.Linear(in_features=input_dim, out_features=target_dim)
+            )
+
             projected_shape = to_shape_no_batch
 
         case "cnn":