Source-linked AI summary

EdgeBERT: Sentence-Level Energy Optimizations for Latency-Aware Multi-Task NLP Inference

Thierry Tambe, Coleman Hooper, Lillian Pentecost, Tianyu Jia, En-Yu Yang, Marco Donato, Victor Sanh, Paul N. Whatmough, Alexander M. Rush, David Brooks, Gu-Yeon Wei

arXiv:2011.14203v5cs.ARcs.CL

TL;DR

BERT’s computational and memory demands hinder low-latency inference on resource-constrained edge devices. EdgeBERT addresses this with latency-aware early exit and DVFS, compression, specialized hardware, and eNVM-based storage. The system reports lower energy than conventional, latency-unaware early-exit, and GPU-based baselines, while its ReRAM choice defines an explored but not exhaustive storage design space.

  • Problem

    BERT’s heavy compute and memory demands make real-time, energy-efficient on-device inference challenging; its FP32 base model consumes 432 MB.

  • Method

    EdgeBERT combines entropy-based early-exit prediction, latency-aware DVFS, adaptive attention span, pruning, quantization, specialized acceleration, and eNVM storage for shared parameters.

  • Results

    Up to 7×, 2.5×, and 53× lower energy is reported than unoptimized baseline inference, latency-unaware early exit, and CUDA adaptations on an Nvidia Jetson Tegra X2, respectively.

  • Takeaways & Limitations

    The reported results support latency-bound multi-task NLP inference on resource-constrained embedded devices through algorithm-hardware co-design.

  • Takeaways & Limitations

    The study specifically considers dense ReRAM arrays for shared embeddings, while a larger eNVM design space remains to be explored.

Abstract

from arXiv · show

Transformer-based language models such as BERT provide significant accuracy improvement for a multitude of natural language processing (NLP) tasks. However, their hefty computational and memory demands make them challenging to deploy to resource-constrained edge platforms with strict latency requirements. We present EdgeBERT, an in-depth algorithm-hardware co-design for latency-aware energy optimization for multi-task NLP. EdgeBERT employs entropy-based early exit predication in order to perform dynamic voltage-frequency scaling (DVFS), at a sentence granularity, for minimal energy consumption while adhering to a prescribed target latency. Computation and memory footprint overheads are further alleviated by employing a calibrated combination of adaptive attention span, selective network pruning, and floating-point quantization. Furthermore, in order to maximize the synergistic benefits of these algorithms in always-on and intermediate edge computing settings, we specialize a 12nm scalable hardware accelerator system, integrating a fast-switching low-dropout voltage regulator (LDO), an all-digital phase-locked loop (ADPLL), as well as, high-density embedded non-volatile memories (eNVMs) wherein the sparse floating-point bit encodings of the shared multi-task parameters are carefully stored. Altogether, latency-aware multi-task NLP inference acceleration on the EdgeBERT hardware system generates up to 7x, 2.5x, and 53x lower energy compared to the conventional inference without early stopping, the latency-unbounded early exit approach, and CUDA adaptations on an Nvidia Jetson Tegra X2 mobile GPU, respectively.

1. INTRODUCTION

EdgeBERT co-designs latency-aware early exit, DVFS, model compression, specialized hardware, and eNVM storage to reduce energy for multi-task NLP inference on resource-constrained edge devices.

  • Motivation: BERT’s compute and memory demands make on-device inference difficult, with the FP32 BERT base model consuming 432 MB.Edge deployment also requires real-time throughput and low inference and data-movement energy.
  • Method: EdgeBERT uses early-exit entropy, latency and entropy constraints, and dynamic voltage-frequency scaling to minimize energy while meeting target latency.Entropy predicts whether inference can exit early, after which voltage and frequency are adjusted.
  • Contributions: The paper’s contributions include latency-bound NLP inference, robust eNVM storage for shared word embeddings, and a compressed sparse accelerator supporting DVFS and entropy prediction.These contributions combine algorithmic optimizations with dedicated hardware and memory-system design.
  • Hardware co-design: The design combines compressed sparse computation, specialized Transformer datapaths, a fast-switching LDO, an ADPLL, and eNVM storage for shared multi-task embeddings.The accelerator is implemented as a 12nm system and addresses intermittent-operation memory costs.
  • Results: Up to 7×, 2.5×, and 53× lower energy are reported versus unoptimized inference, latency-unaware early exit, and CUDA adaptations on an Nvidia Jetson Tegra X2.The evaluation covers four NLP tasks.

2. BACKGROUND

The paper evaluates EdgeBERT on representative GLUE tasks and uses ALBERT as an efficient baseline whose shared encoder parameters and smaller embeddings reduce model size.

  • Benchmarks: GLUE evaluates sentence understanding across Single-Sentence, Similarity and Paraphrase, and Inference categories.EdgeBERT is evaluated on SST-2, QQP, QNLI, and MNLI, covering all three categories.
  • Optimization: Adaptive attention span addresses quadratic attention cost by learning how far each self-attention head needs to inspect the input.Heads can learn reduced spans or be removed when redundant.
  • Baseline: ALBERT shares Transformer encoder parameters and uses a smaller embedding size than BERT, resulting in significant compression.The embedding size is 128 for ALBERT versus 768 for BERT.
  • Optimization: EdgeBERT further reduces latency and storage requirements through adaptive attention span, early exit, network pruning, and floating-point quantization.These optimizations are applied to multi-task ALBERT inference with minimal accuracy degradation.

3. ALLEVIATING TRANSFORMER MEMORY AND COMPUTATION COSTS

EdgeBERT reduces Transformer inference cost through coordinated algorithmic optimizations that lower voltage, computation, switching activity, and capacitance while preserving latency-aware execution. The approach combines entropy-based early exit, adaptive attention spans, pruning, and floating-point quantization.

  • Energy optimization: DVFS lowers accelerator energy by reducing supply voltage, while adaptive attention spans, pruning, and quantization target clock cycles, switching activity, and capacitance.The energy abstraction identifies VDD, Ncycles, α, and C as optimization targets.
  • Entropy-based Early Exit: Early exit predicts the exit layer from the first Transformer layer to run remaining computation in an energy-optimal, latency-bounded manner.Conventional early exit can vary substantially across sentences and violate strict real-time latency constraints.
  • Adaptive Attention Span: Adaptive attention spans can disable more than half of ALBERT’s attention heads, reducing single-batch FLOPS by 1.22× and 1.18×.Eight heads can be disabled for MNLI and QQP, while seven can be disabled for SST-2 and QNLI, with minimal accuracy loss.
  • Network Pruning: Movement pruning and magnitude pruning are selected according to each task’s tolerance for increasing sparsity, while magnitude pruning is retained for embeddings to preserve data reuse.Movement pruning is applied during fine-tuning, whereas embedding-layer magnitude pruning maintains uniformity across multi-domain acceleration.
  • Floating-Point Quantization: Both weights and activations use 8-bit floating-point quantization, with a 4-bit per-layer-scaled exponent space providing the best accuracy across NLP tasks.Floating-point quantization is used because it offers 2× higher dynamic range than integer datatypes.

4. NON-VOLATILE MEMORY STORAGE OF SHARED PARAMETERS

EdgeBERT stores shared word embeddings in dense eNVMs to avoid repeated off-chip transfers and retain parameters during intermittent operation. Fault-injection experiments guide the choice of robust storage density and protect sparse bitmask metadata separately from non-zero values.

  • Shared-parameter storage: Shared word embeddings are stored in eNVMs to avoid reloading them across NLP tasks and to retain them when the system powers off between inferences.Task-specific encoder weights remain distinct because they must be updated when switching tasks.
  • Scope: The storage study focuses on dense, energy-efficient ReRAM arrays, while other emerging multilevel-cell non-volatile memories remain outside the work’s scope.ReRAM is selected for its relative maturity and demonstrated read characteristics.
  • eNVM Modeling Methodology: The study evaluates 8-bit quantized embeddings using 100 fault-injection trials per ReRAM storage configuration.The experiments quantify the trade-off between storage density and task accuracy using 28nm ReRAM cell characteristics.
  • Sparse encoding: Sparse embedding weights use bitmask-style encoding, with the bitmask stored in lower-risk SLC devices and non-zero data values tested in MLC storage.This separation addresses the reported vulnerability of DNN weight bitmasks to MLC faults.
  • Storage selection: MLC2 ReRAM does not degrade accuracy across multiple tasks in the reported fault-injection trials, whereas MLC3 causes potentially catastrophic minimum-accuracy degradation and lowers average QNLI accuracy.EdgeBERT therefore selects MLC2 ReRAM for word-embedding storage.

5. EDGEBERT’S LATENCY-AWARE INFERENCE

EdgeBERT predicts an early exit layer from first-layer entropy, then uses sentence-level DVFS to reduce energy while meeting a per-sentence latency target. The predictor is implemented as a lookup table and inference still checks entropy at each layer before the predicted exit.

  • Latency-aware inference: Sentence-level DVFS uses the predicted exit layer to minimize per-sentence energy while satisfying a real-time latency target.The scheme combines the early-exit predictor with dynamic voltage and frequency scaling.
  • Early-exit prediction: The EE predictor is a ReLU-activated five-layer perceptron with 64 hidden cells per layer, trained from entropy values across all 12 Transformer layers.Its input is the entropy at encoder layer 1, and its output forecasts the early-exit layer.
  • Early-exit prediction: The predictor is distilled into a lookup table, while entropy remains checked after each encoder layer to stop earlier than the predicted layer when the threshold is reached.This preserves early stopping while limiting the predictor’s per-sentence computational overhead.
  • Evaluation: Entropy prediction is evaluated against conventional early exit at fixed accuracy-degradation thresholds of 1%, 2%, and 5%.The comparison increases the entropy threshold until the specified relative accuracy drop is reached.
  • Sentence-level DVFS: The optimal frequency is computed from remaining encoder cycles and the remaining latency budget, after accounting for elapsed frontend time.The DVFS controller selects the corresponding energy-optimal supply voltage from a voltage/frequency lookup table.

6. ALGORITHMIC SYNERGY

EdgeBERT combines pruning, quantization, adaptive attention spans, and entropy prediction in a multi-step ALBERT training and evaluation procedure. The resulting combination maintains comparable average exit layers while enabling latency-aware DVFS and compact parameter storage.

  • Training and evaluation: The training procedure applies pruning during fine-tuning, knowledge distillation from a task-tuned ALBERT teacher, and separate pruning schedules for embeddings and encoder layers.Movement or magnitude pruning is applied to encoder layers, while magnitude pruning is applied to embeddings.
  • Training and evaluation: Evaluation applies 8-bit floating-point quantization, models embeddings in a 2-bit-per-cell MLC2 ReRAM configuration, and deploys adaptive attention masking with entropy-based early exit.The learned attention span mask is multiplied element-wise with attention weights.
  • Synergistic inference: EdgeBERT provides a comparable average exit layer to conventional early exit at the same accuracy threshold while allowing DVFS to adapt voltage and frequency to the predicted exit layer.This links the early-exit prediction to latency-aware energy control.
  • Synergistic inference: The entropy predictor requires a lower entropy threshold than conventional early exit for the same accuracy target because its neural predictor can introduce classification error.The lower threshold reflects conservative prediction.
  • Memory optimization: Across four corpora, embedding density reaches 40%, producing a 1.73MB compact memory baseline for eNVM storage.The result is reported as a uniform embedding-layer density across the evaluated corpora.

7. THE EDGEBERT HARDWARE ACCELERATOR SYSTEM

The EdgeBERT accelerator co-designs specialized datapaths, compressed sparse execution, non-volatile parameter storage, and fast on-demand DVFS. Its processing and special-function units support the Transformer operations required for latency-aware inference.

  • System architecture: The accelerator is designed to exploit EdgeBERT’s latency and memory optimizations for compute and energy efficiency.Its highlighted features include early-exit assessment, attention masking, layer normalization, non-volatile storage, DVFS, and sparse execution.
  • Processing unit: The processing unit executes matrix multiplications for ALBERT linear layers and attention heads using parallel floating-point vector MACs.The datapath accepts two n×n matrices and computes n×n×n MAC operations in n clock cycles.
  • Sparse execution: Bit-mask encoding stores binary zero/non-zero tags so sparse matrices can be compressed, decoded, and processed while skipping zero entries.The decoder reconstructs zero values and identifies input banks that can be neglected.
  • Special function unit: The special function unit computes early-exit assessment, DVFS control, element-wise addition, layer normalization, and softmax, with auxiliary storage for LUTs and attention masks.Its softmax unit applies the learned attention span mask after numerically stable softmax computation.
  • Attention span masking: Null attention-span masks allow the controller to cancel an attention head’s computation and write a zero context vector directly.This mechanism skips the full sequence of operations for heads whose learned mask is null.
  • Early exit and DVFS: The EE assessment unit compares entropy with the threshold and, in latency-aware mode, passes the predicted exit layer to the DVFS controller.The controller uses the prediction to retrieve voltage and frequency settings from a lookup table.
  • Early exit and DVFS: The integrated LDO scales voltage from 0.5V to 0.8V in 25mV steps, while the ADPLL relocks frequency quickly after an update request.These components implement the accelerator’s latency-driven DVFS system.

8. HARDWARE EVALUATION

Hardware evaluation identifies an energy-optimal EdgeBERT configuration and shows how adaptive attention, sparse execution, DVFS, and eNVMs reduce latency and energy under latency constraints. The accelerator achieves substantial energy savings against both conventional inference and a mobile GPU.

  • Design Space Exploration via MAC scaling: A MAC vector size of 16 is energy-optimal because larger datapath power begins to offset throughput gains.Per-sentence latency decreases by roughly 3.5× as vector size doubles.
  • Design Space Exploration via MAC scaling: Adaptive attention span reduces accelerator processing time by up to 1.2×, while compressed sparse execution adds 1.4–1.7× energy savings.QQP receives the greatest benefit from compressed sparse execution.
  • Design Space Exploration via MAC scaling: 53× lower energy is achieved than the Nvidia TX2 mobile GPU when all optimizations are combined.The comparison uses the energy-optimal accelerator design with MAC vector size 16.
  • Hardware Implementation: The placed-and-routed energy-optimal accelerator occupies 1.4mm2 and consumes 86mW average power.The breakdown covers latency, energy, area, and power contributions in the 12nm system.
  • DVFS-based Latency-Aware Inference: 7× and 2.5× per-inference energy savings are achieved versus conventional inference and latency-unbounded early exit, respectively.The results use latency targets between 50ms and 100ms; at some targets, DVFS reaches its scaling floor.
  • DVFS-based Latency-Aware Inference: At a 75ms MNLI target, early exit, adaptive attention span, DVFS, sparse execution, and eNVMs contribute 21%, 12%, 23%, 39%, and 5% of total energy reduction.These contributions quantify the combined sources of accelerator energy savings.
  • Non-Volatile Memory Evaluation: Static eNVM storage of shared embeddings avoids repeated DRAM reads and SRAM transfers after power-on.The configuration stores embedding bitmasks in SLC ReRAM and nonzero parameters in MLC2 ReRAM.

9. RELATED WORK

Related work reduces Transformer computation or storage through attention approximation, quantization, and accelerator design. EdgeBERT differs by learning sparse attention spans during fine-tuning and statically storing shared embeddings in eNVMs.

  • Prior Transformer Acceleration: Prior Transformer accelerators include attention candidate search, quantization, and dense systolic-array designs, but most prior accelerator work focused on CNNs and RNNs.These approaches target computation or memory reduction through different hardware and numerical strategies.
  • Attention Optimization: EdgeBERT learns the optimal attention search radius during fine-tuning and uses a sparse mask to avoid unnecessary computation and data movement.A3 instead fetches full uncompressed data from DRAM before dynamically reducing attention computation.
  • Quantization and Storage: EdgeBERT uses uniform 8-bit storage in SRAM and eNVMs, contrasting with GOBO’s mixed-precision datapath and non-uniform memory storage.GOBO applies 3-bit clustering to most weights while retaining outliers and activations in FP32.
  • Non-Volatile Embedding Storage: EdgeBERT statically stores task-shared word embeddings in high-density eNVMs rather than traditional volatile SRAM.The paper identifies this storage choice as a source of energy and latency benefits.

10. CONCLUSION

EdgeBERT addresses the latency, memory, and energy costs of Transformer inference on resource-constrained edge platforms through algorithm–hardware co-design. Its sentence-level latency-aware accelerator reduces energy relative to conventional inference and latency-unbounded early exit.

  • Conclusion: EdgeBERT combines first-layer early exit prediction with sentence-level DVFS to minimize energy while meeting a prescribed latency target.Adaptive attention span, selective pruning, floating-point quantization, and eNVM storage further reduce latency and memory footprint.
  • Conclusion: 7× and 2.5× lower energy are reported than conventional full-model inference and latency-unbounded early exit, respectively.The conclusion attributes these savings to the integrated algorithmic and hardware methodology.
Loading 2011.14203v5…