AI21 Labs: Rewriting the rules on natural language processing with Google Cloud
About AI21 Labs
AI21 Labs was founded in 2017 to reimagine the way humans read and write. The company uses neural networks fuelled by huge volumes of data and powerful computing to enable machines to understand written language and build tools that complement the skills of humans.
Tell us your challenge. We're here to help.
Contact usAI21 Labs is on a mission to enhance humanity's unique ability to create and understand written language and it uses Google Cloud to power the machine learning tools needed to achieve this.
Google Cloud results
- Scales by a factor of 100 the number A100 GPUs running within a day
- Provides flexible price-performance choice, between 7x higher cost and 2x better performance with GKE changing architecture
- Allows saving of up to 20% due to variation in system load with no performance tradeoffs with Elasticity
Builds language models with 178 billion data parameters
How we read and write is changing. Advancements in machine learning are quickly expanding the tools available to help us craft language. Founded in 2017, AI21 Labs is hoping to become the next big disruptor in this space, leveraging the power of Natural Language Processing (NLP) to give machines the ability to understand and generate natural text.
Yet, far from seeing machines as a replacement for our unique ability to create and understand the written symbols that represent our understanding of the world, AI21 Labs views NLP as a tool as revolutionary as Gutenberg’s printing press. It aims to empower people to become better versions of their writing and reading selves, using the power of Machine Learning (ML). It’s a lofty ambition, one which the AI21 Labs team felt Google Cloud was well placed to enable. "There’ve been huge transformations in NLP even since we entered this arena," explains Barak Peleg, VP of Engineering at AI21 Labs. "We’re a startup among giants, and we knew that we had to adopt the best technology available to achieve our ambitions. For us, Google Cloud and its wide range of advanced machine learning accelerators was an ideal fit."
Building proprietary NLP algorithms with Google Cloud
AI21 Labs uses a wide variety of ML accelerators, including high-performance GPUs to develop and serve its suite of language models, including state-of-the-art Jurassic-1 Jumbo, comprising 178 billion parameters of data. It also uses Cloud Storage to store linguistic data drawn from datasets built from online content.
On the operational side of the business, AI21 Labs uses Google Kubernetes Engine for its production environment, abstracting functionality to Docker container images. This is used across the company, from running internal research projects to powering its commercial products. It's also using Cloud Build for its continuous integration pipelines, where code is pushed into Git repositories where it’s built and tested, before being deployed via Docker images and internal Python packages.
"Google Kubernetes Engine provides a uniformity that crosses all of our technological work," Peleg adds. "It’s allowed us to easily provide our data scientists with the ability to run thousands of GPU hours in a day. We’ve easily been able to take research code or algorithmic code, and integrate it into our production environment. It’s another key offering from Google Cloud and provides significant benefits to the way we work."
"We’re a startup among giants, and we knew that we had to adopt the best technology available to achieve our ambitions. For us, Google Cloud and its wide range of advanced machine learning accelerators was an ideal fit."
—Barak Peleg, VP Engineering at AI21 LabsMeeting the demands of different audiences
One of the key attractions for AI21 Labs in using Google Cloud was the variety of GPUs, and related price points, on offer. That meant AI21 Labs could use the most appropriate GPU for the task at hand. This approach has enabled the development of both a B2C and B2B offering, one of which is Wordtune, the company’s flagship product launched in October 2020. It’s a writing companion that helps you rephrase your sentences to say exactly what you mean through clear, compelling, and authentic writing.
In addition, AI21 Studio is a developer platform providing an API and interactive web environment where businesses and other organizations can query their language models. Applications of this tool include summarizing texts; topic classification tasks, such as classifying news articles based on their title and their summary; ideation; copywriting, and even generating code in programming languages.
"The acceleration requirements for each of these products are different," Peleg explains. "Wordtune products have less demanding requirements compared to AI21 studio. The great thing about Google Cloud is that it enables us to use a variety of GPUs or high-performance chipsets, as and when we need them. That means we can satisfy the demand of our different audiences in a cost-effective way. Because we are based on GKE, for example, we can now give them the flexibility to choose their price-performance point, between 7x higher cost and 2x better performance with changing architecture."
Processing and analyzing data to extract business insight
When it comes to working with data, AI21 Labs uses BigQuery to batch logging data. This allows it to achieve several things, depending on its needs. In AI21 Studio, it can provide customers with information, such as accurate billing information based on their usage. In both Wordtune and AI21 Studio, it allows the team to internally drill down into user usage, providing analysis of behavioral data to offer insights into the benefits they’re gaining from each product. The team also uses Looker Studio as a visualization tool on top of BigQuery to look at data taken from its marketing platforms and payment providers to connect the dots between spend and revenue.
"We’re still in the early stages of using business intelligence and data visualization," adds Peleg. "But BigQuery is an invaluable data warehouse solution. In Wordtune, for example, we’re able to connect behavioral data, such as clicks and pageviews, with API requests, performed to our back-end systems. So, we're able to correlate or connect between these two separate data sources, and better understand our system's behavior."
"BigQuery is an invaluable data warehouse solution. In Wordtune, for example, we’re able to connect behavioral data, such as clicks and pageviews, with API requests, performed to our back-end systems. So, we're able to correlate or connect between these two separate data sources, and better understand our system's behavior."
—Barak Peleg, VP Engineering at AI21 LabsReaping the benefits of a successful partnership
Working with Google Cloud has brought a lot of benefits to AI21 Labs, says Peleg. It’s enabled the team to build the tools that they need while improving specific workflows that its researchers need. And, while technology and having the right tools is clearly a key component for success, the human dimension of these changes cannot be overlooked. Peleg adds: "Our collaboration with Google Cloud has had a really profound effect. There's a really good feeling among the team that they are able to achieve what they want, without needing a huge amount of expertise when it comes to Google Cloud."
Another vital element to this success is the reliability and reproducibility of results in the Google Cloud environment, he adds. That means the team is able to alter specific parameters, or sets of parameters, to repeat experiments and monitor the outcomes of their changes. That’s vital to improving NLP models. It also means that they can create multiple identical production environments, altering their scale depending on need. In practice, that might include a testing area to try out new technologies from Google Cloud, an internal production environment for 'dogfooding' features that are in development, as well as external environments for delivering services to customers.
Accelerating performance to meet future demand
AI21 Labs is off to a great start and says it sees Google Cloud as central to the future success of the business. The company is already seeing great performance-boosting benefits. Elasticity, for example, allows saving of up to 20% due to its variation in system load with no performance tradeoffs. And the team hopes to be able to improve on this in the future, particularly when it comes to accelerating ML. "That’s a key part of our roadmap," adds Peleg. "We’re interested in really optimizing and further automating our processes, particularly when it comes to our training stage and our online inference use cases. Google Cloud is continually improving its offering and communicating those benefits to us, so we’re excited to see what more we can get out of our partnership in the future."
"Our collaboration with Google Cloud has had a really profound effect. There's a really good feeling among the team that they are able to achieve what they want, without needing a huge amount of expertise when it comes to Google Cloud."
—Barak Peleg, VP Engineering at AI21 LabsTell us your challenge. We're here to help.
Contact us
About AI21 Labs
AI21 Labs was founded in 2017 to reimagine the way humans read and write. The company uses neural networks fuelled by huge volumes of data and powerful computing to enable machines to understand written language and build tools that complement the skills of humans.