Transformers with Flax - T5, BERT, RoBERTa, MT3 and more

[8bitmp3]

Just wanted to highlight some work made by others, in case someone has missed it:

T5 (Text-To-Text Transfer Transformer)

As demoed by the JAX team (@skye @jekbradbury) at NeurIPS 2020 this week, T5 (Text-To-Text Transfer Transformer) JAX code is available: https://github.com/google-research/google-research/tree/master/flax_models/t5x - built with Flax.

HuggingFace Transformers

And, in case someone has missed the November 2020 announcement by @avital:

• HuggingFace 🤗 and Google Brain 🧠 (by @marcvanzee) are collaborating to integrate HuggingFace's Transformers library in JAX with Flax (with Linen API):
  https://twitter.com/avitaloliver/status/1326986383983063058
https://twitter.com/MorganFunto/status/1326988188469039104

[avital]

Thanks for noticing and sharing @8bitmp3!

[8bitmp3]

T5x (T5 with JAX) code: https://github.com/google-research/google-research/tree/master/flax_models/t5x

➡️ https://github.com/google-research/t5x "new and improved"

thanks @adarob @hwchung27 @dhgarrette @levskaya @cpgaffney1 @0x0539 @hthu @afrozenator

[8bitmp3]

• JAX on TPUs at NeurIPS 2020 - slides + video: https://twitter.com/jekbradbury/status/1337528357517291520 (by @skye @jekbradbury)
• T5x (T5 with JAX) code: https://github.com/google-research/google-research/tree/master/flax_models/t5x

[8bitmp3]

Just wanted to highlight some work made by others, in case someone has missed it:

More recent work from Google Magenta:

MT3: Multi-Task Multitrack Music Transcription

by @jpgard, @iansimon, @ethman, @cghawthorne, @jesseengel

arXiv: https://arxiv.org/abs/2111.03017
GitHub: https://github.com/magenta/mt3

"Automatic Music Transcription (AMT), inferring musical notes from raw audio, is a challenging task at the core of music understanding. Unlike Automatic Speech Recognition (ASR), which typically focuses on the words of a single speaker, AMT often requires transcribing multiple instruments simultaneously, all while preserving fine-scale pitch and timing information. Further, many AMT datasets are “low-resource”, as even expert musicians find music transcription difficult and time-consuming. Thus, prior work has focused on task-specific architectures, tailored to the individual instruments of each task. In this work, motivated by the promising results of sequence-to-sequence transfer learning for low-resource Natural Language Processing (NLP), we demonstrate that a general-purpose Transformer model can perform multi-task AMT, jointly transcribing arbitrary combinations of musical instruments across several transcription datasets. We show this unified training framework achieves high-quality transcription results across a range of datasets, dramatically improving performance for low-resource instruments (such as guitar), while preserving strong performance for abundant instruments (such as piano). Finally, by expanding the scope of AMT, we expose the need for more consistent evaluation metrics and better dataset alignment, and provide a strong baseline for this new direction of multi-task AMT."

"We use the T5 “small” model architecture described in Raffel et al. (2019). This is a standard Transformer architecture, and we use the implementation available in T5X5, which is built on FLAX (Heek et al., 2020) and JAX (Bradbury et al., 2020)."

"A key contribution of our work is the use of a single generic architecture — Transformers via T5 (Raffel et al., 2019) — to address a variety of tasks previously tackled using complex, handcrafted, dataset-specific architectures. The T5 architecture is an encoder-decoder Transformer model which closely follows the original form in Vaswani et al. (2017).
"In the T5 architecture, a sequence of inputs is mapped to a sequence of learned embeddings plus fixed positional embeddings; we use absolute positional embeddings instead of the bucketed “relative” embeddings used in Raffel et al. (2019) to ensure that all positions can be attended to with equal resolution. The model uses a series of standard Transformer self-attention “blocks” in both the encoder and decoder. In order to produce a sequence of output tokens, the model uses greedy autoregressive decoding: an input sequence is fed in, the output token with the highest predicted probability of occurring next is appended to the sequence, and the process is repeated until an end-of-sequence (EOS) token is produced. We use the T5 “small” model, which contains approximately 60 million parameters. While much larger models are commonly used in language modeling tasks, we found that increasing model size tended to exacerbate overfitting."

"In this work we have shown that posing multi-instrument music transcription as a sequence-to-sequence task and training a generic Transformer architecture simultaneously on a variety of datasets advances the state of the art in multi-instrument transcription, most notably in the low-resource scenario. We also introduced and applied a consistent evaluation method using note onset+offset+instrument F1 scores, using a standard instrument taxonomy."

---

Sequence-to-Sequence Piano Transcription with Transformers

by @cghawthorne, @iansimon, @rigeljs, @ethman, @jesseengel

arXiv: https://arxiv.org/abs/2107.09142

"Automatic Music Transcription has seen significant progress in recent years by training custom deep neural networks on large datasets. However, these models have required extensive domain-specific design of network architectures, input/output representations, and complex decoding schemes. In this work, we show that equivalent performance can be achieved using a generic encoder-decoder Transformer with standard decoding methods. We demonstrate that the model can learn to translate spectrogram inputs directly to MIDI-like output events for several transcription tasks. This sequence-to-sequence approach simplifies transcription by jointly modeling audio features and language-like output dependencies, thus removing the need for task-specific architectures. These results point toward possibilities for creating new Music Information Retrieval models by focusing on dataset creation and labeling rather than custom model design."

"The model is implemented using the T5X framework, which is built on Flax and JAX. We also use SeqIO for data preprocessing and evaluation."

"While some recent research using Transformers has favored very large models, such as GPT-3 with 175B parameters, we found that a comparatively small model is sufficient for these tasks. With the configuration described above, our model has only 54M parameters, only roughly twice that of Onsets and Frames, which has 28M parameters."

"We have shown that a generic Transformer architecture trained to map spectrograms to MIDI-like output events with no pretraining can achieve state-of-the-art performance on automatic piano transcription. We see this as an appeal to simplicity; we used standard formats and architectures as much as possible and were able to achieve results on par with models customized for piano transcription. Possibly the main source of complexity in our setup is the splitting of examples into segments; future work could include the investigation of sparse attention mechanisms to enable the transcription of an entire piece of music in a single encoding and decoding pass. Also worth exploring
would be the use of distillation or related techniques to enable models like this to run in realtime on mobile devices or the web.
"Our results suggest that a generic sequence-to-sequence framework with Transformers might also be beneficial for other MIR tasks, such as beat tracking, fundamental frequency estimation, chord estimation, etc. The field of Natural Language Processing has seen that a single large language model, such as GPT-3 or T5, has been capable of solving multiple tasks by leveraging the commonalities between tasks..."

[8bitmp3]

https://github.com/google/flaxformer

Flaxformer: transformer architectures in JAX/Flax

Flaxformer is a transformer library for primarily NLP and multimodal research at Google. It is used for many NLP research use cases, providing both off-the-shelf BERT and T5 models, and several research projects built on shared components.

[8bitmp3]

Nested Hierarchical Transformer

NesT is a simple method, which aggregates nested local transformers on image blocks. The idea makes vision transformers attain better accuracy, data efficiency, and convergence on the ImageNet benchmark. NesT can be scaled to small datasets to match convnet accuracy.

Blog post: https://ai.googleblog.com/2022/02/nested-hierarchical-transformer-towards.html
GitHub: https://github.com/google-research/nested-transformer
arXiv: https://arxiv.org/pdf/2105.12723.pdf

[8bitmp3]

Multiview Transformers for Video Recognition (2022)

Video understanding requires reasoning at multiple spatiotemporal resolutions – from short fine-grained motions to events taking place over longer durations. Although transformer architectures have recently advanced the stateof-the-art, they have not explicitly modelled different spatiotemporal resolutions. To this end, we present Multiview Transformers for Video Recognition (MTV). Our model consists of separate encoders to represent different views of the input video with lateral connections to fuse information across views. We present thorough ablation studies of our model and show that MTV consistently performs better than single-view counterparts in terms of accuracy and computational cost across a range of model sizes. Furthermore, we achieve state-of-the-art results on six standard datasets, and improve even further with largescale pretraining.

GitHub: https://github.com/google-research/scenic/tree/main/scenic/projects/mtv
arXiv: https://arxiv.org/abs/2201.04288

[8bitmp3]

FitVid: Overfitting in Pixel-Level Video Prediction

...we introduce a new architecture, named FitVid, which is capable of severe overfitting on the common benchmarks while having similar parameter count as the current state-of-the-art models. We analyze the consequences of overfitting, illustrating how it can produce unexpected outcomes such as generating high quality output by repeating the training data, and how it can be mitigated using existing image augmentation techniques. As a result, FitVid outperforms the current state-of-the-art models across four different video prediction benchmarks on four different metrics.

Site: https://sites.google.com/corp/view/fitvidpaper
GitHub: https://github.com/google-research/fitvid
arXiv: https://arxiv.org/abs/2106.13195

[8bitmp3]

Visual Prompt Tuning for Generative Transfer Learning

We present a method for learning image generation models from diverse data distributions and varying amount of training data via knowledge transfer from the source model trained on a large dataset. A simple modification on prompt token designs allows to learn a parameter and compute efficient class and instance conditional image generation models of autoregressive and non-autoregressive vision transformers. Comprehensive experimental results of image synthesis are provided across diverse visual domains, tasks, and the number of training images. In addition, we show how to use learned prompts for novel image synthesis in the form of marquee header prompts, which is particularly useful when synthesizing images using generative models learned from a few images.

arXiv: https://arxiv.org/pdf/2210.00990.pdf
GitHub: https://github.com/google-research/generative_transfer

[8bitmp3]

Scaling Vision with Sparse Mixture of Experts

Sparsely-gated Mixture of Experts networks (MoEs) have demonstrated excellent scalability in Natural Language Processing. In Computer Vision, however, almost all performant networks are “dense”, that is, every input is processed by every parameter. We present a Vision MoE (V-MoE), a sparse version of the Vision Transformer, that is scalable and competitive with the largest dense networks. When applied to image recognition, V-MoE matches the performance of state-ofthe-art networks, while requiring as little as half of the compute at inference time. Further, we propose an extension to the routing algorithm that can prioritize subsets of each input across the entire batch, leading to adaptive per-image compute. This allows V-MoE to trade-off performance and compute smoothly at test-time. Finally, we demonstrate the potential of V-MoE to scale vision models, and train a 15B parameter model that attains 90.35% on ImageNet.

arXiv: https://arxiv.org/pdf/2106.05974.pdf
GitHub: https://github.com/google-research/vmoe (by @jpuigcerver and many others)