DAPA: Distribution Aware Piecewise Activation Functions for On-Device Transformer Inference and Training

Maoyang Xiang; Bo Wang

cs.LGMar 19, 2026

DAPA: Distribution Aware Piecewise Activation Functions for On-Device Transformer Inference and Training

Authors: Maoyang Xiang, Bo Wang

Organizations: Information Systems Technology and Design Singapore University of Technology and Design Singapore

Abstract

Non-linear activation functions play a pivotal role in on-device inference and training, as they not only consume substantial hardware resources but also impose a significant impact on system performance and energy efficiency. In this work, we propose Distribution-Aware Piecewise Activation (DAPA), a differentiable and hardware-friendly activation function for Transformer architectures by exploiting the distribution of pre-activation data. DAPA employs a non-uniform piecewise approximation that allocates finer segments to high-probability regions of the distribution, improving generalizability over prior piecewise linear methods. The resulting approximation is further quantized using Distribution-Weighted Mean Square Error to reduce latency and resource utilization for hardware deployment. Our HLS implementation demonstrates that DAPA speeds up GELU computation by 16 $\times$ and decreases DSP utilization by 16 $\times$ while maintaining comparable or better performance across vision Transformers and GPT-2 models.

Figures & tables

Figure 1: The relationship between MSE/DWMSE and the difference in Top-1 ImageNet-1K classification accuracy relative to the FP32 baseline, evaluated across three Vision Transformer variants.

Figure 2: The relationship between MSE/DWMSE and the difference in WikiText2 PPL relative to the FP32 baseline, evaluated on GPT2.

Model	Metric	Correlation Coefficients			Fisher CI
Model	Metric	r	$\rho$	$\tau$	Fisher CI
ViT-Base	MSE	0.191	0.359	0.312	[-0.43, 0.69]
ViT-Base	DWMSE	0.646	0.754	0.562	[0.12, 0.89]
ViT-Small	MSE	0.323	0.169	0.250	[-0.31, 0.76]
ViT-Small	DWMSE	0.938	0.923	0.812	[0.79, 0.98]
ViT-Tiny	MSE	0.490	0.528	0.438	[-0.12, 0.83]

Table 1: Comparison of correlation coefficients and 95% Fisher CIs.

Figure 3: The DAPA method approximates the GELU function and its derivative by partitioning the probability density function (PDF) into

N

quantile regions.

Figure 4: DAPA functions generated from distributions with different numbers of images are evaluated on the ViT-Small.

Model	Torch, a F	b M, c A, F	d D, A, F	D, e S, F	D, S&A, F	D, S&A, f Q	[ 11 ]	[ 14 ]	[ 20 ]	[ 13 ]
ViT-Tiny	75.51%	73.82%	75.51%	75.44%	75.42%	75.17% (Q9.7) g	-	-	-	-
ViT-Small	81.40%	80.81%	81.41%	81.37%	81.40%	81.30% (Q8.8)	-	-	-	-
ViT-Base	81.66%	81.35%	81.69%	81.69%	81.70%	81.61% (Q7.7)	-	-	-	-
DeiT-Tiny	74.53%	73.92%	74.55%	74.47%	74.43%	74.37% (Q6.9)	-	-	71.47%	-
DeiT-Small	81.17%	80.85%	81.11%	81.21%	81.16%	81.11% (Q6.8)	79.36%	79.11%	79.44%	80.34%
DeiT-Base	83.32%	83.13%	82.90%	83.28%	82.88%	82.97% (Q7.5)	-	-	82.92%	-

Table 2: Performance of Vision and NLP Transformer models with different approximation methods

Method	SST-2	MRPC (F1)	STS-B	QQP	MNLI-m	MNLI-mm	QNLI	Average
Baseline in [ 3 ]	91.51	88.01	88.32	90.77	83.83	84.43	90.88	88.24
[ 3 ]	91.40	88.05	88.27	90.76	83.77	84.33	90.88	88.21
[ 24 ]	91.86	88.22	88.30	90.62	83.82	84.17	90.68	88.24
Baseline in [ 20 ]	92.08	85.66	-	91.35	84.23	-	91.49	88.96
[ 20 ]	91.96	86.27	-	91.12	84.10	-	91.23	88.94
PyTorch, F	92.43	91.35	88.05	90.91	84.57	84.45	91.54	89.05

Table 7

Figure 5: Training loss curves for three ViT variants comparing standard PyTorch GELU activation with the DAPA(16) approximation function across epochs.

	ViT-Base	ViT-Small	ViT-Tiny
GELU (PyTorch)	69.76%	67.70%	66.27%
DAPA(16)	69.90%	68.35%	66.47%

Table 3: Classification results on ImageNet-1K: ViT models trained with GELU and DAPA activation functions

Figure 6: Hardware architecture of the proposed DAPA with $N{=}8$ segments with four-stage pipeline. Stages I–III realize a pipelined comparator tree and encoder for segment identification, while stage IV performs the final linear computation with a single MAC unit using the selected $(a_{n},b_{n})$ .

	DType a	L(ns) b	B c	DSP	FF	LUT
Gelu (Ours)	FP32	580	5	62	12158	7228
Gelu (Ours)	Fix16	320	-	8	5858	6302
D(16) (Ours) d	FP32	150	2	7	1304	1100
D(16) (Ours) d	Fix16	20	-	1	100	401
DS(16) (Ours) e	Fix16	155	-	1	2243	1614
Gelu [ 11 ]	Fix16	-	-	16	2951	2940

Table 4: Comparison of FPGA Hardware Implementations for Activation Functions.

Neural Network Activation Functions

Neural network activation functions shape how models represent nonlinear patterns, learn, remain stable, and map efficiently to hardware.

See all 78 papers

TRACE: Transformer Residual Activation Codebook Encoding for Communication-efficient Distributed Edge Inference
Oct 15, 2026Muhammad Azlan Qazi, Alexandros Iosifidis, Qi ZhangEfficient Transformer Inference Transformer Inference
DIGIT and the Department of Electrical and Computer Engineering, Aarhus University, Denmark · the Data Science Research center, Tampere University, Finland
Profiling-Driven Adaptive Distributed Transformer Inference on Embedded Edge Deployment
May 25, 2026Muhammad Azlan Qazi, Alexandros Iosifidis, Qi ZhangEfficient Transformer Inference Transformer Inference
Aarhus University · Aarhus, Denmark · Tampere University +1
Recent Developments in Transformer Inference Deployment on FPGA Platforms: A Survey
Sep 1, 2026Arjan Blankestijn, Uraz Odyurt, Amirreza YousefzadehEfficient Transformer Inference Transformer Inference
Computer Architecture for Embedded Systems, University of Twente Enschede, The Netherlands · Faculty of Engineering Technology, University of Twente Enschede, The Netherlands

Oct 15, 2026J/K move · Enter open · S save

TRACE: Transformer Residual Activation Codebook Encoding for Communication-efficient Distributed Edge Inference

Muhammad Azlan Qazi, Alexandros Iosifidis, Qi Zhang

DIGIT and the Department of Electrical and Computer Engineering, Aarhus University, Denmark · the Data Science Research center, Tampere University, Finland

Efficient Transformer Inference Transformer Inference Edge Inference Vector Quantization

Open paper AskPDF arXiv

Distributed Transformer inference across co-located edge devices is fundamentally communication-bound, as exchanging activations between sequence-parallel partitions dominates inference latency under commodity edge bandwidths. Existing compression schemes reduce this cost but retain a structural bottleneck: the received data must still undergo key/value projection before attention can proceed, serializing computation behind communication. This paper presents Transformer Residual Activation Codebook Encoding (Trace), a framework that addresses these limitations. Trace places a Vector Quantizer (VQ) at the post-attention residual of every Transformer layer and precomputes codebook-indexed key/value representations. During inference, exchanged activations are encoded into compact codebook indices, allowing receivers to directly recover keys and values through table lookup instead of executing the downstream computations required to derive them. By eliminating the received data downstream computation, Trace enables communication--computation overlap and reduces a key latency bottleneck in distributed Transformer inference. A hierarchical clustering further structures the learned codebook into a multi-resolution hierarchy, enabling runtime adaptation of the communication bitrate without retraining. The framework supports two operating variants: one prioritizing accuracy and the other balancing accuracy and efficiency through buffer-compatible quantization. Experiments on ViT-Base (ImageNet-1K and CIFAR-100) and ViViT-B (UCF-101) show that Trace achieves up to 1.89% higher accuracy and up to 1.27x end-to-end inference speedup compared with the state-of-the-art VQ-based distributed inference framework. Compared with uncompressed sequence-parallel inference, Trace reduces communication overhead by up to 22,624x.

Abstract

Neural Network Activation Functions

Explore similar work

TRACE: Transformer Residual Activation Codebook Encoding for Communication-efficient Distributed Edge Inference

Profiling-Driven Adaptive Distributed Transformer Inference on Embedded Edge Deployment

Recent Developments in Transformer Inference Deployment on FPGA Platforms: A Survey