💻 computer science

STARFISH: faST Accuracy Recovery in pruned networks From Internal State Healing

The paper introduces STARFISH, an efficient healing method that recovers most of the accuracy lost during neural network pruning by aligning the pruned model's internal state representations with the original model using a tiny set of unlabeled calibration data, significantly outperforming state-of-the-art techniques especially under aggressive pruning conditions.

Original authors: Shir Maon, Odelia Melamed, Adi Shamir

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shir Maon, Odelia Melamed, Adi Shamir

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, incredibly smart library (a neural network) filled with millions of books (weights). This library can answer almost any question perfectly. But, the library is so huge that it takes up too much space, costs too much to run, and is slow to search through.

To fix this, you decide to throw away 50%, 75%, or even 90% of the books. This is called pruning.

The Problem:
When you throw away so many books, the library stops working well. It starts giving wrong answers. Usually, to fix this, you have to hire a team of librarians to re-read the entire original collection of books and re-learn how to organize the remaining ones. This is expensive, takes a long time, and often, you don't even have access to the original books anymore (they might be private or proprietary).

The Solution: STARFISH
The authors of this paper introduce a new method called STARFISH. They named it after the sea creature that can regrow its lost arms. Just like a starfish can heal itself even after losing limbs, STARFISH helps a "pruned" (damaged) network heal itself and regain its intelligence.

Here is how it works, using simple analogies:

1. The "Ghost" vs. The "Real"

Imagine the original, full library is a Master Chef who knows exactly how to cook a perfect meal. After you throw away most of the ingredients (pruning), you are left with a Junior Chef who is confused and can't cook well.

Usually, to fix the Junior Chef, you'd make them taste the original dishes and try to memorize the final taste (the output). But the paper says this isn't enough when the damage is severe.

Instead, STARFISH looks at the internal thoughts of the Master Chef.

It asks the Master Chef: "When you see a tomato, what do you think about before you decide it's a tomato? What about when you think about the sauce? What about the heat?"
It records these internal "thoughts" (called internal states or representations) for a few sample dishes.
Then, it tells the Junior Chef: "Don't just try to guess the final taste. Try to think exactly like the Master Chef did at every single step of the process."

2. The Tiny "Calibration" Set

The magic of STARFISH is that it doesn't need the whole library of books to do this. It only needs a tiny handful of examples (like 1,000 images) to act as a "calibration set."

It's like giving the Junior Chef a single, small recipe card with a few notes on how the Master Chef thought.
The Junior Chef practices on these few examples, adjusting their thinking process until their internal "thoughts" match the Master Chef's thoughts perfectly.

3. Why It's a Big Deal

The paper claims this method is a game-changer, especially when you cut away a lot of the network:

The "50% Cut" Scenario: If you remove half the weights, STARFISH recovers the network's accuracy almost perfectly (up to 99.8% of the original). It beats the current best methods by a huge margin.
The "75% Cut" Scenario: This is where it gets really impressive. If you remove 75% of the weights, other methods usually fail miserably, leaving the network with only 40% of its original smarts. STARFISH, however, manages to recover 82% of the original accuracy.
The "85% Cut" Scenario: Even when you remove 85% of the weights (leaving almost nothing), STARFISH can still restore the network to about 92% of its original performance.

4. The "No Original Data" Superpower

Most repair methods require you to have the original training data (the millions of books) to re-teach the network.

STARFISH's Advantage: It works even if you don't have the original training data. It only needs the damaged network, the list of what was removed (the mask), and a tiny, random set of unlabeled pictures. This makes it perfect for modern situations where companies might share a model but keep their training data secret.

Summary

Think of STARFISH as a "mind-meld" technique. Instead of forcing a damaged network to guess the right answer, it forces the damaged network to think exactly like the original, healthy network did. By aligning these internal thoughts using just a few examples, it can heal even the most severely "amputated" networks, bringing them back to near-perfect performance without needing the original training data.

Technical Summary: STARFISH

Problem Statement
Pruning neural networks is a standard technique for reducing model size, memory footprint, and inference latency by removing weights, neurons, or channels. However, aggressive pruning often leads to a substantial degradation in predictive accuracy. While recovery processes (healing) are typically employed to regain lost performance, existing methods face significant limitations. Traditional approaches often require full retraining or fine-tuning on the original training dataset, which is computationally expensive and frequently impossible in modern settings where training data is proprietary or inaccessible. Furthermore, state-of-the-art recovery methods often struggle to restore accuracy under high-sparsity conditions (e.g., removing >75% of weights), particularly for modern Vision Transformer (ViT) architectures.

Methodology: STARFISH
The authors propose STARFISH (faST Accuracy Recovery in pruned networks From Internal State Healing), a post-pruning recovery method that operates without access to the original training data or labels. The core hypothesis is that optimizing the alignment of intermediate representations (internal states) between the pruned network and the original dense network is a more effective proxy for accuracy restoration than optimizing solely on final output predictions.

The algorithm proceeds in two stages using a small, unlabeled calibration set ( $S_{cal}$ ) sampled from the data distribution but distinct from the training set:

Caching Target Representations: The original dense model is run on $S_{cal}$ to cache the vectorized outputs of every computational block (e.g., after normalization, non-linearities, and residual additions).
Alignment Optimization: The pruned model ( $f_{\theta \odot m}$ ) is trained to minimize the alignment loss between its intermediate representations and the cached targets. The loss function is defined as the average cosine distance between the pruned model's hidden states ( $f^\ell_{\theta \odot m}(x)$ ) and the dense model's hidden states ( $r^\ell_x$ ) across all blocks $\ell$ :
$\mathcal{L}_{align}(\theta \odot m, x) = \frac{1}{L} \sum_{\ell \in [L]} \left( 1 - \frac{\langle f^\ell_{\theta \odot m}(x), r^\ell_x \rangle}{\|f^\ell_{\theta \odot m}(x)\| \|r^\ell_x\|} \right)$
The pruned model's remaining weights are updated to minimize the empirical loss over the calibration set.

Theoretical motivation is provided by showing that the distance between the final hidden representations (inputs to the classification head) serves as an upper bound for the Kullback–Leibler (KL) divergence of the output distributions. This suggests that aligning internal states inherently constrains the output error.

Key Contributions

Representation-Based Recovery: The paper introduces a recovery method independent of specific pruning types, masks, or architectures. It empirically demonstrates that intermediate representation alignment outperforms direct output-based objectives (such as Cross-Entropy or KL divergence matching) as sparsity levels increase.
High-Sparsity Efficiency: The method is the first to efficiently restore high percentages of original accuracy at extreme sparsity levels (e.g., 85% sparsity). It achieves this using a tiny calibration set (as few as 1,000 examples, or 0.4% of the training set size) and requires no access to the original training data.
Modularity and Speed: The approach is computationally lightweight, requiring only a few forward and backward passes on a small dataset. It is applicable to structured, semi-structured, and unstructured pruning patterns.

Experimental Results
The authors evaluated STARFISH on ImageNet-1K using various Vision Transformer (DeiT, ViT) and Convolutional Neural Network (MobileNet) architectures.

High Sparsity (0.75–0.85): On a DeiT-B model pruned to 75% sparsity, STARFISH recovered 82% of the original dense accuracy, whereas competing methods (like CORP) recovered only 40%. On a DeiT3-H model at 85% sparsity, STARFISH retained 92.08% of the dense accuracy, improving the pruned model's performance by 72.36 percentage points.
Moderate Sparsity (0.5): Across various DeiT architectures and pruning scopes (MLP, QK, and joint MLP+QK), STARFISH consistently outperformed state-of-the-art methods. For example, in a DeiT-T model with 50% structured pruning on MLP+QK, STARFISH achieved 81.95% retention of dense accuracy, a 31.54 percentage-point improvement over the CORP method.
Comparison: STARFISH demonstrated superior performance against methods like SNOWS (which uses Hessian approximation) and CORP (which uses closed-form affine recovery), particularly in regimes where pruning severely degrades the model.

Significance and Claims
The paper claims that STARFISH represents a significant advancement in the field of model compression by shifting the recovery paradigm from output matching to internal state healing. The authors emphasize that this approach is particularly vital for modern large-scale models where training data is unavailable and full retraining is infeasible. By leveraging the principle that internal representations capture the model's knowledge more robustly than final outputs under high sparsity, STARFISH enables the deployment of highly efficient, pruned networks without sacrificing predictive performance. The method is presented as a general, plug-and-play solution that can be applied to any pruned network regardless of the initial pruning strategy.