PaperlayerView exact revision on arXiv

Source-linked AI summary

Fast and Robust Early-Exiting Framework for Autoregressive Language Models with Synchronized Parallel Decoding

Sangmin Bae, Jongwoo Ko, Hwanjun Song, Se-Young Yun

arXiv:2310.05424v1cs.CL

TL;DR

Autoregressive language models face high inference latency, and conventional early-exiting frameworks have challenges involving state copying, exit paths, and confidence-threshold selection. FREE combines shallow-deep processing, synchronized parallel decoding, and adaptive threshold estimation; across generation tasks, the authors report faster inference without compromising output quality.

  • Problem

    Autoregressive language models have high inference latency, while early-exiting frameworks face state-copying degradation, costs from numerous exits, and effortful confidence-threshold selection.

    Source: §1 · p. 1

  • Method

    FREE combines a shallow-deep module and synchronized parallel decoding with adaptive confidence-threshold estimation using a Beta mixture model.

    Source: §1 · p. 2, §5.3 · p. 6

  • Results

    Across extensive generation tasks, FREE achieved superior performance with significant latency acceleration without compromising generated-output quality.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The paper concludes that FREE accelerates inference without compromising generated-output quality across the evaluated generation tasks.

    Source: §7 · p. 9

    FREE requires additional computational resources to fine-tune the shallow model, and further investigation is needed to select its depth for different language models.

    Source: §7 · p. 9

Abstract

from arXiv · show

To tackle the high inference latency exhibited by autoregressive language models, previous studies have proposed an early-exiting framework that allocates adaptive computation paths for each token based on the complexity of generating the subsequent token. However, we observed several shortcomings, including performance degradation caused by a state copying mechanism or numerous exit paths, and sensitivity to exit confidence thresholds. Consequently, we propose a Fast and Robust Early-Exiting (FREE) framework, which incorporates a shallow-deep module and a synchronized parallel decoding. Our framework enables faster inference by synchronizing the decoding process of the current token with previously stacked early-exited tokens. Furthermore, as parallel decoding allows us to observe predictions from both shallow and deep models, we present a novel adaptive threshold estimator that exploits a Beta mixture model to determine suitable confidence thresholds. We empirically demonstrated the superiority of our proposed framework on extensive generation tasks.

1 Introduction

Autoregressive language models incur high inference latency because their Transformer layers are computed sequentially for each token. FREE combines a shallow-deep module, synchronized parallel decoding, and adaptive threshold estimation to address shortcomings of conventional early-exiting methods.

1 Introduction

  • Early-exiting methods adapt computation per token based on generation complexity, aiming to accelerate inference without reducing generation quality relative to a full model.Autoregressive models compute multiple Transformer layers sequentially for each token, creating high latency and computational burdens for real-time adaptability.Source: §1 · p. 1
  • State copying performs poorly with larger models and longer outputs, while numerous exit paths can add computation or produce abnormally long outputs, and threshold choice requires effort and overhead.These are challenges identified for conventional early-exiting frameworks.Source: §1 · p. 1
  • FREE combines a shallow-deep module, synchronized parallel decoding, and adaptive threshold estimation, reporting consistent speedup and performance without expensive threshold search.The framework synchronizes the current token with previously stacked early-exited tokens.Source: §1 · p. 2
  • Synchronized parallel decoding stacks consecutive shallow-exited tokens until a non-exiting token, then processes them together using actual attention key and value states.This avoids approximated states from state copying and is more efficient than decoding each token autoregressively.Source: §1 · p. 2
  • A Beta mixture model uses shallow- and deep-model prediction alignment to determine dataset-specific confidence thresholds during inference.Parallel decoding provides predictions from the deep model even for early-exited tokens.Source: §1 · p. 2

2 Related Work

Prior work studied early exiting and parallel decoding, while FREE combines synchronized decoding within one network with adaptive threshold estimation. The related methods and their constraints differ in model setting, calibration needs, and decoding strategy.

2.1 Early-exiting Framework

  • Elbayad et al. (2020) adjusted decoder depth using intermediate-layer predictions, while Schuster et al. (2022) explored confidence measures, decaying thresholds, and calibration methods.These approaches target a trade-off between inference speed and accuracy.Source: §2.1 · p. 2
  • Earlier early-exiting experiments were primarily conducted on small decoder models, and their calibration procedures required extra training time on additional calibration sets.The cited discussion says larger-model validation was needed and that the calibration overhead hindered real deployment.Source: §2.1 · p. 2

2.2 Parallel Decoding

  • Prior parallel-decoding work includes non-autoregressive generation and speculative decoding, where a small approximation model predicts autoregressively and a large target model verifies predictions in parallel.Related approaches also use models of varying depth with speculative sampling or rollback refinement.Source: §2.2 · pp. 2–3
  • FREE introduces synchronized parallel decoding within one network, alongside a shallow-deep module, and estimates dataset-specific thresholds from predictions at different model depths.The cited discussion contrasts this approach with refinement strategies that may cause unbounded latency increases by restarting after incorrect predictions.Source: §2.2 · p. 3

3 Preliminary

The preliminary section describes Transformer computation and how early exiting predicts tokens before the final layer when confidence exceeds a threshold. It also introduces state copying as a way to approximate deeper-layer attention states for subsequent tokens.

3 Preliminary

  • A Transformer has L layers, each with multi-head attention and a feed-forward sublayer, and the final decoder layer produces token predictions through a softmax classifier.The hidden-state computation described is for the next time step.Source: §3 · p. 3
  • Early exiting permits a next-token prediction at an earlier layer when its confidence score exceeds a predefined threshold.The prediction is made from the earlier layer's hidden state.Source: §3 · p. 3
  • Table 1 compares ROUGE-L scores for a full model and oracle-exiting, and reports cosine similarity between their final-layer and oracle-exited hidden states.The caption identifies the compared model outputs and hidden-state measure, but supplies no values.Source: Table 1 · p. 3
  • After a token exits early, state copying duplicates its early-layer hidden state into deeper blocks to approximate key and value states needed for later self-attention.The cited prior work found no detrimental effect for state copying from lower layers in small T5 models.Source: §3 · p. 3

4 Re-evaluating Early-exit Framework

The re-evaluation examines early-exiting across T5 and LongT5 models and summarization datasets, finding vulnerabilities tied to model size, output length, exit position, computational overhead, and threshold choice.

4 Re-evaluating Early-exit Framework

  • The re-evaluation used T5 model sizes on SAMSum and CNN/DailyMail, and LongT5-base on Multi-News and BIGPATENT.It reports four new findings about the early-exiting framework.Source: §4 · p. 3

4.1 Lack of Robustness to Model Size and Output Sequence Length

  • Generation-quality drops with oracle exiting and state copying were 0.67 on SAMSum for T5-small, 1.24 on SAMSum for T5-large, 7.99 on Multi-News, and 4.69 on BIGPATENT.The hidden-state cosine similarity between final and oracle-exited layers also decreased as models grew and output sequences lengthened.Source: §4.1 · pp. 3–4

4.2 Performance Drop by Exit Position

  • Exiting T5-large tokens at the first or second layer yielded nearly zero ROUGE-L and abnormally long outputs; the study excluded the first two or four candidate layers for base or large models, respectively.Training objectives combine losses across decoder layers; prior unweighted or weighted averages facilitated earlier-layer learning without harming full-model performance on small decoders.Source: §4.2 · p. 4, Figure 2 · p. 4

4.3 Non-negligible Computational Cost

  • Early-exiting often increased total decoding time versus the full model, with confidence softmax and duplicated-hidden-state key/value computation adding overhead.Figure 3 breaks down component costs across three summarization datasets and reports ROUGE-L scores above the bars.Source: §4.3 · p. 4, Figure 3 · p. 4

4.4 Disparate Optimal Confidence Threshold

  • Optimal confidence thresholds for lowest latency at the same performance varied by dataset: SQuAD and CNN/DailyMail tolerated lower thresholds, while IWSLT required higher ones.Table 2 presents selected thresholds and corresponding performance.Source: §4.4 · p. 5, Table 2 · p. 5

5 Novel Early-Exiting Framework: FREE

FREE combines a shallow-deep module with synchronized parallel decoding and adaptive confidence-threshold estimation. Its three components define the framework's model paths, token processing, and threshold selection.

5 Novel Early-Exiting Framework: FREE

  • FREE combines a shallow-deep module, synchronized parallel decoding, and confidence estimation to make early exiting more robust.Source: §5 · p. 5

5.1 Shallow-Deep Module

  • The shallow-deep module assigns predetermined early layers to the shallow model and all layers to the deep model, using dynamically mapped layerwise distillation in experiments.Dynamic mapping aligns each deep-model layer with its closest shallow-model counterpart while preserving layer order; the paper uses KD-dyna in all shallow-deep experiments.Source: §5.1 · p. 5

5.2 Synchronized Parallel Decoding

  • Synchronized decoding stacks consecutive tokens exited by the shallow model, then processes them with the next non-exiting token through the deep model.This parallel computation synchronously computes the stacked tokens' key and value states, avoiding approximated states from copying.Source: §5.2 · pp. 5–6

5.3 Adaptive Threshold Estimation

  • The adaptive estimator fits a beta mixture model to shallow-model confidence and prediction alignment, then selects a threshold where the posterior reaches ζ.It gathers samples during early decoding rather than using an extra calibration set; the authors set ζ to 0.4 with class priors fixed at 0.5.Source: §5.3 · p. 6

6 Experiments

Experiments compare FREE and early-exiting baselines across generation tasks, model sizes, and ablations. FREE with adaptive threshold estimation achieved up to 2.16× speedup while retaining 99% of full-model performance.

6.1 Experimental Setup

  • The setup covered question answering, machine translation, and four summarization datasets, using LongT5-base for Multi-News and BIGPATENT and T5-large for the other datasets.The work evaluates both FREE without adaptive threshold estimation (FREE†) and FREE with it.Source: §6.1 · p. 6, §6.2 · p. 7

6.2 Experimental Results

  • FREE† showed robust performance and larger quality-latency AUC than baselines, while CALM performed worse than static exiting across datasets.FREE† adjusts exit thresholds; the authors attribute CALM’s weaker performance likely to state copying and numerous exit positions.Source: §6.2 · p. 7
  • Up to 2.16× speedup was achieved while retaining 99% of full-model performance with adaptive threshold estimation.Table 3 compares early-exiting frameworks across datasets, reporting performance at the smallest threshold that reaches 99% of full-model performance.Source: §6.2 · p. 7, Table 3 · p. 7
  • On T5-3B SAMSum and CNN/DailyMail, FREE maintained latency and ROUGE-L superiority over baselines.The experiments used LoRA adapters targeting self-attention and feed-forward layers, with rank 64.Source: §6.2 · p. 8

6.3 Ablation Study

  • Six-layer shallow models retained robust performance and speedup, whereas reducing depth too far, such as to four layers, decreased performance and speed gain.Lower shallow-model depth tended toward higher speedup, but very low depth was associated with incorrect and redundant output sentences.Source: §6.3 · p. 8
  • Synchronized parallel decoding outperformed state copying in ROUGE-L across three datasets, with the authors attributing the gain to accurate Transformer-layer computation of updated hidden states.Both decoding methods used the shallow-deep module.Source: §6.3 · p. 8, Table 5 · p. 8
  • Using only 3% of total samples let the adaptive estimator approximate the threshold measured on the full sample set, with minimal extra computation time.The calibration set was formed from early-stage instances and updated iteratively.Source: §6.3 · p. 8, Table 6 · p. 8
  • Low refinement thresholds produced minimal performance improvement relative to added latency, and methods without an upper bound on latency increase may not suit early-exiting.The evaluation tested a refinement method that removes subsequent tokens and restarts generation after detecting a wrong token.Source: §6.3 · pp. 8–9, Table 7 · p. 9

7 Conclusion

FREE combines a shallow-deep module, synchronized parallel decoding, and adaptive threshold estimation to address conventional early-exiting challenges. Experiments report latency acceleration without compromising generated-output quality, while the approach requires shallow-model fine-tuning and further depth-selection study.

7 Conclusion

  • FREE combines a shallow-deep module, synchronized parallel decoding, and adaptive threshold estimation.The conclusion presents these as the framework’s three key components.Source: §7 · p. 9
  • Across generation tasks, FREE achieved significant latency acceleration without compromising generated-output quality.The conclusion reports this as the empirical outcome of extensive experiments.Source: §7 · p. 9
  • FREE requires additional computation to fine-tune the shallow model, although parameter-efficient fine-tuning is described as a promising solution.The appropriate shallow-model depth for different language models remains to be investigated.Source: §7 · p. 9

A Dataset Description

The experiments use datasets for summarization, question answering, and machine translation. The appendix describes six datasets and specifies German-to-English translation and the Fixed Constructions subset of BIGPATENT.

A Dataset Description

  • The experimental tasks include summarization, question answering, and machine translation.The dataset descriptions cover SAMSum, CNN/DailyMail, Multi-News, BIGPATENT, SQuAD, and IWSLT 2017.Source: A Dataset Description · p. 13
  • SAMSum contains 16K messenger-like conversations paired with third-person summaries, while CNN/DailyMail contains over 300K English news articles and supports abstractive and extractive summarization.Source: A Dataset Description · p. 13
  • Multi-News comprises 45K news articles with professionally crafted summaries linked to their source articles.Source: A Dataset Description · p. 13
  • BIGPATENT contains 1.3M U.S. patent records with human-written abstractive summaries; this work uses its Fixed Constructions category.Source: A Dataset Description · p. 13
  • SQuAD comprises 87.6K reading-comprehension tasks based on questions about Wikipedia articles, while IWSLT 2017 is evaluated for German-to-English translation.Source: A Dataset Description · p. 13

B Detailed Experimental Setup

The detailed setup reports the training optimizer and learning rate, adaptive-threshold settings, and task-specific quality metrics. Table 8 summarizes optimized hyperparameters for training shallow-deep T5 models.

B Detailed Experimental Setup

  • Table 8 summarizes optimized hyperparameters for training shallow-deep T5 models, including batch and maximum input and output lengths.The caption defines ‘# Batch’ as batch size per GPU multiplied by GPU count.Source: Table 8 · p. 13
  • All datasets use AdaFactor with a learning rate of 1e-4.Source: B Detailed Experimental Setup · p. 13
  • Adaptive-threshold estimation uses an initial threshold value λ0_c of 0.9, ζ of 0.4, and T equal to 3% of the total sample number.Source: B Detailed Experimental Setup · p. 13
  • Output quality is measured with F1 for SQuAD, BLEU for IWSLT 2017, and ROUGE for the four summarization tasks.Source: B Detailed Experimental Setup · p. 13

C Inference Latency Evaluation

The inference-latency evaluation measures decoding time under a batch-size-1 setup and presents quality–latency trade-offs. The additional experiments assess the method and its components.

C Inference Latency Evaluation

  • The evaluation uses 500 predictions per dataset and configuration, batch size 1, greedy decoding with beam size 1, and measures time through completion.Runs use a compiled PyTorch function on one NVIDIA GeForce RTX 3039 GPU and an Intel Core i7-12700K CPU.Source: C Inference Latency Evaluation · p. 13
  • Figure 9 compares output quality and normalized latency across exit conditions; its dashed line marks the full model's F1 and BLEU scores.The inner point of the Pareto curve is excluded.Source: Figure 9 · p. 14
  • Figure 10 shows the performance–normalized-latency trade-off across exit thresholds from 0.0 to 0.9, with latency normalized to a fine-tuned full-model baseline.The tested thresholds are 0.0, 0.1, 0.3, 0.5, 0.7, and 0.9.Source: Figure 10 · p. 14

D Additional Experimental Results

D.1 Performance on Different Datasets

FREE† and FREE outperform CALM and static-exiting baselines consistently on SQuAD, while their advantage is somewhat reduced on IWSLT. Further experiments compare distillation losses, model sizes, and decoding strategies.

D.1 Performance on Different Datasets

  • FREE† and FREE consistently outperform CALM and static-exiting baselines on SQuAD, and FREE outperforms other baselines across various datasets.Source: §D.1 · p. 14
  • IWSLT performance advantages are slightly reduced; the authors attribute this to mT5's larger vocabulary increasing confidence-measurement time.Source: §D.1 · p. 14

D.2 Layerwise Knowledge Distillation

  • KD-dyna offers better performance–latency trade-offs than weighted-average or other KD-based losses, though the differences are not significant.The lower KD-unif performance on SAMSum suggests dynamically determining layer mapping can facilitate more effective knowledge transfer.Source: §D.2 · p. 14

D.3 Comparison with Small-sized Models

  • On CNN/DailyMail, FREE has competitive inference speedup versus T5-base; on SQuAD, it reports superior F1 and significantly higher speedup.On SQuAD, T5-base's ROUGE-L score is 90.50, versus 90.24 for FREE's six-layer shallow model.Source: §D.3 · pp. 14–15, Table 9 · p. 14

D.4 Various Decoding Strategies

  • Across top-k and nucleus sampling, FREE maintains consistent, robust performance and achieves larger speedup than CALM.Source: §D.4 · p. 15, Table 10 · p. 15
Loading 2310.05424v1…