AERO: Softmax-Only LLMs for Efficient Private Inference
The pervasiveness of proprietary language models has raised privacy concerns
for users' sensitive data, emphasizing the need for private inference (PI),
where inference is performed directly on encrypted inputs. However, current PI
methods face prohibitively higher communication and latency overheads,
primarily due to nonlinear operations. In this paper, we present a
comprehensive analysis to understand the role of nonlinearities in
transformer-based decoder-only language models. We introduce AERO, a four-step
architectural optimization framework that refines the existing LLM architecture
for efficient PI by systematically removing nonlinearities such as LayerNorm
and GELU and reducing FLOPs counts. For the first time, we propose a
Softmax-only architecture with significantly fewer FLOPs tailored for efficient
PI. Furthermore, we devise a novel entropy regularization technique to improve
the performance of Softmax-only models. AERO achieves up to 4.23times
communication and 1.94times latency reduction. We validate the effectiveness
of AERO by benchmarking it against the state-of-the-art.
6711c3fea79b431fea2c3b48
Nandan Kumar Jha
admin_assigned
2024-10-18T16:02:37.993Z
670ec3f6db1a6bcfe832e0a6
https://cdn-avatars.huggingface.co/v1/production/uploads/no-auth/-mZNLeLJoXzkwPgYO38lF.png
Nandan Kumar Jha
user
nandan523
6711c3fea79b431fea2c3b49
Brandon Reagen
AERO: Softmax-Only LLMs for Efficient Private Inference
The pervasiveness of proprietary language models has raised privacy concerns
for users' sensitive data, emphasizing the need for private inference (PI),
where inference is performed directly on encrypted inputs. However, current PI
methods face prohibitively higher communication and latency overheads,
primarily due to nonlinear operations. In this paper, we present a
comprehensive analysis to understand the role of nonlinearities in
transformer-based decoder-only language models. We introduce AERO, a four-step
architectural optimization framework that refines the existing LLM architecture
for efficient PI by systematically removing nonlinearities such as LayerNorm
and GELU and reducing FLOPs counts. For the first time, we propose a
Softmax-only architecture with significantly fewer FLOPs tailored for efficient
PI. Furthermore, we devise a novel entropy regularization technique to improve
the performance of Softmax-only models. AERO achieves up to 4.23times
communication and 1.94times latency reduction. We validate the effectiveness
of AERO by benchmarking it against the state-of-the-art.
AERO framework optimizes transformer-based LLMs for efficient private inference by removing nonlinearities and using a Softmax-only architecture with reduced FLOPs and improved entropy regularization.
private inference LLM architecture AERO LayerNorm GELU FLOPs Softmax-only architecture entropy regularization