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

US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
LU, Luxembourg
Are you a talented and inventive scientist with a strong passion about modern data technologies and interested to improve business processes, extracting value from the data? Would you like to be a part of an organization that is aiming to use self-learning technology to process data in order to support the management of the procurement function? The Global Procurement Technology, as a part of Global Procurement Operations, is seeking a skilled Data Scientist to help build its future data intelligence in business ecosystem, working with large distributed systems of data and providing Machine Learning (ML) and Predictive Modeling expertise. You will be a member of the Data Engineering and ML Team, joining a fast-growing global organization, with a great vision to transform the Procurement field, and become the role model in the market. This team plays a strategic role supporting the core Procurement business domains as well as it is the cornerstone of any transformation and innovation initiative. Our mission is to provide a high-quality data environment to facilitate process optimization and business digitalization, on a global scale. We are supporting business initiatives, including but not limited to, strategic supplier sourcing (e.g. contracting, negotiation, spend analysis, market research, etc.), order management, supplier performance, etc. We are seeking an individual who can thrive in a fast-paced work environment, be collaborative and share knowledge and experience with his colleagues. You are expected to deliver results, but at the same time have fun with your teammates and enjoy working in the company. In Amazon, you will find all the resources required to learn new skills, grow your career, and become a better professional. You will connect with world leaders in your field and you will be tackling Data Science challenges to ensure business continuity, by taking the right decisions for your customers. As a Data Scientist in the team, you will: -be the subject matter expert to support team strategies that will take Global Procurement Operations towards world-class predictive maintenance practices and processes, driving more effective procurement functions, e.g. supplier segmentation, negotiations, shipping supplies volume forecast, spend management, etc. -have strong analytical skills and excel in the design, creation, management, and enterprise use of large data sets, combining raw data from different sources -provide technical expertise to support the development of ML models to facilitate intelligent digital services, such as Contract Lifecycle Management (CLM) and Negotiations platform -cooperate closely with different groups of stakeholders, e.g. data/software engineers, product/program managers, analysts, senior leadership, etc. to evaluate business needs and objectives to set up the best data management environment -create and share with audiences of varying levels technical papers and presentations -deal with ambiguity, prioritizing needs, and delivering results in a dynamic environment Basic qualifications -Master’s Degree in Computer Science/Engineering, Informatics, Mathematics, or a related technical discipline -3+ years of industry experience in data engineering/science, business intelligence or related field -3+ years experience in algorithm design, engineering and implementation for very-large scale applications to solve real problems -Very good knowledge of data modeling and evaluation -Very good understanding of regression modeling, forecasting techniques, time series analysis, machine-learning concepts such as supervised and unsupervised learning, classification, random forest, etc. -SQL and query performance tuning skills Preferred qualifications -2+ years of proficiency in using R, Python, Scala, Java or any modern language for data processing and statistical analysis -Experience with various RDBMS, such as PostgreSQL, MS SQL Server, MySQL, etc. -Experience architecting Big Data and ML solutions with AWS products (Redshift, DynamoDB, Lambda, S3, EMR, SageMaker, Lex, Kendra, Forecast etc.) -Experience articulating business questions and using quantitative techniques to arrive at a solution using available data -Experience with agile/scrum methodologies and its benefits of managing projects efficiently and delivering results iteratively -Excellent written and verbal communication skills including data visualization, especially in regards to quantitative topics discussed with non-technical colleagues
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, CA, San Francisco
About Twitch Launched in 2011, Twitch is a global community that comes together each day to create multiplayer entertainment: unique, live, unpredictable experiences created by the interactions of millions. We bring the joy of co-op to everything, from casual gaming to world-class esports to anime marathons, music, and art streams. Twitch also hosts TwitchCon, where we bring everyone together to celebrate, learn, and grow their personal interests and passions. We’re always live at Twitch. Stay up to date on all things Twitch on Linkedin, Twitter and on our Blog. About the role: Twitch builds data-driven machine learning solutions across several rich problem spaces: Natural Language Processing (NLP), Recommendations, Semantic Search, Classification/Categorization, Anomaly Detection, Forecasting, Safety, and HCI/Social Computing/Computational Social Science. As an Intern, you will work with a dedicated Mentor and Manager on a project in one of these problem areas. You will also be supported by an Advisor and participate in cohort activities such as research teach backs and leadership talks. This position can also be located in San Francisco, CA or virtual. You Will: Solve large-scale data problems. Design solutions for Twitch's problem spaces Explore ML and data research
US, WA, Seattle
We are a team of doers working passionately to apply cutting-edge advances in deep learning in the life sciences to solve real-world problems. As a Senior Applied Science Manager you will participate in developing exciting products for customers. Our team rewards curiosity while maintaining a laser-focus in bringing products to market. Competitive candidates are responsive, flexible, and able to succeed within an open, collaborative, entrepreneurial, startup-like environment. At the leading edge of both academic and applied research in this product area, you have the opportunity to work together with a diverse and talented team of scientists, engineers, and product managers and collaborate with others teams. Location is in Seattle, US Embrace Diversity Here at Amazon, we embrace our differences. We are committed to furthering our culture of inclusion. We have ten employee-led affinity groups, reaching 40,000 employees in over 190 chapters globally. We have innovative benefit offerings, and host annual and ongoing learning experiences, including our Conversations on Race and Ethnicity (CORE) and AmazeCon (gender diversity) conferences. Amazon’s culture of inclusion is reinforced within our 14 Leadership Principles, which remind team members to seek diverse perspectives, learn and be curious, and earn trust Balance Work and Life Our team puts a high value on work-life balance. It isn’t about how many hours you spend at home or at work; it’s about the flow you establish that brings energy to both parts of your life. We believe striking the right balance between your personal and professional life is critical to life-long happiness and fulfillment. We offer flexibility in working hours and encourage you to find your own balance between your work and personal lives Mentor & Grow Careers Our team is dedicated to supporting new members. We have a broad mix of experience levels and tenures, and we’re building an environment that celebrates knowledge sharing and mentorship. Our senior members enjoy one-on-one mentoring and thorough, but kind, code reviews. We care about your career growth and strive to assign projects based on what will help each team member develop into a better-rounded engineer and enable them to take on more complex tasks in the future. Key job responsibilities • Manage high performing engineering and science teams • Hire and develop top-performing engineers, scientists, and other managers • Develop and execute on project plans and delivery commitments • Work with business, data science, software engineer, biological, and product leaders to help define product requirements and with managers, scientists, and engineers to execute on them • Build and maintain world-class customer experience and operational excellence for your deliverables
US, Virtual
The Amazon Economics Team is hiring Interns in Economics. We are looking for detail-oriented, organized, and responsible individuals who are eager to learn how to work with large and complicated data sets. Some knowledge of econometrics, as well as basic familiarity with Stata, R, or Python is necessary. Experience with SQL, UNIX, Sawtooth, and Spark would be a plus. These are full-time positions at 40 hours per week, with compensation being awarded on an hourly basis. You will learn how to build data sets and perform applied econometric analysis at Internet speed collaborating with economists, data scientists and MBAʼs. These skills will translate well into writing applied chapters in your dissertation and provide you with work experience that may help you with placement. Roughly 85% of interns from previous cohorts have converted to full time economics employment at Amazon. If you are interested, please send your CV to our mailing list at econ-internship@amazon.com.
US, WA, Seattle
Amazon internships are full-time (40 hours/week) for 12 consecutive weeks with start dates in May - July 2023. Our internship program provides hands-on learning and building experiences for students who are interested in a career in hardware engineering. This role will be based in Seattle, and candidates must be willing to work in-person. Corporate Projects (CPT) is a team that sits within the broader Corporate Development organization at Amazon. We seek to bring net-new, strategic projects to life by working together with customers and evolving projects from ZERO-to-ONE. To do so, we deploy our resources towards proofs-of-concept (POCs) and pilot programs and develop them from high-level ideas (the ZERO) to tangible short-term results that provide validating signal and a path to scale (the ONE). We work with our customers to develop and create net-new opportunities by relentlessly scouring all of Amazon and finding new and innovative ways to strengthen and/or accelerate the Amazon Flywheel. CPT seeks an Applied Science intern to work with a diverse, cross-functional team to build new, innovative customer experiences. Within CPT, you will apply both traditional and novel scientific approaches to solve and scale problems and solutions. We are a team where science meets application. A successful candidate will be a self-starter comfortable with ambiguity, strong attention to detail, and the ability to work in a fast-paced, ever-changing environment. As an Applied Science Intern, you will own the design and development of end-to-end systems. You’ll have the opportunity to create technical roadmaps, and drive production level projects that will support Amazon Science. You will work closely with Amazon scientists, and other science interns to develop solutions and deploy them into production. The ideal scientist must have the ability to work with diverse groups of people and cross-functional teams to solve complex business problems.