AxLaM: energy-efficient accelerator design for language models for edge computing

(a) Background—BERT: use cases and internal structure

BERT [1] is a pre-trained language model introduced by Google in 2018. Unlike traditional CNN models, BERT uses bidirectional training, leveraging both preceding and succeeding contexts to predict words, which results in a deeper understanding of language. This approach has led to state-of-the-art performance across various NLP tasks. BERT excels in tasks such as text classification, named entity recognition, question answering, sentiment analysis and language translation [16]. Its pre-trained nature allows for fine-tuning, making it versatile and well-suited for edge applications.

BERT is based on the Transformer architecture, consisting of an input embedding layer, multiple Transformer encoder blocks and an output layer. Each encoder block includes a multi-head self-attention mechanism and a feed-forward neural network. Tokens are embedded with positional encodings and processed through Transformer blocks. The BERT-large model has specific parameters: a model dimensionality (${\displaystyle d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}}$) of 1024, 24 Transformer blocks ($N$), each with 16 attention heads ($H$). The key/query ($d_k$) and value ($d_v$) vectors have dimensions ${\displaystyle d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}/H=64}$. Each attention head computes a 1024 × 1024 matrix, with concatenated outputs passed through a linear transformation to produce the final layer output:

where $W_O$ is the weight matrix that maps the concatenated output to the embedding size, ${\displaystyle d_{\mathrm{m}\mathrm{o}\mathrm{d}\mathrm{e}\mathrm{l}}=1024}$. After multi-head attention, the output undergoes a residual connection and layer normalization and is processed through a feed-forward neural network, which first increases and then reduces dimensionality. The matrix operations are listed in table 1. These operations enhance computational efficiency and consistency. BERT’s effect on NLP tasks has made it a leading model, with scalable configurations such as BERT-base and BERT-large suitable for various environments. BERT’s performance is recognized in MLPerf [16], making it the baseline workload for our design.

(b) Background—spatial and temporal reuse in hardware accelerator

While efforts exist to accelerate BERT [17,18], no accelerators have been specifically optimized for edge scenarios, as noted in [2,3]. In edge computing, the Simba architecture is a leading solution for deep neural network inference [8,19], thanks to its advanced vector multipliers and accumulators that share input activations while outputting to multiple channels. Notably, Simba features a data-agnostic design, making it robust against data distribution variations.

At the core of Simba, illustrated in figure 1, is the PE, responsible for executing convolutional layers, fully connected layers and post-processing operations such as bias addition, ReLU and pooling. Each PE contains eight lanes, each using a distinct weight tensor to generate elements for a single output channel (K). An 8-bit precision vector MAC unit within each lane multiplies eight input elements from different channels (C) with eight weight elements, summing them to produce a single output. This approach is highly efficient, given typical channel counts (64–1024). The architecture is supported by local input activation, output activation and accumulation SRAMs, which buffer the datapath. Energy efficiency is optimized by minimizing SRAM access: input activation SRAM is read every cycle, and its elements are forwarded across lanes. The weight SRAM, wider than the input activation SRAM, supplies distinct vectors to each lane and reuses weights across multiple inputs, reusing them P $\times$ Q times (where P and Q are output dimensions). Accumulation SRAM stores intermediate sums, conserving energy by accumulating the eight-wide vector from C channels. The output size, often exceeding the accumulation buffer’s capacity, requires temporal tiling to generate portions of output activations sequentially. This buffer also supports cross-PE reductions when the weight kernel spans multiple PEs. Post-processing tasks such as ReLU, bias addition and pooling are performed after accumulation. The buffer is dual-banked, allowing simultaneous access by MACs, routers and post-processing units, with an arbitration crossbar to resolve bank conflicts.

• Download figure
• Open in new tab
• Download PowerPoint

(c) Motivation for a new architecture design

The growing complexity and computational demands of NLP, particularly models such as BERT, have exposed the limitations of existing hardware accelerators for machine learning (ML). Accelerators for BERT, which are primarily based on FPGA platforms [18] or the more recent ASIC designs [2], often require significant modifications to the neural network architecture to achieve operational efficiency. Additionally, such modifications, including extreme quantization or mixed-signal processing techniques, frequently result in a notable loss of accuracy [2,13,20]. While the Simba architecture, owing to silicon realization and data-agnostic operation, has been recognized as a leading solution for deep neural network inference, its efficacy diminishes when confronted with the distinct computational demands of BERT’s matrix operations. The various matrix multiplications, particularly those involving the QKV (Query, Key, Value) operations, impose a substantial power burden, primarily owing to the significant DRAM accesses required, as illustrated in figure 2. These multiplications, as listed in table 1, highlight the computational challenges that Simba faces when deployed for models such as BERT.

• Download figure
• Open in new tab
• Download PowerPoint

The QKV operations, though not the largest in magnitude, play a critical role in BERT’s performance owing to relatively frequent computation (16 times each per encoder block). The power consumption associated with these operations is exacerbated by the limited data reuse inherent in their matrix multiplications. In energy-efficient computing, small buffer sizes are often preferred to optimize the power consumption of such operations. However, as the buffer sizes increase, the energy required for data access also rises, leading to further inefficiencies. Conversely, operations such as Bottleneck Contraction (BtlCont), Bottleneck Expansion (BtlExp) and Multi-Head Attention (MHead) exhibit a more pronounced data reuse pattern. Each output word in these operations depends on approximately 1024 input words, allowing for a significant reduction in access counts as buffer sizes grow. This characteristic makes large buffer sizes essential to fully exploit data reuse, thereby improving power efficiency in these layers.

Despite these insights, the Simba architecture’s current design is not well-suited to meet the diverse computational requirements of BERT’s matrix operations. Simba’s average power consumption, 9 and 19 W, respectively, for INT8 and FP16 pipelines, is far beyond the acceptable range for edge and portable applications, which typically impose stringent power constraints, often below 5 W [21]. Therefore, while Simba excels in certain scenarios, its inability to efficiently handle the unique demands of BERT’s matrix multiplications underscores the need for a new architecture that can better address these challenges.

(d) Design consideration: challenges to be addressed

Efficient hardware accelerators for neural networks such as BERT must address several key challenges: high DRAM access count, on-chip buffer energy consumption and on-chip compute energy efficiency.

1. Reducing high DRAM access count and energy consumption: high DRAM access significantly effects overall energy consumption. To mitigate this, increasing the on-chip buffer size can reduce the frequency of DRAM accesses by enhancing data reuse and minimizing energy-intensive memory transactions. This approach leverages larger buffers to store more data locally, thereby reducing the need for frequent access to external memory.

2. Reducing energy at the on-chip buffer level: while larger on-chip buffers can reduce DRAM access, they can also lead to increased energy consumption within the chip. This challenge can be managed by optimizing memory hierarchies with multi-level buffering strategies and exploiting opportunities for temporal and data sharing within the application. Effective buffer management ensures that the energy benefits of reduced DRAM access are not offset by increased internal energy costs.

3. The energy required for on-chip computations, particularly MAC operations, is a significant consideration for BERT’s intensive computations. Reducing this energy footprint can be achieved by employing low-power arithmetic units with reduced precision and by implementing operand-sharing techniques to maximize the reuse of input data. These strategies help lower the energy consumption associated with each computation, making the overall processing more efficient.

(e) Design consideration: on-chip buffer size and memory accesses

The unique matrix multiplications within BERT’s encoder exhibit distinct DRAM access patterns that significantly effect the power consumption associated with different buffer sizes. The primary objective is to select a buffer size that minimizes the overall system energy consumption (considering on-chip data reuse) during data accesses and computational processes. As shown in figure 3a , the DRAM access count varies with increasing total buffer size, offering valuable insights into data reuse and how frequently each matrix multiplication operation requires access to main memory. To complement this, figure 3b illustrates the energy per byte access data for different buffer sizes, highlighting the energy implications of choosing specific buffer configurations. Balancing the frequency of DRAM accesses against the energy overhead of larger buffers is critical in determining the most energy-efficient buffer size. To achieve optimal energy efficiency, both the frequency of DRAM accesses for each matrix multiplication and the energy cost associated with each access—given a specific buffer size—must be considered. Additionally, the repetitive nature of BERT’s matrix operations within the encoder should be taken into account, as certain operations recur more frequently, potentially skewing the overall energy profile. For instance, while QKV operations may require fewer DRAM accesses in some buffer configurations, their central role in the encoder and frequent occurrence contribute significantly to the total energy consumption.

• Download figure
• Open in new tab
• Download PowerPoint

Comparing the DRAM access count data with buffer access energy, it becomes evident that smaller buffer sizes generally exhibit lower energy costs per access. However, as buffer sizes increase, while DRAM access counts typically decrease, the energy cost per access rises. Considering all available metrics, an intermediate buffer size appears to strike the best balance between reducing DRAM accesses and maintaining manageable energy costs per access. Analysis suggests that buffer sizes in the range of 64–128 KB may be appropriate, as indicated by the trends in the provided data. However, further refinement of this range is necessary, particularly by considering the specific frequency of each matrix operation in BERT’s encoder.

(f) Opportunity: HBM logic layer for bandwidth and low access energy

Traditional DDR DRAM (such as DDR4 DRAM) is insufficient for handling the high frequency memory accesses required by BERT’s encoders, where DRAM access energy is a major concern [22]. HBM3 [6] offers a tenfold reduction in access energy per bit (4.2 versus 46 pJ), making it a superior choice for BERT’s frequent matrix multiplications, significantly enhancing accelerator efficiency. HBM3’s design is well-suited for BERT-like models, allowing direct integration of accelerator designs within its logic layer (up to 30 mm² and 8.5 W power [23]), as supported by recent manufacturer initiatives [24]. This integration minimizes data movement, reducing latencies and further cutting energy consumption, crucial for the high data reuse in BERT’s encoders. HBM3’s structure stacks DRAM dies atop a logic layer, connected via through-silicon vias, achieving a compact memory footprint. It vastly outperforms DDR4 in bandwidth, offering approximately 512 GB/s per stack compared to DDR4’s 25.6 GB/s, making HBM3 ideal for high-throughput, rapid-access applications [6]. However, this presents three main challenges: ensuring balanced access to all HBM memory channels, managing design area constraints and addressing thermal and power limitations.

(g) Addressing the high MAC energy

In accelerators for complex language models such as BERT, MAC operations are the second highest contributors to power consumption after DRAM access, as shown in figure 2. Traditional 8-bit integer operations, while efficient, do not meet the precision demands of BERT, necessitating a solution that combines floating-point precision with reduced power overheads (figure 4). Recent advancements in multiplier design, such as the POSIT number system introduced by Gustafson & Yonemoto [5], offer a promising alternative. POSITs, defined by parameters $(N,es)$, provide dynamic flexibility, allowing efficient space usage by encoding exponent and fraction parts only when necessary. However, the increased area and power requirements of POSITs present challenges. To mitigate these, the fixed-POSIT representation was developed [25], which standardizes the regime length ($rf$) alongside $N$ and $es$. The AFPOS multiplier further refines this approach, optimizing for area, energy and latency while requiring 25% less bit storage than fixed-POSITs [22]. Our experimentation has shown that fixing the $2^{{}^{\beta }}$ term for AFPOS (with 1-bit sign, 4-bit exponent and 3-bit mantissa) to $2^{-7}$ results in negligible loss of accuracy compared to FP32, as shown in figure 5. For the rest of the paper, AFPOS refers to this ($N=10,es=4,\beta =-7$) configuration. Table 2 lists the latency, area and energy of these multipliers. We have highlighted the designs with the highest quantization yet having similar accuracy and Matthews correlation coefficient (MCC) score as FP32. Compared to a quantized 16-bit floating-point multiplier, AFPOS reduces area by a factor of 4.04, energy by a factor of 8.9 and latency by a factor of 2.1. We use a 16-wide vector MAC, with area, latency and power as listed in table 3, for efficient spatial reduction.

• Download figure
• Open in new tab
• Download PowerPoint

• Download figure
• Open in new tab
• Download PowerPoint

(h) PE design details and innovations

BERT’s computational architecture, particularly its vector multipliers, poses unique challenges distinct from frameworks such as Simba. In Simba, shared input activations among a set of vector multipliers, along with exclusive weights, allow weights to be read once and reused multiple times. This resource optimization, however, is not inherent in BERT. When dealing with matrices $L$ and $R$ (as seen in table 1), if the elements of $R$ are temporarily stationary, those of $L$ require constant cycling and vice versa. This dynamic underscores the need to efficiently amortize read energy costs, leading to a redesigned vector MAC and PE, as shown in figure 6. Instead of employing 16 separate PEs, we propose a unified PE that optimizes shared resource utilization while maintaining an acceptable fan-out. For each read of an element in $L$, the element contributes to the partial product of eight distinct columns, while efficient reads from $R$ enable participation in eight different rows of the output matrix. A simplified representation of this concept is illustrated in figure 7, where words read together are annotated with alphanumeric labels and sources are colour-coded equally compared to figure 6. Furthermore, within a column or row, 16 elements are multiplied together and accumulated using the vector MAC and written to the accumulation buffer (A.SRAM). We choose this configuration, as the baseline Simba instance, in efficiency mode, offers 1 TOPS throughput (which offers sub-1 s inference of BERT-large). Our goal is to achieve a similar throughput with our vector MAC configuration, with the parallel adder tree. The operational strategy dictates that either the $L$ or $R$ matrix elements must be refreshed upon completing each operand set. To balance energy efficiency and performance, our design incorporates an 8 KB buffer for all $L$ buffers and an equivalent capacity for the $R$ buffers. The A.SRAM is designed to cache partial products, enabling efficient summation.

• Download figure
• Open in new tab
• Download PowerPoint

• Download figure
• Open in new tab
• Download PowerPoint

(a) Evaluation infrastructure

At the software level, we model BERT using a PyTorch-based GPU accelerated framework, modified to support different number representations, including AFPOS. In the framework, the flow of data is altered such that the data are encoded into the specified encoding scheme (such as BFloat, POSIT or AFPOS) before multiplication. This allows the framework to use the IEEE 754 Float-based pre-trained BERT model without alteration. After fine-tuning the model with the CoLA dataset, we measure the accuracy and MCC scores against the full-precision model (result shown in figure 5). We use this software framework for reproducibility and for exact computation and score calculation.

For hardware modelling, we use Cadence Genus to synthesize the MAC units (of all encoding schemes) at a 65 nm node, ensuring timing constraints are met at 500 MHz frequency. SRAM buffers are modelled using CACTI [26]. System-level modelling and optimal data search (for least latency followed by least energy) are performed with Timeloop, Accelergy and Alladin framework [7,27], defining both Simba and our proposed architecture based on synthesized values. For Simba, the external DRAM memory used is DDR4 DRAM, and the hardware model is identical to the hardware-realized version of Simba [8]. Our proposed work is modelled as a design instantiated on the logic layer of the HBM3 memory. We use this hardware model framework as it has been validated against real hardware and allows for flexible extensions.

(b) Workload description

We focus on the unique matrix operations within BERT’s encoder, as listed in table 1. These operations include the QKV computations, attention mechanisms, multi-head concatenations and feed-forward networks. The network level performance of any variation of BERT-large can be extrapolated from these matrices.

(a) Latency and energy

Figure 8 compares the latency during the processing of each unique matrix by Simba and the proposed design. Although both designs feature an identical RAW throughput of 1 TOPS, Simba’s reduced data reuse, particularly in the QKV and attention matrices, limits its ability to fully utilize all MAC units owing to bandwidth constraints. This inefficiency results in a 1.2 times speedup for the proposed design. Each encoder is processed within 27 ms, enabling the complete BERT-large model to be processed in 0.6 s, well within the acceptable latency threshold for handheld/edge devices performing natural language tasks, even when processing up to 1000 tokens.

• Download figure
• Open in new tab
• Download PowerPoint

Figure 9 shows the amortized energy cost of a MAC operation at the system level, comparing Simba, a variant of Simba (Simba with AFPOS multipliers, Simba implemented in HBM3’s logic layer and Simba implemented with AFPOS multiplier in HBM3’s logic layer) and the proposed designs implemented with HBM3 and the candidate multipliers. The average of each implementation in figure 9 represent weighted average operations corresponding to an encoder in BERT-large. The AFPOS-adapted version provides insight into the effect of changing the number system alone. The proposed design (AxLaM (AFPOS)) demonstrates the lowest energy per MAC operation with a ninefold reduction in energy consumption, primarily owing to decreased DRAM access and multiplication energy compared to Simba. By only altering the multiplier in Simba, the energy difference narrows to sevenfold for the proposed design. Additionally, the internal buffer energy is halved in the proposed design, which is attributed to the improved PE organization.

• Download figure
• Open in new tab
• Download PowerPoint

(b) Area and power constraints comparison: Simba versus proposed design

Integrating the proposed design within the available area of the HBM3 memory’s logic die offers substantial benefits in bandwidth and energy efficiency. Figure 10 illustrates the area distribution for the various components of the design. Compared to the Simba design, the proposed design not only adheres to the 36 mm² area constraint of HBM3 but also significantly reduces it to just 2.5 mm². This efficiency is primarily due to the streamlined single PE design of the proposed system, as opposed to Simba’s more expansive 16 PE architecture, resulting in an impressive 58% reduction in area.

• Download figure
• Open in new tab
• Download PowerPoint

In terms of power consumption, the proposed design shows marked improvement, consuming only 1.103 W, which is a ninefold reduction compared to Simba’s 9.34 W (INT8). This significant decrease ensures that the proposed design comfortably fits within HBM3’s 8.5 W power limit [6]. Power consumption is calculated based on the energy and time required to process one full encoder of BERT.

(c) Comparison with other recent accelerators

Table 4 lists a comparison between our proposed design, AxLaM and recent accelerators such as OPTIMUS, A³, Transformer Engine, SwiftTron and others. Energy efficiency has been normalized to a 65 nm node based on voltage and capacitance scaling models. AxLaM distinguishes itself through the use of the AFPOS system, achieving a throughput of 1 TOPS with a normalized energy efficiency of 1.3 TOPS/W. This positions AxLaM as one of the most energy-efficient designs, particularly for BERT, all within a compact 2.5 mm² area on a 65 nm process node.

While other accelerators, such as the Transformer Engine, offer higher raw throughput, they fall short in terms of energy efficiency. AxLaM also demonstrates superior area efficiency, especially when compared to designs such as Swiftron and H3D, which require significantly more space. AxLaM effectively balances throughput, energy efficiency and area, making it a highly competitive solution for edge devices focused on NLP tasks.

In terms of energy efficiency, H3D [13] and FACT [20] are the closest competitors. However, FACT requires modifications to the language model. The actual benefit comes from algorithm changes applied to BERT before porting it on the hardware. On the other hand, H3D, being based on analogue computing, suffers from high variability and is not suitable for mass manufacturability. This further emphasizes AxLaM’s robustness and suitability for practical deployment.