decouple mhsa residual

Author: mmz33

users/zeineldeen/models/asr/encoder/conformer_encoder_v2.py

@@ -903,12 +903,7 @@ def _create_mhsa_module(self, prefix_name, source, layer_index):
 
         drop = self.network.add_dropout_layer("{}_dropout".format(prefix_name), mhsa_linear, dropout=self.dropout)
 
-        res_inputs = [drop, source]
-
-        mhsa_res = self.network.add_combine_layer(
-            "{}_res".format(prefix_name), kind="add", source=res_inputs, n_out=self.enc_value_dim
-        )
-        return mhsa_res
+        return drop
 
     def _create_convolution_module(self, prefix_name, source, layer_index, half_step=False):
         """
@@ -1071,6 +1066,10 @@ def _create_conformer_block(self, i, source):
         if self.convolution_first:
             conv_module_ = self._create_convolution_module(prefix_name, ff_module1, i)
             mhsa_module = self._create_mhsa_module(prefix_name, conv_module_, i)
+            mhsa_module = self.network.add_combine_layer(
+                "{}_res".format(prefix_name), kind="add", source=[mhsa_module, conv_module_], n_out=self.enc_value_dim
+            )
+
             ff_module2_input = mhsa_module
         else:
             if self.no_mhsa_module:
@@ -1083,6 +1082,9 @@ def _create_conformer_block(self, i, source):
                     )
                     mhsa_input = conv_module1
                 mhsa = self._create_mhsa_module(prefix_name, mhsa_input, i)
+                mhsa = self.network.add_combine_layer(
+                    "{}_res".format(prefix_name), kind="add", source=[mhsa, mhsa_input], n_out=self.enc_value_dim
+                )
 
             conv_module = self._create_convolution_module(prefix_name, mhsa, i, half_step=self.sandwich_conv)
             ff_module2_input = conv_module

users/zeineldeen/models/asr/encoder/ebranchformer_encoder.py

@@ -14,6 +14,8 @@ def __init__(self, cgmlp_ff_dim, **kwargs):
         self.cgmlp_ff_dim = cgmlp_ff_dim
 
     def _create_conv_spatial_gating_unit(self, prefix_name, source, layer_index):
+        # Half split input into [A,B] -> A * DwConv(LN(B)) -> dropout
+        #
         # see also here: https://github.com/espnet/espnet/blob/master/espnet2/asr/layers/cgmlp.py#L15
 
         split_size = self.cgmlp_ff_dim // 2
@@ -47,6 +49,8 @@ def _create_conv_spatial_gating_unit(self, prefix_name, source, layer_index):
         return dropout
 
     def _create_conv_gating_mlp(self, prefix_name, source, layer_index):
+        # GeLU(FF(LN(x))) -> Half split input into [A,B] -> A * DwConv(LN(B)) -> dropout -> FF
+
         prefix_name = "{}_cgmlp".format(prefix_name)
 
         ln = self.network.add_layer_norm_layer("{}_ln".format(prefix_name), source)
@@ -65,6 +69,7 @@ def _create_conv_gating_mlp(self, prefix_name, source, layer_index):
 
         gelu_act = self.network.add_activation_layer("{}_gelu".format(prefix_name), ff1, activation="gelu")
 
+        # Half split input into [A,B] -> A * DwConv(LN(B)) -> dropout
         csgu = self._create_conv_spatial_gating_unit(f"{prefix_name}_csgu", gelu_act, layer_index)
 
         br_merge_ff = self.network.add_linear_layer(