The Larger the Merrier? Efficient Large AI Model Inference in Wireless Edge Networks

TL;DR

This work tackles efficient LAIM inference in wireless edge networks by pruning and splitting a pre-trained LAIM into on-device and on-server sub-models, reducing computation and communication. It establishes that LAIM output distortion is upper-bounded by parameter distortion and derives a rate-distortion lower bound for pruning, enabling a distortion-aware design. It then formulates a joint optimization of pruning ratios, transmit power, and computation frequencies under latency and energy constraints, solved via a successive convex approximation algorithm. Simulations show the proposed design achieves better inference quality-latency-energy trade-offs than fully on-device/on-server benchmarks and highlights the split point’s critical role in heterogeneous edge environments.

Abstract

The growing demand for large artificial intelligence model (LAIM) services is driving a paradigm shift from traditional cloud-based inference to edge-based inference for low-latency, privacy-preserving applications. In particular, edge-device co-inference, which partitions LAIMs between edge devices and servers, has emerged as a promising strategy for resource-efficient LAIM execution in wireless networks. In this paper, we investigate a pruning-aware LAIM co-inference scheme, where a pre-trained LAIM is pruned and partitioned into on-device and on-server sub-models for deployment. For analysis, we first prove that the LAIM output distortion is upper bounded by its parameter distortion. Then, we derive a lower bound on parameter distortion via rate-distortion theory, analytically capturing the relationship between pruning ratio and co-inference performance. Next, based on the analytical results, we formulate an LAIM co-inference distortion bound minimization problem by jointly optimizing the pruning ratio, transmit power, and computation frequency under system latency, energy, and available resource constraints. Moreover, we propose an efficient algorithm to tackle the considered highly non-convex problem. Finally, extensive simulations demonstrate the effectiveness of the proposed design. In particular, model parameter distortion is shown to provide a reliable bound on output distortion. Also, the proposed joint pruning ratio and resource management design achieves superior performance in balancing trade-offs among inference performance, system latency, and energy consumption compared with benchmark schemes, such as fully on-device and on-server inference. Moreover, the split point is shown to play a critical role in system performance optimization under heterogeneous and resource-limited edge environments.

Figures (6)

• Figure 1: Illustration of the considered LAIM co-inference system.
• Figure 2: Distribution of weights of various pre-trained models.
• Figure 3: Distortions of model outputs and parameters w.r.t. pruning ratios.
• Figure 4: System performance w.r.t. different delay thresholds. (a) Pruning BART with the magnitude paradigm under the inference energy consumption threshold $E_0 =0.0225 ~ {\rm J}$. (b) Pruning BART with the random paradigm under the inference energy consumption threshold $E_0 =0.0225 ~ {\rm J}$. (c) Pruning BERT with the magnitude paradigm under the inference energy consumption threshold $E_0 =0.0150 ~ {\rm J}$. (d) Pruning BERT with the random paradigm under the inference energy consumption threshold $E_0 =0.0011 ~ {\rm J}$.
• Figure 5: System performance w.r.t. different energy consumption thresholds. (a) Pruning BART with the magnitude paradigm under the inference delay threshold $T_0 =0.22 ~ {\rm s}$. (b) Pruning BART with the random paradigm under the inference delay threshold $T_0 =0.22 ~ {\rm s}$. (c) Pruning BERT with the magnitude paradigm under the inference delay threshold $T_0 =0.35 ~ {\rm s}$. (d) Pruning BERT with the random paradigm under the inference delay threshold $T_0 =0.21 ~ {\rm s}$.

Theorems & Definitions (14)

• Lemma 3.1
• Lemma 3.2
• Proposition 3.1
• Theorem 3.1
• Remark 3.1
• Proposition 3.2
• Remark 3.2
• Remark 4.1
• Remark 4.2
• Lemma 4.1