ConsistentEE: A Consistent and Hardness-Guided Early Exiting Method for Accelerating Language Models Inference

TL;DR

ConsistentEE introduces a reinforcement learning-based early exiting method for language models that aligns training and inference, improving efficiency by adaptively deciding when to exit based on instance hardness.

Contribution

It formulates early exiting as a reinforcement learning problem with a policy network and introduces the Memorize Layer to better handle instance hardness.

Findings

Outperforms baseline methods on multiple NLP tasks

Balances accuracy and speed effectively for easy and hard instances

Demonstrates significant inference acceleration without accuracy loss

Abstract

Early Exiting is one of the most popular methods to achieve efficient inference. Current early exiting methods adopt the (weighted) sum of the cross entropy loss of all internal classifiers during training, imposing all these classifiers to predict all instances correctly. However, during inference, as long as one internal classifier predicts an instance correctly, it can accelerate without losing accuracy. Thus, there is a notable gap between training and inference. We propose ConsistentEE, an early exiting method that is consistent in training and inference. ConsistentEE formulates the early exiting process as a reinforcement learning problem. A policy network is added to decide whether an instance should exit or continue. The training objective of ConsistentEE only require each instance to be predicted correctly by one internal classifier. Additionally, we introduce the concept Memorize Layer to measure the hardness of an instance. We incorporate memorized layer into reward function design, which allows "easy" instances to focus more on acceleration while "hard" instances to focus more on accuracy. Experimental results show that our method outperforms other baselines on various natural language understanding and generation tasks.