Add project_encoder and related layers

vit_gpt2/configuration_vit_gpt2.py

@@ -16,7 +16,10 @@ class ViTGPT2Config(PretrainedConfig):
 
     def __init__(self, vision_config_dict=None, text_config_dict=None, **kwargs):
         super().__init__(
-            text_config_dict=text_config_dict, vision_config_dict=vision_config_dict, **kwargs)
+            vision_config_dict=vision_config_dict, text_config_dict=text_config_dict, **kwargs
+        )
+
+        project_encoder = kwargs.pop("project_encoder", None)
 
         if vision_config_dict is None:
             vision_config_dict = {}
@@ -41,6 +44,19 @@ class ViTGPT2Config(PretrainedConfig):
         self.decoder_start_token_id = self.text_config.bos_token_id
         self.forced_eos_token_id = self.text_config.eos_token_id
 
+        _project_encoder = getattr(self.text_config, "project_encoder", None)
+        if project_encoder is not None and _project_encoder is not None:
+            assert project_encoder == _project_encoder
+        elif project_encoder:
+            _project_encoder = project_encoder
+        elif _project_encoder:
+            project_encoder = _project_encoder
+        else:
+            project_encoder = False
+
+        self.config.project_encoder = project_encoder
+        self.text_config.project_encoder = project_encoder
+
     @classmethod
     def from_vision_text_configs(cls, vision_config: ViTConfig, text_config: GPT2Config, **kwargs):
 

vit_gpt2/modeling_flax_gpt2.py

@@ -308,9 +308,13 @@ class FlaxGPT2Block(nn.Module):
         self.attn = FlaxGPT2Attention(self.config, dtype=self.dtype)
 
         if not self.only_self_attn:
-            self.encoder_ln = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
+            self.cross_attn_ln = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
             # [IMPORTANT] Cross attention requires ``causal=False``! This is a bug I made previously.
-            self.encoder_attn = FlaxGPT2Attention(config=self.config, dtype=self.dtype, causal=False, self_attn=False)
+            self.cross_attn = FlaxGPT2Attention(config=self.config, dtype=self.dtype, causal=False, self_attn=False)
+
+            if self.config.project_encoder:
+                self.encoder_projection_ln = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
+                self.encoder_projection_mlp = FlaxGPT2MLP(self.config, self.config.hidden_size, dtype=self.dtype)
 
         self.ln_2 = nn.LayerNorm(epsilon=self.config.layer_norm_epsilon, dtype=self.dtype)
         self.mlp = FlaxGPT2MLP(self.config, inner_dim, dtype=self.dtype)
@@ -348,10 +352,19 @@ class FlaxGPT2Block(nn.Module):
         cross_attn_weights = None
         if encoder_hidden_states is not None:
 
+            if self.project_encoder:
+                residual = encoder_hidden_states
+                encoder_hidden_states = self.encoder_projection_ln(encoder_hidden_states)
+                feed_forward_hidden_states = self.encoder_projection_mlp(
+                    encoder_hidden_states, deterministic=deterministic
+                )
+                # residual connection
+                encoder_hidden_states = residual + feed_forward_hidden_states
+
             residual = hidden_states
-            hidden_states = self.encoder_ln(hidden_states)
+            hidden_states = self.cross_attn_ln(hidden_states)
 
-            cross_attn_outputs = self.encoder_attn(
+            cross_attn_outputs = self.cross_attn(
                 hidden_states=hidden_states,
                 key_value_states=encoder_hidden_states,
                 attention_mask=encoder_attention_mask,

vit_gpt2/modeling_flax_vit_gpt2_lm.py

@@ -541,6 +541,7 @@ class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
 
             if "config" not in text_kwargs:
                 text_config = GPT2Config.from_pretrained(text_pretrained_model_name_or_path)
+                text_config.project_encoder = text_kwargs.pop("project_encoder", None)
                 text_kwargs["config"] = text_config
 
             text_kwargs["config"].add_cross_attention = True