Google open sources Switch Transformer models in T5X/JAX

Google open sources Switch Transformer models in T5X/JAX
Listen to this story

Google Brain has open sourced Switch Transformer models including 1.6T param Switch-C and the 395B param Switch-XXL in T5X/JAX.

Check out the GitHub repository here

What is JAX?


Sign up for your weekly dose of what's up in emerging technology.

JAX (Just After eXecution) is a machine/deep learning library developed by DeepMind. All JAX operations are based on XLA or Accelerated Linear Algebra. XLA, developed by Google, is a domain-specific compiler for linear algebra that uses whole-program optimisations to accelerate computing. XLA makes BERT’s training speed faster by almost 7.3 times.


Download our Mobile App

T5X is a modular, composable, research-friendly framework for high-performance, configurable, and inference of sequence models (starting with language) at many scales.

T5X allows many different classes of language tasks to be expressed uniformly, and a single encoder-decoder architecture can handle them without any task-specific parameters. 

Switch Transformer 

In deep learning, models usually reuse the same parameters for all inputs. However, the Switch Transformer uses a Mixture of Experts (MoE) and selects different parameters for each incoming example.

The model is primarily used for NLP research. The Transformer uses an algorithm called Switch Routing. Instead of activating several experts and combining the output, it chooses a single expert to work on an input. The algorithm simplifies the routing computation and reduces communication costs since individual expert models are hosted on different GPU devices.

Support independent technology journalism

Get exclusive, premium content, ads-free experience & more

Rs. 299/month

Subscribe now for a 7-day free trial

More Great AIM Stories

Tasmia Ansari
Tasmia is a tech journalist at AIM, looking to bring a fresh perspective to emerging technologies and trends in data science, analytics, and artificial intelligence.

AIM Upcoming Events

Early Bird Passes expire on 3rd Feb

Conference, in-person (Bangalore)
Rising 2023 | Women in Tech Conference
16-17th Mar, 2023

Conference, in-person (Bangalore)
Data Engineering Summit (DES) 2023
27-28th Apr, 2023

3 Ways to Join our Community

Telegram group

Discover special offers, top stories, upcoming events, and more.

Discord Server

Stay Connected with a larger ecosystem of data science and ML Professionals

Subscribe to our Daily newsletter

Get our daily awesome stories & videos in your inbox