feat: force final ln in encoder

src/dalle_mini/model/configuration.py

@@ -60,12 +60,14 @@ class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
         # transformer variants
         ln_type="layernorm",  # layer normalization type, "rmsnorm", "layernorm"
         ln_positions="normformer",  # layer normalization positions, "normformer", "swinv2", "cogview", "postln", "deepnet" (same as postln)
-        head_scale=True,  # used in NormFormer
+        head_scale=False,  # used in NormFormer
         use_cosine_attention=False,  # used in Swin v2
         tau_init=0.05,  # used only in cosine attention (Swin v2)
         use_deepnet_scaling=False,  # used in Deepnet
-        use_glu=True,  # "GLU Variants Improve Transformer"
-        use_all_scale=True,  # use scale in layernorm even when seemingly unnecessary
+        use_glu=False,  # "GLU Variants Improve Transformer"
+        # parameters that should not be necessary but could affect results
+        force_ln_scale=True,  # force scale in layernorm even when followed by dense layers
+        force_final_ln_encoder=False,  # force layer normalization in encoder final layer even when followed by dense layers
         **kwargs,
     ):
         # text normalizer
@@ -91,7 +93,8 @@ class DalleBartConfig(PretrainedFromWandbMixin, PretrainedConfig):
         self.tau_init = tau_init
         self.use_deepnet_scaling = use_deepnet_scaling
         self.use_glu = use_glu
-        self.use_all_scale = use_all_scale
+        self.force_ln_scale = force_ln_scale
+        self.force_final_ln_encoder = force_final_ln_encoder
 
         # common parameters
         self.encoder_vocab_size = encoder_vocab_size

src/dalle_mini/model/modeling.py

@@ -378,7 +378,7 @@ class GLU(nn.Module):
                 self.config.ln_type,
                 dtype=self.dtype,
                 epsilon=1e-05,
-                use_scale=self.config.use_all_scale,
+                use_scale=self.config.force_ln_scale,
             )(x)
         w = nn.Dense(
             self.ffn_dim,
@@ -403,7 +403,7 @@ class GLU(nn.Module):
                 self.config.ln_type,
                 dtype=self.dtype,
                 epsilon=1e-05,
-                use_scale=self.config.use_all_scale,
+                use_scale=self.config.force_ln_scale,
             )(x)
         x = nn.Dropout(rate=self.config.activation_dropout)(
             x, deterministic=deterministic
@@ -443,7 +443,7 @@ class FFN(nn.Module):
                 self.config.ln_type,
                 dtype=self.dtype,
                 epsilon=1e-05,
-                use_scale=self.config.use_all_scale,
+                use_scale=self.config.force_ln_scale,
             )(x)
         x = nn.Dense(
             self.ffn_dim,
@@ -459,7 +459,7 @@ class FFN(nn.Module):
                 self.config.ln_type,
                 dtype=self.dtype,
                 epsilon=1e-05,
-                use_scale=self.config.use_all_scale,
+                use_scale=self.config.force_ln_scale,
             )(x)
         x = nn.Dropout(rate=self.config.activation_dropout)(
             x, deterministic=deterministic
@@ -512,7 +512,7 @@ class FlaxBartEncoderLayer(nn.Module):
                 self.config.ln_type,
                 dtype=self.dtype,
                 epsilon=1e-05,
-                use_scale=self.config.use_all_scale,
+                use_scale=self.config.force_ln_scale,
             )(hidden_states)
         hidden_states, attn_weights = FlaxBartAttention(
             config=self.config,
@@ -561,7 +561,7 @@ class FlaxBartEncoderLayer(nn.Module):
             use_scale = (
                 self.use_scale
                 or self.config.ln_positions == "postln"
-                or self.config.use_all_scale
+                or self.config.force_ln_scale
             )
             hidden_states = norm(
                 self.config.ln_type,
@@ -617,7 +617,7 @@ class FlaxBartDecoderLayer(nn.Module):
                 self.config.ln_type,
                 dtype=self.dtype,
                 epsilon=1e-05,
-                use_scale=self.config.use_all_scale,
+                use_scale=self.config.force_ln_scale,
             )(hidden_states)
         hidden_states, attn_weights = FlaxBartAttention(
             config=self.config,
@@ -656,7 +656,7 @@ class FlaxBartDecoderLayer(nn.Module):
                     self.config.ln_type,
                     dtype=self.dtype,
                     epsilon=1e-05,
-                    use_scale=self.config.use_all_scale,
+                    use_scale=self.config.force_ln_scale,
                 )(hidden_states)
             hidden_states, cross_attn_weights = FlaxBartAttention(
                 config=self.config,
@@ -709,7 +709,7 @@ class FlaxBartDecoderLayer(nn.Module):
             use_scale = (
                 self.use_scale
                 or self.config.ln_positions == "postln"
-                or self.config.use_all_scale
+                or self.config.force_ln_scale
             )
             hidden_states = norm(
                 self.config.ln_type,
@@ -761,8 +761,9 @@ class FlaxBartEncoderLayerCollection(nn.Module):
             # or every 6 layers for Swin v2
             # not needed for other models which use layernorm before x-attention
             # ignored args for deepnet which always add a norm with scale
-            add_norm = self.config.ln_positions == "swinv2" and (
-                (i == n_layers - 1) or ((i + 1) % 6 == 0)
+            add_norm = self.config.force_final_ln_encoder or (
+                self.config.ln_positions == "swinv2"
+                and ((i == n_layers - 1) or ((i + 1) % 6 == 0))
             )
             # we don't need to scale the norm for the last layer
             use_scale = i != n_layers - 1