🤖 machine learning

Reduced Matrix Multiplication: Input-Adaptive Matrix-Product Reduction for LLM Inference

This paper introduces Reduced Matrix Multiplication (RMM), a training-free, input-adaptive inference method that reduces computational costs in Transformer-based models by selectively retaining informative slices of matrix products, achieving scalable accuracy-efficiency trade-offs across various model sizes, tasks, and modalities without modifying model weights.

Original authors: Zixuan Lan, Yanhong Li, Jiawei Zhou

Published 2026-08-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zixuan Lan, Yanhong Li, Jiawei Zhou

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 the brain of a modern computer as a massive, hyper-intelligent librarian. This librarian doesn't just read books; it writes them, solves math problems, and even looks at pictures to tell stories. To do this, it uses a special kind of brain architecture called a "Transformer." Think of a Transformer as a giant factory where information flows through assembly lines. At every station on these lines, the factory performs a massive calculation: it multiplies huge grids of numbers together. These are the "matrix multiplications" that power the magic. The problem is, these calculations are incredibly heavy. They require a lot of energy and time, like trying to move a mountain of sand with a teaspoon. As these AI models get smarter and bigger, the "mountain" gets taller, making it expensive and slow to run them. Scientists have long wondered: Is the librarian actually using every single grain of sand in that mountain for every single task, or is there a lot of unnecessary sand we could sweep away without the librarian noticing?

This paper introduces a clever new trick called Reduced Matrix Multiplication (RMM). Instead of trying to permanently shrink the factory or delete parts of the librarian's memory (which can break things), RMM acts like a smart, on-the-fly filter. Every time the librarian is about to do a massive calculation, RMM pauses and asks, "Which specific numbers in this grid are actually doing the heavy lifting right now?" It then picks only the most important ones—say, the top 50% or 70%—and ignores the rest for that specific moment. It's like a chef who, instead of chopping every single vegetable in the fridge for a soup, quickly tastes the broth and decides to only use the carrots and onions that are currently needed for this specific flavor. The best part? The librarian doesn't need to be retrained or taught how to do this; it just works naturally with the numbers it's already generating.

The researchers tested this idea on a wide variety of AI models, ranging from small ones with 1 billion parameters to giants with 70 billion. They found that these models are surprisingly flexible. When they reduced the amount of math being done by keeping only a portion of the numbers, the models didn't fall apart. In fact, the bigger the model, the more sand it seemed willing to let go of without losing its mind. For example, the giant 70-billion-parameter model could handle a 50% reduction in calculations while still answering questions and writing stories almost as well as before. However, the paper also discovered a funny quirk in how these factories are built: the "attention" parts of the brain (where the model decides what to focus on) are very chill and can handle losing half their work easily. But the "MLP" parts (the parts that actually process and transform the information) are much more sensitive; if you cut too much there, the model starts to stumble.

The team also showed that this trick works for models that can see and talk, not just text-only ones. They even built special software tools to prove that skipping these numbers actually makes the computer run faster in real life, especially when the stories or sentences get very long. While the paper suggests this is a promising way to make AI cheaper and quicker, it also notes that there isn't a single "perfect" setting for everyone. You have to tune how much you cut based on the specific model and what you're asking it to do. But overall, RMM offers a fresh, training-free way to make these digital giants run lighter, proving that sometimes, doing less math can actually help you think just as clearly.

Technical Summary: Reduced Matrix Multiplication (RMM) for LLM Inference

Problem Statement
Transformer-based language models achieve strong performance but incur substantial inference costs driven by repeated high-dimensional matrix multiplications in attention and feed-forward (MLP) layers. While prior work has addressed inference efficiency through structured pruning, low-rank approximation, token compression, or KV-cache management, these methods typically modify fixed model structures or shorten inputs/caches. They do not directly address whether the contracted computation inside each matrix product can be adaptively reduced for a specific input without modifying model weights. Existing activation-sparsity methods (e.g., TEAL, CATS) skip low-magnitude activation entries but focus on sparsifying hidden states rather than reducing the shared contraction axes of general matrix products.

Methodology: Reduced Matrix Multiplication (RMM)
The authors propose Reduced Matrix Multiplication (RMM), a training-free, input-adaptive inference method. RMM operates on the unified form of Transformer computations, $Y = AB$, where $A$ represents the activation matrix and $B$ represents the weight matrix (or intermediate representation).

Core Mechanism: Instead of computing the full product $AB$, RMM dynamically selects a subset of indices $I$ along the shared contraction dimension $d$ based on the current input's activation magnitudes. It computes the reduced product $A_{:,I} B_{I,:}$ .
Selection Strategy: The method employs activation-aware dimension selection. For each feature dimension $j$ , an importance score $s_j = \|A_{:,j}\|_2$ (the column norm of the activation matrix) is calculated. The top- $k$ indices with the highest scores are retained, where $k = \lceil \rho d \rceil$ and $\rho$ is a user-controlled retention ratio.
Theoretical Basis: The authors prove that TopK selection by column norm is minimax optimal for minimizing the worst-case approximation error over all possible weight matrices $B$ when only $A$ is known at selection time. This provides a principled justification for using activation magnitudes to guide the reduction.
Application Scope: RMM is applied to:
- Attention Layers: Reducing the feature dimension in $QK^\top$ (query-key interaction) and optionally the token dimension in $PV$ (attention-value interaction).
- MLP and Linear Projections: Reducing the hidden dimension in feed-forward projections ($XW$).
Implementation: The method requires computing feature scores and performing top- $k$ selection, which are lightweight vector operations compared to the dense matrix multiplications they replace. Custom Triton kernels are implemented to realize these savings in practice.

Key Contributions

Input-Adaptive Reduction: Introduces a novel perspective on inference optimization that reduces the contracted computation of matrix products themselves, distinct from pruning weights, compressing tokens, or managing caches.
Training-Free and Controllable: The method requires no retraining and offers a smooth, predictable accuracy-efficiency trade-off controlled by a simple retention ratio ( $\rho$ ).
Mechanistic Insights: Through ablation studies, the paper reveals a structural asymmetry in Transformers:
- Attention-side computations (Q, K, V projections, $QK^\top$ , $PV$) are substantially more reducible and robust to aggressive pruning.
- MLP components are significantly more sensitive to reduction, with the "Up" projection being the most sensitive and the "Down" projection being more robust. Pruning the entire MLP block leads to severe performance collapse.
Generalization: Demonstrates that the principle extends beyond text-only models to multimodal vision-language inference (e.g., Qwen2.5-VL).

Experimental Results
The authors evaluated RMM across models ranging from 1B to 70B parameters (including LLaMA 3.1/3.2 and Qwen families) and diverse tasks (QA, reasoning, summarization, long-context, and vision-language).

Performance vs. Static/Random Baselines: Under fixed retention ratios (e.g., $\rho=0.5$ ), RMM consistently outperforms static pruning methods (SparseGPT, Wanda, SliceGPT) and random pruning. Static methods suffer from unstable degradation, while RMM maintains coherent generation and reasoning capabilities.
Scaling Trends: Larger models generally tolerate more aggressive reduction. For instance, LLaMA 3.1 70B remains close to full-model performance at $\rho=0.8$ on most benchmarks, whereas smaller models degrade more quickly, particularly on challenging tasks like GSM8K and HumanEval.
Robustness: RMM remains stable in autoregressive generation and long-context settings (up to 30K tokens), showing smooth degradation rather than abrupt failure.
Runtime Efficiency: Benchmarks on NVIDIA A100 GPUs using custom kernels show that computational savings translate into practical wall-clock speedups, especially at longer sequence lengths (e.g., 1.40x speedup at 4096 tokens). Notably, at 4096 tokens, the dense implementation ran out of memory, while RMM completed inference.
Compatibility: The method is compatible with INT8 weight quantization and extends to various vision-language model backbones.

Significance and Claims
The paper positions RMM as a scalable direction for input-adaptive inference-time optimization. Its primary significance lies in demonstrating that redundancy in Transformer inference is not uniformly distributed; rather, it is concentrated in specific components (particularly attention-side operations) and varies by input. By leveraging this structural asymmetry and input-dependent activation information, RMM achieves controllable efficiency gains without the need for model retraining or structural modification. The authors conclude that matrix-product-level adaptive reduction is a promising avenue for efficient Transformer inference, suggesting that future methods should account for component-specific redundancy patterns.

Limitations
The work focuses on a training-free formulation. The authors acknowledge they have not explored extensions beyond this setting (e.g., integrating adaptive reduction into pretraining objectives) nor conducted a detailed study of redundancy patterns across all components of complex vision-language models beyond the initial evaluation. Additionally, while custom kernels show speedups, full integration with all inference frameworks and quantized backends remains future work.

Technical Summary: Reduced Matrix Multiplication (RMM) for LLM Inference

More like this