🤖 machine learning

CoCurve: Cross-Module Co-Pruning Curvature for Training-Free Structured LLM Pruning

CoCurve is a training-free structured pruning method for LLMs that jointly optimizes the removal of attention and FFN units by leveraging a second-order Taylor expansion to capture cross-module dependencies through a Fisher matrix, enabling efficient, label-free compression via a single-shot quadratic program without fine-tuning.

Original authors: Zhiren Gong, Zihao Zeng, Zijie Wang, Tiantong Wang, Chau Yuen, Wei Yang Bryan Lim

Published 2026-07-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhiren Gong, Zihao Zeng, Zijie Wang, Tiantong Wang, Chau Yuen, Wei Yang Bryan Lim

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 have a giant, super-smart robot brain made of billions of tiny gears and springs. This is a Large Language Model (LLM), a type of AI that can write stories, solve math problems, and even code software. But here's the catch: these brains are so huge and heavy that they are incredibly expensive to run and slow to think. To make them useful for everyday people, scientists need to shrink them down, like packing a massive library into a pocket-sized book. This process is called "pruning."

Think of pruning like editing a long movie. You want to cut out the boring scenes to make the movie shorter and faster to watch, but you don't want to ruin the plot. For a long time, editors (scientists) thought the best way to do this was to look at each scene individually. If a scene looked dull on its own, they'd cut it. They assumed that if you cut a bunch of dull scenes, the movie would just be a little shorter, but the story would stay the same. They treated the movie like a simple list of independent clips.

However, this paper argues that a movie (or an AI brain) isn't just a list of independent clips. It's a complex web where scenes talk to each other. Sometimes, a scene that looks boring on its own is actually the secret glue holding two other important scenes together. If you cut that "boring" scene, the whole story might fall apart, even if you didn't cut the other scenes. The old way of editing missed these hidden connections.

Enter CoCurve: The Detective Who Maps the Connections

The researchers behind this paper, from Nanyang Technological University, introduced a new method called CoCurve (Cross-Module Co-Pruning Curvature). Instead of just looking at individual gears to see if they are weak, CoCurve acts like a detective who maps out how every gear is connected to every other gear.

Here is how they did it, using a simple analogy: Imagine the AI brain is a giant city with two types of buildings: Attention Towers (which help the AI focus on specific words) and FFN Factories (which process and understand those words).

The Old Way (The "Node-First" Mistake): Traditional methods would walk through the city, check each building individually, and say, "This tower looks a bit empty, let's knock it down," or "This factory seems slow, let's close it." They assumed that knocking down a few empty buildings wouldn't hurt the city. But they missed the fact that some "empty" towers were actually the only bridges connecting two busy factories. Knocking them down caused traffic jams that stopped the whole city.
The CoCurve Way: CoCurve doesn't just check the buildings; it checks the bridges between them. It uses a clever mathematical trick (a second-order Taylor expansion of a "KL divergence," which is just a fancy way of measuring how much the AI's answers change when you hide a part of it) to draw a map of the city.
- On this map, the diagonal shows how important a single building is on its own (like the old method).
- The off-diagonal lines show the "co-pruning curvature." This is the extra damage caused if you knock down two specific buildings together. It reveals that sometimes, two weak buildings are actually a super-strong team, and you must keep them both.

The Magic Trick: No Training Required

Usually, figuring out these connections requires a lot of heavy lifting, like retraining the AI or running millions of tests. But CoCurve is a "training-free" method. It's like a one-shot magic trick.

The researchers take a small, unlabeled sample of text (like a few pages of a book).
They run the AI through this text, but they "mask" (hide) one building at a time to see how the AI's answers wiggle.
Using a neat mathematical shortcut (a Gram product), they can calculate the entire map of connections just from these single tests. They don't need to test every possible pair of buildings (which would take forever), nor do they need to update the AI's weights. It's all done in one go.

What They Found

The team tested CoCurve on four different AI models, ranging from 3 billion to 24 billion parameters. They compared it against nine other popular pruning methods.

The Results: CoCurve consistently outperformed the others. It kept the AI smarter and more capable, especially on tricky tasks like writing code or solving math problems, where other methods often made the AI "collapse" and fail completely.
The Big Win: On a massive 24-billion-parameter model, when they pruned away 40% of the brain, CoCurve was 14.5 times better at predicting the next word (a measure called perplexity) than the next-best method.
The "Bridge" Discovery: The paper proved that the gains didn't come from just finding the "best" buildings to keep. The gains came specifically from the cross-module edges—the connections between the Attention Towers and the FFN Factories. When they removed these connections from the math, the method performed much worse, proving that these hidden bridges are the secret sauce.

When Does It Work?

The researchers also found a rule for when this detective work is most useful. If a model has too many redundant buildings (like a city with 100 identical factories that all do the same thing), the map of connections gets noisy and confusing. In those cases, CoCurve has a "safety switch" (a damping factor) that turns off the complex bridge-mapping and just sticks to the simple "check each building" method. This ensures it never makes things worse, even on models where the connections aren't as critical.

The Bottom Line

CoCurve changes the game by teaching us that in a complex AI brain, you can't just prune the nodes (the parts); you have to prune the edges (the connections). By respecting the hidden teamwork between different parts of the AI, CoCurve creates smaller, faster models that don't lose their smarts. It's a smarter way to edit the movie, ensuring the plot stays intact even when the runtime is cut in half.

Technical Summary: CoCurve

Problem Statement

Structured pruning compresses Large Language Models (LLMs) by removing whole computational units, such as attention heads and feed-forward network (FFN) channel groups. While training-free methods are preferred for their lack of data and compute requirements, they typically suffer from a critical blind spot: they rank removable units independently.

Most existing approaches assume that the loss from pruning a set of units is the sum of the individual losses of those units ( $\Delta L(S) \approx \sum \Delta L(u)$ ). This assumption fails in Transformer architectures where sublayers are coupled through a shared residual stream. Consequently, two units that appear weak individually (low saliency) may be jointly indispensable. Independent scoring is blind to this coupling and often removes such "bridge units" alongside the units they support, causing disproportionate damage to the model's generative capabilities, particularly in code and math tasks.

Methodology: CoCurve

The authors introduce CoCurve (Cross-Module Co-Pruning Curvature), a calibration-only, training-free method that jointly prunes attention and FFN units by modeling their interactions.

1. Pruning Risk as Second-Order Self-Distillation

The method defines pruning risk as the token-level Kullback-Leibler (KL) divergence between the frozen full model (teacher) and its masked copy (student) over unlabeled text.

Zeroth and First Order: At zero pruning, the distributions coincide, so the risk and its gradient vanish.
Second Order: The leading term of the Taylor expansion is purely quadratic: $R(s) \approx \frac{1}{2} s^\top H s$ .
The Fisher Matrix ( $H$ ): This single matrix $H$ $H$ contains:
- Diagonal ( $H_{uu}$ ): Classical node saliency (similar to Optimal Brain Damage).
- Off-Diagonal ( $H_{uv}$ ): Co-pruning curvature edges, representing the extra distortion incurred when removing units $u$ and $v$ together.

2. Practical Estimation via Single-Unit Ablations

Computing a full $M \times M$ interaction matrix typically requires $O(M^2)$ pairwise ablations. CoCurve avoids this by deriving a closed-form estimator:

The off-diagonal entries $H_{uv}$ are shown to be a Gram product of single-unit logit perturbation features.
By performing only $M$ single-unit ablation passes (one per unit) and computing the Fisher-weighted dot products of their logit shifts, the full interaction matrix is recovered.
This requires no gradients, no Hessian-vector products, and no pairwise sweeps.

3. One-Shot Greedy Solver

Pruning is formulated as a budgeted binary quadratic program. The authors employ a cost-normalized greedy solver that:

Selects units to remove based on marginal risk, which includes both the node saliency and the accumulated interaction with already-selected units.
Operates under a shared budget for attention and FFN units, allowing the optimization to determine the optimal split rather than fixing it manually.
Includes a trust-region scalar ( $\lambda$ ) to dampen the edge term at high pruning ratios where the second-order approximation degrades. This $\lambda$ is determined by a label-free rule based on the model's FFN channel redundancy.

4. Compensation (Optional)

For aggressive ratios, a closed-form compensation step can re-scale the outputs of kept FFN groups to absorb the residual stream shift caused by removed units. This step is also gradient-free and uses the same edge blocks computed during selection.

Key Contributions

A Node-to-Edge Objective: Deriving a quadratic pruning risk where the off-diagonal terms represent cross-module co-pruning curvature between attention and FFN units, moving beyond independent node ranking.
A Practical Estimator and Solver: Demonstrating that the full interaction matrix is a Gram product of single-unit ablation features, recoverable in $O(M)$ passes. The method solves the selection problem in a single shot without weight updates or recovery.
An Interpretable Mechanism: The method explicitly protects "bridge units" (low saliency but high connectivity) that node-first methods discard. It places structured pruning on a continuum with classical second-order methods (OBD/OBS).
A Predictive Rule: A single calibration statistic (FFN channel redundancy) predicts a priori when the edge term should be trusted or damped, turning an empirical observation into an architecture-level rule.

Results

The authors evaluated CoCurve on four dense LLMs (3B–8B) and a scaling study on a 24B model, comparing against nine training-free baselines across 14 benchmarks (including language modeling, commonsense, knowledge, math, and code).

Performance: CoCurve consistently outperforms strong training-free baselines. On the headline model (Llama-3.1-8B-Instruct) at 20% pruning, it achieves the best scores in language modeling, world knowledge, and code generation, and the highest average for commonsense tasks.
Fragile Tasks: The method shows the largest gains on fragile generative tasks (code and math), where other methods often collapse. At 40% pruning on the 24B model, CoCurve achieves 14.5× lower perplexity than the strongest valid baseline.
Ablations:
- Removing the off-diagonal edge term reduces performance to that of a diagonal-only (OBD-style) selector, confirming the edges drive the gain.
- Joint attention-FFN pruning outperforms pruning either module alone.
- The edge benefit is inversely correlated with FFN channel redundancy; the method includes a gate to dampen edges on high-redundancy models (e.g., Qwen2.5-32B) to prevent degradation.

Significance and Claims

The paper reframes structured LLM pruning: "Prune the edges, not just the nodes."

Mechanism: The authors claim that the failure of existing methods stems from ignoring the coupling between heterogeneous modules (attention and FFN) via the residual stream. CoCurve captures this coupling through a derived Fisher curvature matrix.
Efficiency: The method is strictly training-free, requiring no labels, fine-tuning, or iterative re-estimation. It is computationally efficient, requiring only forward passes for single-unit ablations.
Robustness: The approach provides a "graceful fallback." On high-redundancy models where edges might mislead, the method automatically reverts to a strong diagonal (OBD-style) selector, ensuring it never performs worse than a strong classical baseline.
Scope: The authors position CoCurve as a solution for moderate compression regimes (up to ~30-40%). They acknowledge that like all training-free methods, it degrades super-linearly at extreme ratios without recovery, but it significantly extends the usable compression frontier compared to independent ranking methods.

The paper concludes that cross-module interactions are not merely decorative but are load-bearing for preserving the functional computation required for complex generation, and that modeling these interactions is essential for effective structured pruning.