🖼️ When ViT meets GPT-2 📝

• The source code of Flax's version of GPT-2 is modified to be able to accept an encoder's outputs, so it can be used as a decoder in an encoder-decoder architecture.
• Originally, we created FlaxViTGPT2ForConditionalGenerationModule, which is FlaxViTGPT2Module (ViT + [GPT-2 without LM head]) with an extra LM head. However, when loading the pretrained French GPT-2 model, the LM head's weigths are not loaded. We therefore created FlaxViTGPT2LMForConditionalGenerationModule which is ViT + [GPT-2 with LM head], and we no longer need to add a LM head over it. By doing so, the pretrained LM head's weights are also loaded, and the only randomly initialized weigths are the cross-attention weights.
• The provided training script run_summarization.py is modified to send pixel values to the model instead of a sequence of input token ids, and a necessary change due to the ViT model not accepting an attention_mask argument.
• We first tried to use WIT : Wikipedia-based Image Text Dataset, but found it is a very changeling task since, unlike traditional image captioning tasks, it requires the model to be able to generate different texts even if two images are similar (for example, two famous dogs might have completely different Wikipedia texts).
• We finally decided to use COCO image dataset at the final day of this Flax community event. We were able to translate only about 65000 examples to French for training, and the model is trained for only 5 epochs (beyond this, it started to overfit). This leads to the poor performance.