🤖 machine learning

Prox: Training-Free FFN Activation Sparsity via Approximate Intermediate-Channel Salience in LLMs

Prox is a training-free framework that enables efficient sparse SwiGLU FFN execution in LLMs by constructing channel masks from approximate intermediate-channel salience, thereby achieving significant speedups and superior model quality compared to existing methods without requiring retraining.

Original authors: Jinyi Liu, Wei Chen, Pengyu Chen, Xinyi Yuan, Minghe Bai, Guoquan Wu, Jun Wei

Published 2026-07-31

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Jinyi Liu, Wei Chen, Pengyu Chen, Xinyi Yuan, Minghe Bai, Guoquan Wu, Jun Wei

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to run a massive, hyper-intelligent robot brain on a small, battery-powered device. This brain, known as a Large Language Model (LLM), is incredible at writing stories, solving math problems, and chatting, but it is also incredibly hungry. It gobb up memory and processing power just to think about the next word in a sentence. To make these robots faster and less thirsty, scientists have been trying to teach them to "skip" parts of their thinking process that aren't important. This is called activation sparsity. Think of it like a chef who, instead of chopping every single vegetable in the kitchen for a soup, quickly tastes the ingredients and only chops the ones that will actually add flavor, leaving the rest untouched. The goal is to get the same delicious soup (the correct answer) but with much less work. However, the tricky part is figuring out which vegetables to skip without accidentally throwing away the secret spice that makes the dish taste right. If you guess wrong, the soup tastes bland, or worse, inedible.

This is the puzzle tackled by a new method called Prox, developed by researchers Jinyi Liu and his team. They focused on a specific type of robot brain architecture called SwiGLU, which is the current standard for many powerful AI models. The researchers discovered that while the brain's "middle thought" (an intermediate state) holds the perfect map for deciding what to skip, calculating that map exactly is too slow and expensive. So, they came up with a clever two-step trick. First, they use a fast, low-quality "rough draft" of the brain's thought to quickly decide which parts are important. Then, and this is the key, they only do the heavy, precise math on those specific important parts, ignoring the rest entirely. Their experiments showed that this method is much better than previous tricks that tried to guess the importance without doing the full math first. In tests across ten different AI models, Prox managed to speed up the robot's thinking by nearly double (up to 1.99 times faster) while keeping the answers just as smart as the slow, full-speed version. It's like having a super-fast scout who runs ahead to mark the path, allowing the main army to march only on the safe, paved roads, saving time and energy without getting lost.

The Problem: The Robot's Heavy Lunch

Large Language Models are like giant libraries of knowledge, but they are also very heavy. When they try to answer a question, they have to move a massive amount of data around inside their memory. A big chunk of this traffic comes from a specific part of the brain called the Feed-Forward Network (FFN). You can think of the FFN as the part of the brain that takes an idea, processes it, and passes it along. In modern models, this part uses a special recipe called SwiGLU, which involves three different math operations (projections) working together.

The problem is that these three operations are so heavy that they slow everything down. Scientists have tried to fix this by making the model "sparse," meaning they try to skip the math for parts of the brain that seem unimportant. However, existing methods for doing this without retraining the model (which takes forever) have a flaw. They often guess which parts to skip based on just one piece of information, like looking at only the "gate" or only the "up" part of the recipe. This is like trying to decide if a song is good by listening to only the drums or only the vocals. If you guess wrong, you skip a crucial note, and the song (or the AI's answer) falls apart. As the researchers tried to skip more and more parts to get faster, the quality of the answers dropped significantly.

The Insight: The Map vs. The Journey

The researchers, led by Jinyi Liu, noticed something interesting about the SwiGLU recipe. There is a specific moment in the process called the intermediate state (let's call it the "middle thought"). This middle thought is the perfect signal for deciding which parts of the brain to keep and which to skip. If you could see this middle thought exactly, you could skip up to 70% of the work and still get a perfect answer.

But here's the catch: to see this middle thought exactly, you have to do all the heavy math first. It's like trying to see the finish line of a race, but you have to run the whole race to get there. That defeats the purpose of trying to save time.

The team realized they didn't actually need the exact value of the middle thought. They only needed to know the order of importance. They needed to know which parts were the "big" ones and which were the "small" ones, not the exact numbers. This is the difference between knowing exactly how many dollars are in your wallet versus just knowing which bills are the biggest. If you can guess the order correctly, you can pick the right things to keep.

The Solution: Prox's Two-Stage Trick

To solve this, the team created Prox, a two-stage framework that acts like a smart scout and a precise builder.

Stage 1: The Quick Scout (Proxy Construction)
In the first stage, Prox uses a "cheap" version of the math. It takes the input and runs it through a simplified, lower-quality version of the brain (using quantized weights, which are like compressed, smaller numbers). It doesn't try to get the perfect answer; it just wants to get a rough idea of which channels are loud and which are quiet. It creates a "mask," which is basically a list of "Keep" and "Discard" tags. Because this stage uses simplified math, it's very fast and doesn't require much memory. Crucially, the errors in this rough calculation don't ruin the final answer because the rough numbers are never used for the final output—they are only used to make the list.

Stage 2: The Precise Builder (Exact Sparse Computation)
Once the list is made, the second stage kicks in. This is where the real magic happens. The model looks at the list from Stage 1 and says, "Okay, we are going to ignore 70% of the work and only do the math for the top 30%." But here is the difference from other methods: for those top 30%, Prox uses the original, full-precision, heavy-duty math. It calculates the exact values for the important parts. This ensures that the final answer is just as accurate as if the model had done all the work, but it only did the work for the parts that mattered.

Why It Works Better

The researchers tested Prox on ten different large language models, including popular ones like Qwen, Mistral, and Llama-3. They compared it against other methods that try to skip work without retraining.

The results were clear: Prox was the winner. At high levels of skipping (like 70% sparsity), other methods started to make mistakes, and the AI's answers got worse. Prox, however, kept the answers sharp. In fact, at 70% sparsity, Prox was able to speed up the decoding process (the speed at which the AI generates text) by up to 1.99 times on a standard graphics card. That means the AI could talk almost twice as fast without losing its smarts.

The team also checked if this new method would break if you tried to combine it with other speed-up tricks, like making the numbers smaller (quantization) or skipping parts of the attention mechanism. They found that Prox works perfectly alongside these other methods, like a universal adapter that fits into any setup.

The Takeaway

The paper concludes that Prox is a powerful, training-free way to make AI faster. It proves that you don't need to retrain a giant model to make it efficient. Instead, you can use a clever two-step process: use a fast, rough guess to find the important parts, and then do the heavy lifting only on those parts. This approach avoids the "guessing game" that caused other methods to fail, ensuring that the robot brain stays fast and smart, even when it's running on a small device. While the current version is best for single-user conversations (like chatting on a phone), the researchers suggest that with some more work, this idea could eventually help run even larger groups of users efficiently. For now, it stands as a major step forward in making AI faster without sacrificing its intelligence.

Technical Summary: Prox - Training-Free FFN Activation Sparsity via Approximate Intermediate-Channel Salience

Problem Statement

Large Language Models (LLMs) rely heavily on Feed-Forward Networks (FFNs), particularly the SwiGLU variant, which dominate memory traffic and computation during inference. While activation sparsity offers a pathway to mitigate these bottlenecks by skipping low-salience channels, existing training-free methods face a critical trade-off: they suffer substantial model-quality degradation at high sparsity levels.

Current approaches typically rely on single-branch signals (e.g., the gate activation $h$ or the up-branch activation $u$ ) to select channels for the intermediate state $s = u \odot h$ . This partial information leads to misranked salient channels. Alternatively, methods that apply sparsity sequentially (e.g., sparsifying input $x$ then the intermediate state $s$ ) compound approximation errors, degrading downstream performance. Furthermore, while the exact intermediate state $s$ provides an ideal signal for channel selection, computing it requires dense execution of the up and gate projections, negating the benefits of sparsity.

Methodology: The Prox Framework

The authors propose Prox, a two-stage, training-free framework designed to enable sparse execution of SwiGLU FFNs without retraining. The core insight is that sparse execution requires only the channel mask (determined by the magnitude ranking of the intermediate state), not the exact values of the state itself.

Stage 1: Lightweight Proxy Construction

The goal of Stage 1 is to construct a shared channel mask ( $m_s$ ) with minimal computational cost.

Input Sparsity & Quantization: Prox applies input sparsity to the FFN input $x$ and utilizes quantized proxy weights (INT4) for the up and gate projections.
Proxy Activation: It computes sparse proxy activations $\tilde{u}$ and $\tilde{h}$ using the quantized weights and sparse input.
Mask Generation: The proxy intermediate state is calculated as $\tilde{s} = \tilde{u} \odot \tilde{h}$ . The magnitude ranking of $\tilde{s}$ is used to generate a binary channel mask $m_s$ .
Error Isolation: Crucially, the approximation errors in $\tilde{s}$ are isolated to the selection process; these proxy values are not propagated downstream.

Stage 2: Exact Sparse Computation

Once the mask $m_s$ is established, Prox performs exact computation on the retained channels using the original full-precision weights.

Exact Projections: The selected channels are computed exactly for the up and gate projections ( $u_{sp}, h_{sp}$ ) using full-precision weights.
Final Output: The exact intermediate state $s_{sp}$ is computed and passed through the down projection (also sparsified by the same mask) to produce the final output.

Sparsity Allocation

Prox dynamically allocates the computational budget between Stage 1 and Stage 2.

Stage 2 Anchor: Based on "oracle" experiments showing that sparsifying ~70% of exact intermediate channels incurs negligible perplexity increase, Stage 2 sparsity ( $s_2$ ) is anchored around 70%.
Stage 1 Constraint: Stage 1 sparsity ( $s_1$ ) is constrained (e.g., retaining at least 30% of input entries) to ensure the proxy signal remains reliable for channel selection.
Cost Model: The framework optimizes $s_1$ and $s_2$ to meet a target effective sparsity while minimizing the normalized computational cost, which accounts for the overhead of quantized proxy projections.

Hardware Implementation

Prox utilizes custom CUDA and Triton kernels to realize acceleration:

Stage 1: A fused split-N kernel evaluates proxy projections, skipping rejected inputs and performing INT4 unpacking directly in registers.
Stage 2: A fused output-sparse kernel computes exact up/gate projections and an input-sparse GEMV for the down projection, restricting weight accesses to active channels.

Key Contributions

Insight on Intermediate State: The authors identify the SwiGLU intermediate state ( $s$ ) as the optimal channel-selection signal. They demonstrate that while exact $s$ is costly, a low-cost input-sparse proxy preserves the necessary magnitude ranking to recover the optimal mask.
Prox Framework: A novel two-stage method that decouples channel selection (via quantized proxies) from value computation (via exact weights). This prevents error accumulation common in prior training-free methods.
Comprehensive Evaluation: Extensive experiments across ten LLMs from six model families (including Qwen3, Llama-3, Mistral, and Gemma-3) show Prox outperforms state-of-the-art training-free baselines (CATS, COUNTDOWN, TEAL) at all sparsity levels, particularly at 60–70% sparsity.

Experimental Results

Accuracy: At 70% FFN sparsity, Prox achieves up to a 14.4% improvement in average downstream task performance compared to the best baseline (TEAL). It maintains model quality with less than a 3% relative perplexity increase in most cases, whereas other methods degrade significantly.
Speedup: On an NVIDIA A6000 GPU, Prox achieves an end-to-end decoding speedup of 1.51× to 1.99× at 60–70% sparsity.
Compatibility: Prox is orthogonal to model quantization (e.g., AWQ, FP8, W4A16) and sparse attention mechanisms (e.g., RocketKV, HShare), allowing for synergistic efficiency gains.

Significance and Claims

The paper claims that Prox represents a significant advancement in training-free LLM acceleration by resolving the accuracy-efficiency trade-off inherent in previous channel-selection strategies. By leveraging the observation that ranking stability is sufficient for mask generation, Prox enables aggressive sparsity (up to 70%) without the accuracy collapse seen in prior art. The framework is presented as a practical solution for resource-constrained, latency-sensitive inference, offering near-dense accuracy with nearly double the decoding throughput, while remaining fully compatible with existing quantization and sparse-attention ecosystems.

The authors note current limitations, specifically that Prox targets single-batch autoregressive decoding and does not yet support large-batch serving, and that storing INT4 proxy weights introduces a ~12% memory overhead, which future work could address via sliding-window schemes.