A version of the BERT language model that’s 20 times as fast

Determining the optimal architectural parameters reduces network size by 84% while improving performance on natural-language-understanding tasks.

In natural-language understanding (NLU), the Transformer-based BERT language model is king. Its high performance on multiple tasks has strongly influenced contemporary NLU research. 

On the other hand, it is a relatively big and slow model, which makes it unsuitable for some applications. Multiple efforts have been made to compress the BERT architecture, but the choice of architectural parameters (the number of layers, the number of processing nodes per layer, and so on) has been somewhat arbitrary, and the resulting models are rarely much better than the original at optimizing the balance between the model’s size, speed, and error rate.

A few weeks ago, we released part of the code for Bort, a highly optimized language model (LM) extracted from the BERT architecture through a combination of two rigorous algorithmic techniques especially designed for neural-network compression.

We tested Bort on 23 NLU tasks, and on 20 of them, it improved on BERT’s performance — by 31% in one case — even though it’s 16% the size and about 20 times as fast.

The code doesn’t contain the full versions of the algorithms we used to extract Bort from BERT, but I will discuss them here briefly.

Principled approach

A solid understanding of the intrinsic difficulty of a problem allows us to design efficient and correct algorithms for solving that problem that can be trusted to perform consistently. Think of trying to sort an array of numbers by just trying out random permutations: this approach will quickly become intractable as the size of the input grows. Thankfully, sorting is a well-studied problem with fast and correct solutions and a nicely bounded runtime growth. Why can’t compressing BERT, or any network, be the same way? 

We wanted to produce a version of BERT whose architectural parameters minimized its parameter sizeinference speed, and error rate, and I wanted to show that these architectural parameters were optimal. I call this problem “optimal subarchitecture extraction”, or OSE. Note that this is a different compression approach than weight pruning, which heuristically removes connections from a previously trained network. 

An algorithm solving OSE should also work for any input, not just BERT; a general-purpose algorithm would save a lot of time — and GPU cycles — otherwise spent on trial and error approaches.

Model showing comparison of weight pruning and optimal subarchitecture extractionfor a toy source network.
Comparison of weight pruning (left, bottom) and optimal subarchitecture extraction (or OSE, bottom right) for a toy source network (top). Weight pruning removes edges from a (usually trained) network, while OSE reparametrizes the layers. For example, the source network has two fully connected linear layers of dimensions 4x3 and 4x4 (second and third from the left), making its architectural parameters (2, 4, 3, 4, 4). The optimal subarchitecture ended up with two layers of dimensions 3x3 and 3x2, so the architectural parameters are (2, 3, 3, 3, 2). The parameter size was brought down from 2(4*3 + 4 + 4*4 + 4) = 72 to 2(3*3 + 3 + 3*2 + 2) = 40. 
Credit: Glynis Condon

OSE is a computationally hard problem, so the best we can hope for is an algorithm that returns something in the ballpark of the optimum. But even a ballpark algorithm could be impractically time consuming with a large enough input. 

My research showed that there is an efficient algorithm for extracting a subarchitecture whose functions — parameter size, inference speed, and error rate — have a polynomial correlation with the architectural parameters, meaning that they’re bounded by some polynomial function of the architectural parameters. 

This is easy to see for the parameter size of the network, since it boils down to a counting argument based on the number of nodes per layer, times the number of layers (see image caption above). The same argument can be made for the inference speed — the rate at which a network outputs a result given an input — as it is related to the parameter size. 

Under some circumstances, the algorithm’s error rate has a similar correlation. Furthermore, whenever the “cost” associated with the first (call it A) and last layers of the network is lower than that of the middle layers (B), the runtime of a solution will not explode. I call all these assumptions the ABnC property, which BERT turns out to have.

For inputs having the ABnC property, the algorithm I designed behaves like a fully polynomial-time approximation scheme, or FPTAS. Being an FPTAS means that an approximation parameter — which defines how far short of the optimal solution you’re willing to fall — can now be given as an input to the algorithm, and the algorithm’s execution time depends polynomially on that parameter. 

The more accurate you want the approximation to be, the longer the algorithm takes to execute. But at least the trade-off can be precisely specified, and the runtime of the algorithm is guaranteed to not grow too quickly.

I ran the FPTAS on BERT and obtained a set of architectural parameters (Bort), and from the proofs we knew that it would be Pareto optimal. That is, it optimized the balance between inference speed, parameter size, and error rate: a faster model would necessarily be bigger or more error prone; a more accurate model would be larger or slower. Any model that broke that balance would necessarily be suboptimal.

Generalizability

The true strength of BERT lies in its generalizability across tasks. BERT learns to represent words of a language as points in a shared space, where proximity in the space implies semantic similarity. That general model can then be fine-tuned on specific tasks, such as question answering or text classification.

To see whether Bort generalizes as well as BERT, we needed to pre-train it and then test it on several different natural-language-understanding tasks. The FPTAS doesn’t return a trained model, and it didn’t know that our ultimate goal was fine-tuning.

When fine-tuning, I opted to follow BERT’s approach and attach a single linear classifier to Bort. I supposed that a good LM would have no problems with this setup, but the fine-tuning process was hard, and several variations on this strategy (e.g., knowledge distillation, deeper classifiers, and hyperparameter search) failed to yield a model that reached the theoretical limits predicted by the FPTAS.

In machine learning, it is crucial to have a good selection of features to make the data representative of the task, which in turn makes the task learnable. Deep networks operate in a hierarchical fashion, which makes them fantastic at selecting features automatically. If the model is too small and the data too scarce, however, the task may not be learnable. And Bort is, by design, a small model.

The Agora algorithm

To address this problem, I developed a second algorithm, called Agora, which leverages the development set for the task Bort is being fine-tuned on. In machine learning, a labeled data set is usually split into three components: the training set is used to train a model; the development (or dev) set is used to check for over-/underfitting; and the test set is used to assess the model’s generalizability.

With Agora, Bort is first fine-tuned on the training set, then applied to the data in the dev set. Agora finds the data points in the dev set that the pretrained Bort model labeled incorrectly and samples new points near them in a chosen representation space. Those samples are labeled by a second model, and then they’re added into the training set.

It might sound silly to add randomized points from the dev set to the training set. You might expect the fine-tuned model to overfit the dev set data, so that it doesn’t generalize well to unfamiliar inputs — which means that it didn’t actually learn.

But this is a powerful approach to situations with scarce or ill-formed data. The proof of this is non-trivial, but intuitively, this algorithm works because it “reconstructs” the input in a way that is equivalent, in a precise sense, to the distribution of the task we are modeling in the first place. I believe the best part about Agora is not its effectiveness but that it is another proof of what the great Patrick Winston showed fifty years ago: you can’t learn something that you do not already know almost completely.

Graphic of Agora sampling the development set and generating, labeling and adding new points back to the training set.
In this toy example, we are attempting to learn a data set that looks like a torus. However, the training set is not representative of the distribution. Agora samples from the development set, generates new points, labels them, and adds them back to the training set. Given enough rounds, it is able to generate a data set learnable by any model — even a random guesser!

Credit: Glynis Condon

The application of both algorithms to BERT yielded Bort, which has an effective (not counting the embedding layer) size of 5.5% of the original BERT architecture and a net size of 16%. The effective size is more important because the first layer, by the ABnC property, is not as expensive as the other layers when performing inference. Indeed, Bort is up to 20 times faster, on a CPU, than BERT. It is also able to be pretrained much more rapidly than usual — likely due to the FPTAS’s preferring faster-converging architectures. It also obtained improvements of up to 31%, absolute, with respect to BERT, across multiple NLU benchmarks.

Acknowledgments: Daniel J. Perry, my coauthor for the Bort paper.

Related content

GB, London
How can Amazon improve the advertising experience for customers around the world? How can we help advertisers and customers find each other in a meaningful way? Amazon Advertising creates and transforms the connection between retailers/service providers and customers. Our teams strive to reinvent the way advertisers and agencies build brands and drive performance in their advertising. By using Amazon's foundation in e-commerce, we help brands connect with the right customers through creative solutions and formats across screens and devices, and in the physical world. Amazon Advertising seeks a Data Scientist with strong Data Analysis skills to join the ADSP engineering team split across Edinburgh and London. We make Guidance products that help optimise our customer's advertising campaign workflows and performance. As a scientist on the team, you will be involved in many aspects of the process - from idea generation, business analysis and scientific research, through to development - giving you a real sense of ownership. The systems that you help to build will operate at massive scale to advertising customers around the world. Our ideal candidate is an experienced Data scientist who has a track-record of performing analysis, applying statistical techniques and building basic ML models to solve real business problems, who has great leadership and communication skills, and who is motivated to achieve results in a fast-paced environment. Key job responsibilities Rapidly design, prototype and test many possible hypotheses in a high-ambiguity environment, making use of both quantitative analysis and business judgment. Collaborate with software engineering teams to integrate successful experimental results into large-scale, highly complex Amazon production systems. Report results in a manner which is both statistically rigorous and compellingly relevant, exemplifying good scientific practice in a business environment. Promote the culture of experimentation at Amazon.
US, NY, New York
Amazon is looking for a passionate, talented, and inventive Applied Scientist with a strong machine learning background to help build industry-leading language technology. AWS Utility Computing (UC) provides product innovations — from foundational services such as Amazon’s Simple Storage Service (S3) and Amazon Elastic Compute Cloud (EC2), to consistently released new product innovations that continue to set AWS’s services and features apart in the industry. As a member of the UC organization, you’ll support the development and management of Compute, Database, Storage, Internet of Things (Iot), Platform, and Productivity Apps services in AWS, including support for customers who require specialized security solutions for their cloud services. Our mission is to provide a delightful experience to Amazon’s customers by pushing the envelope in Natural Language Processing (NLP), Generative AI, Large Language Model (LLM), Natural Language Understanding (NLU), Machine Learning (ML), Retrieval-Augmented Generation, Responsible AI, Agent, Evaluation, and Model Adaptation. As part of our AI team in Amazon AWS, you will work alongside internationally recognized experts to develop novel algorithms and modeling techniques to advance the state-of-the-art in human language technology. Your work will directly impact millions of our customers in the form of products and services, as well as contributing to the wider research community. You will gain hands on experience with Amazon’s heterogeneous text and structured data sources, and large-scale computing resources to accelerate advances in language understanding. The Science team at AWS Bedrock builds science foundations of Bedrock, which is a fully managed service that makes high-performing foundation models available for use through a unified API. We are adamant about continuously learning state-of-the-art NLP/ML/LLM technology and exploring creative ways to delight our customers. In our daily job we are exposed to large scale NLP needs and we apply rigorous research methods to respond to them with efficient and scalable innovative solutions. At AWS Bedrock, you’ll experience the benefits of working in a dynamic, entrepreneurial environment, while leveraging AWS resources, one of the world’s leading cloud companies and you’ll be able to publish your work in top tier conferences and journals. We are building a brand new team to help develop a new NLP service for AWS. You will have the opportunity to conduct novel research and influence the science roadmap and direction of the team. Come join this greenfield opportunity! About the team Diverse Experiences AWS values diverse experiences. Even if you do not meet all of the qualifications and skills listed in the job description, we encourage candidates to apply. If your career is just starting, hasn’t followed a traditional path, or includes alternative experiences, don’t let it stop you from applying. Why AWS? Amazon Web Services (AWS) is the world’s most comprehensive and broadly adopted cloud platform. We pioneered cloud computing and never stopped innovating — that’s why customers from the most successful startups to Global 500 companies trust our robust suite of products and services to power their businesses. Inclusive Team Culture Here at AWS, it’s in our nature to learn and be curious. Our employee-led affinity groups foster a culture of inclusion that empower us to be proud of our differences. Ongoing events and learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (gender diversity) conferences, inspire us to never stop embracing our uniqueness. Mentorship & Career Growth We’re continuously raising our performance bar as we strive to become Earth’s Best Employer. That’s why you’ll find endless knowledge-sharing, mentorship and other career-advancing resources here to help you develop into a better-rounded professional. Work/Life Balance We value work-life harmony. Achieving success at work should never come at the expense of sacrifices at home, which is why we strive for flexibility as part of our working culture. When we feel supported in the workplace and at home, there’s nothing we can’t achieve in the cloud.
US, CA, Sunnyvale
The Artificial General Intelligence (AGI) team is looking for a highly skilled and experienced Senior Applied Scientist, to lead the development and implementation of cutting-edge algorithms and models for supervised fine-tuning and reinforcement learning through human feedback; with a focus across text, image, and video modalities. As a Senior Applied Scientist, you will play a critical role in driving the development of Generative AI (GenAI) technologies that can handle Amazon-scale use cases and have a significant impact on our customers' experiences. Key job responsibilities - Collaborate with cross-functional teams of engineers, product managers, and scientists to identify and solve complex problems in GenAI - Design and execute experiments to evaluate the performance of different algorithms and models, and iterate quickly to improve results - Think big about the arc of development of GenAI over a multi-year horizon, and identify new opportunities to apply these technologies to solve real-world problems - Communicate results and insights to both technical and non-technical audiences, including through presentations and written reports - Mentor and guide junior scientists and engineers, and contribute to the overall growth and development of the team
JP, 13, Tokyo
Amazon Japan is seeking an experienced Sr. Data Scientist to join our growing team. In this critical role, you will leverage your strong quantitative and analytical skills to drive data-driven insights that shape our FMCG (fast-moving consumer goods) business and other key strategic initiatives. Your responsibilities will include: - Solving complex, ambiguous business problems using appropriate statistical methodologies, modeling techniques, and data science best practices to lead business insights for FMCG business growth. You will work closely with cross-functional partners to translate business requirements into actionable data science solutions. - Designing and implementing scalable, reliable, and efficient data pipelines to extract valuable insights from diverse data sources. This includes making appropriate trade-offs between short-term and long-term needs. - Communicating your findings and recommendations clearly and persuasively to technical and non-technical stakeholders. You will document your work to the highest standards and ensure your solutions have a measurable impact on the business. - Mentoring and developing more junior data scientists on your team. You will actively participate in the hiring process and contribute to the growth of Amazon's data science community. - Staying abreast of the latest advancements in data science and applying innovative techniques where appropriate to tackle challenging business problems.
US, WA, Seattle
We are seeking a talented applied researcher to join the Whole Page Planning and Optimization (WPPO) Science team in Search. The latest data from Business Insider shows that almost 50% of online shoppers visit Amazon first. The Search WPPO Science team is responsible for developing reinforcement learning systems for the next generation Amazon shopping experience and delivering it to millions of customers. We believe that shopping on Amazon should be simple, delightful, and full of WOW moments for EVERYONE, whether you are technically savvy or new to online shopping. As an Applied Scientist, you will be working closely with a team of applied scientists and engineers to build systems that shape the future of Amazon's shopping experience by automatically generating relevant content and building a whole page experience that is coherent, dynamic, and interesting. You will improve ranking and optimization in our algorithm. You will participate in driving features from idea to deployment, and your work will directly impact millions of customers. You are going to love this job because you will: * Apply state-of-the-art Machine Learning (ML) algorithms, including Deep Learning and Reinforcement Learning, to improve hundreds of millions of customers’ shopping experience. * Have measurable business impact using A/B testing. * Work in a dynamic team that provides continuous opportunities for learning and growth. * Work with leaders in the field of machine learning.
US, WA, Bellevue
Conversational AI ModEling and Learning (CAMEL) team is part of Amazon Devices organization where our mission is to build a best-in-class Conversational AI that is intuitive, intelligent, and responsive, by developing superior Large Language Models (LLM) solutions and services which increase the capabilities built into the model and which enable utilizing thousands of APIs and external knowledge sources to provide the best experience for each request across millions of customers and endpoints. We are looking for a passionate, talented, and resourceful Applied Scientist in the field of LLM, Artificial Intelligence (AI), Natural Language Processing (NLP), Recommender Systems and/or Information Retrieval, to invent and build scalable solutions for a state-of-the-art context-aware conversational AI. A successful candidate will have strong machine learning background and a desire to push the envelope in one or more of the above areas. The ideal candidate would also have hands-on experiences in building Generative AI solutions with LLMs, enjoy operating in dynamic environments, be self-motivated to take on challenging problems to deliver big customer impact, moving fast to ship solutions and then iterating on user feedback and interactions. Key job responsibilities As an Applied Scientist, you will leverage your technical expertise and experience to collaborate with other talented applied scientists and engineers to research and develop novel algorithms and modeling techniques to reduce friction and enable natural and contextual conversations. You will analyze, understand and improve user experiences by leveraging Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in artificial intelligence. You will work on core LLM technologies, including Prompt Engineering and Optimization, Supervised Fine-Tuning, Learning from Human Feedback, Evaluation, Self-Learning, etc. Your work will directly impact our customers in the form of novel products and services.
US, WA, Bellevue
Conversational AI ModEling and Learning (CAMEL) team is part of Amazon Artificial General Intelligence (AGI) organization where our mission is to create a best-in-class Conversational AI that is intuitive, intelligent, and responsive, by developing superior Large Language Models (LLM) solutions and services which increase the capabilities built into the model and which enable utilizing thousands of APIs and external knowledge sources to provide the best experience for each request across millions of customers and endpoints. We are looking for a passionate, talented, and resourceful Applied Scientist in the field of LLM, Artificial Intelligence (AI), Natural Language Processing (NLP), Recommender Systems and/or Information Retrieval, to invent and build scalable solutions for a state-of-the-art context-aware conversational AI. A successful candidate will have strong machine learning background and a desire to push the envelope in one or more of the above areas. The ideal candidate would also have hands-on experiences in building Generative AI solutions with LLMs, enjoy operating in dynamic environments, be self-motivated to take on challenging problems to deliver big customer impact, moving fast to ship solutions and then iterating on user feedback and interactions. Key job responsibilities As an Applied Scientist, you will leverage your technical expertise and experience to collaborate with other talented applied scientists and engineers to research and develop novel algorithms and modeling techniques to reduce friction and enable natural and contextual conversations. You will analyze, understand and improve user experiences by leveraging Amazon’s heterogeneous data sources and large-scale computing resources to accelerate advances in artificial intelligence. You will work on core LLM technologies, including Supervised Fine-Tuning (SFT), In-Context Learning (ICL), Learning from Human Feedback (LHF), etc. Your work will directly impact our customers in the form of novel products and services.
IL, Tel Aviv
Are you an inventive, curious, and driven Applied Scientist with a strong background in AI and Deep Learning? Join Amazon’s AWS Multimodal generative AI science team and be a catalyst for groundbreaking advancements in Computer Vision, Generative AI, and foundational models. As part of the AWS Multimodal generative AI science team, you’ll lead innovative research projects, develop state-of-the-art algorithms, and pioneer solutions that will directly impact millions of Amazon customers. Leveraging Amazon’s vast computing power, you’ll work alongside a supportive and diverse group of top-tier scientists and engineers, contributing to products that redefine the industry. Key job responsibilities * Lead research initiatives in Multimodal generative AI, pushing the boundaries of model efficiency, accuracy, and scalability. * Design, implement, and evaluate deep learning models in a production environment. * Collaborate with cross-functional teams to transfer research outcomes into scalable AWS services. * Publish in top-tier conferences and journals, keeping Amazon at the forefront of innovation. * Mentor and guide other scientists and engineers, fostering a culture of scientific curiosity and excellence.
US, WA, Seattle
The Search Supply & Experiences team, within Sponsored Products, is seeking an Applied Scientist to solve challenging problems in natural language understanding, personalization, and other areas using the latest techniques in machine learning. In our team, you will have the opportunity to create new ads experiences that elevate the shopping experience for our hundreds of millions customers worldwide. As an Applied Scientist, you will partner with other talented scientists and engineers to design, train, test, and deploy machine learning models. You will be responsible for translating business and engineering requirements into deliverables, and performing detailed experiment analysis to determine how shoppers are responding to your changes. We are looking for candidates who thrive in an exciting, fast-paced environment and who have a strong personal interest in learning, researching, and creating new technologies with high customer impact. Key job responsibilities As an Applied Scientist on the Search Supply & Experiences team you will: - Perform hands-on analysis and modeling of enormous datasets to develop insights that increase traffic monetization and merchandise sales, without compromising the shopper experience. - Drive end-to-end machine learning projects that have a high degree of ambiguity, scale, and complexity. - Build machine learning models, perform proof-of-concept, experiment, optimize, and deploy your models into production; work closely with software engineers to assist in productionizing your ML models. - Run A/B experiments, gather data, and perform statistical analysis. - Establish scalable, efficient, automated processes for large-scale data analysis, machine-learning model development, model validation and serving. - Stay up to date on the latest advances in machine learning. About the team We are a customer-obsessed team of engineers, technologists, product leaders, and scientists. We are focused on continuous exploration of contexts and creatives where advertising delivers value to shoppers and advertisers. We specifically work on new ads experiences globally with the goal of helping shoppers make the most informed purchase decision. We obsess about our customers and we are continuously innovating on their behalf to enrich their shopping experience on Amazon
US, WA, Seattle
We are seeking a highly skilled economist to measure and understand how each Customer Service activity impacts customers. This candidate's analysis will assist teams across Amazon to prioritize defect elimination efforts and optimize how we respond to customer contacts. This candidate will partner closely with our product, program, and tech teams to deliver their findings to users via systems and dashboards that guide Customer Service planning and policy rules. Key job responsibilities - Develop Causal, Economic, and Machine Learning models at scale. - Engage in economic analysis and raise the bar for research. - Inform strategic discussions with senior leaders across the company to guide policies. A day in the life If you are not sure that every qualification on the list above describes you exactly, we'd still love to hear from you! At Amazon, we value people with unique backgrounds, experiences, and skillsets. If you’re passionate about this role and want to make an impact on a global scale, please apply! Amazon offers a full range of benefits that support you and eligible family members, including domestic partners and their children. Benefits can vary by location, the number of regularly scheduled hours you work, length of employment, and job status such as seasonal or temporary employment. The benefits that generally apply to regular, full-time employees include: 1. Medical, Dental, and Vision Coverage 2. Maternity and Parental Leave Options 3. Paid Time Off (PTO) 4. 401(k) Plan About the team The Worldwide defect elimination team's mission is to understand and resolve all issues impacting customers at scale. The Customer Service Economics and Optimization team is a force multiplier within this group, helping to understand the impact of these issues and our actions to optimize the customer experience.