Accelerating inference of transformer-based models

Abstract

Methods, systems, and computer program products for accelerating inference of transformer-based models are provided herein. A computer-implemented method includes obtaining a machine learning model comprising a plurality of transformer blocks, a task, and a natural language dataset; generating a compressed version of the machine learning model based on the task and the natural language dataset, wherein the generating comprises: obtaining at least one set of tokens, wherein each token in the set corresponds to one of the items in the natural language dataset, identifying and removing one or more redundant output activations of different ones of the plurality of transformer blocks for the at least one set of tokens, and adding one or more input activations corresponding to the one or more removed output activations into the machine learning model at subsequent ones of the plurality of the transformer blocks; and outputting the compressed version of the machine learning model to at least one user.