Latency Adjustable Transformer Encoder for Language Understanding

TL;DR

This paper introduces a Transformer architecture that adaptively adjusts inference latency by removing less important word-vectors during fine-tuning, enabling flexible speed-accuracy trade-offs without additional training.

Contribution

It proposes a novel Attention Context Contribution metric and an offline-tuning method for latency adjustment in Transformer models, enhancing efficiency without significant performance loss.

Findings

Up to 2.9x improvement in Time-to-First-Token for Llama3

Effective removal of less important word-vectors in higher layers

Minimal impact on global context and task performance

Abstract

Adjusting the latency, power, and accuracy of natural language understanding models is a desirable objective of an efficient architecture. This paper proposes an efficient Transformer architecture that adjusts the inference computational cost adaptively with a desired inference latency speedup. In fine-tuning phase, the proposed method detects less important hidden sequence elements (word-vectors) and eliminates them in each encoder layer using a proposed Attention Context Contribution (ACC) metric. After the fine-tuning phase, with the novel offline-tuning property, the inference latency of the model can be adjusted in a wide range of inference speedup selections without any further training. Extensive experiments reveal that most word-vectors in higher Transformer layers contribute less to subsequent layers, allowing their removal to improve inference latency. Experimental results on various language understanding, text generation, and instruction tuning tasks and benchmarks demonstrate the approach's effectiveness across diverse datasets, with minimal impact on the input's global context. The technique improves Time-to-First-Token (TTFT) of Llama3 by up to 2.9x, with minor performance drop. The suggested approach posits that in Large Language Models (LLMs), although the complete network is necessary for training, it can be truncated during the fine-tuning phase.