TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

TL;DR

TraceNAS introduces a training-free neural architecture search framework for non-uniform depth and width pruning of large language models. It relies on a scale-invariant gradient trace proxy, Φ, to measure functional inheritance by aligning pruned sub-networks with the pretrained base on a low-rank gradient manifold. An in-place masking strategy and an evolutionary search navigate the discrete depth and continuous width space, yielding high-performing pruned architectures with dramatically reduced search cost. Across Llama and Qwen families, TraceNAS matches or exceeds training-aware baselines while reducing GPU hours by roughly an order of magnitude, and demonstrates robust generalization to different architectures and reasoning benchmarks. The work advances practical compression of LLMs by enabling efficient, high-fidelity architecture discovery that preserves pretrained knowledge and post-pruning recovery potential.

Abstract

Structured pruning is essential for efficient deployment of Large Language Models (LLMs). The varying sensitivity of LLM sub-blocks to pruning necessitates the identification of optimal non-uniformly pruned models. Existing methods evaluate the importance of layers, attention heads, or weight channels in isolation. Such localized focus ignores the complex global structural dependencies that exist across the model. Training-aware structured pruning addresses global dependencies, but its computational cost can be just as expensive as post-pruning training. To alleviate the computational burden of training-aware pruning and capture global structural dependencies, we propose TraceNAS, a training-free Neural Architecture Search (NAS) framework that jointly explores structured pruning of LLM depth and width. TraceNAS identifies pruned models that maintain a high degree of loss landscape alignment with the pretrained model using a scale-invariant zero-shot proxy, effectively selecting models that exhibit maximal performance potential during post-pruning training. TraceNAS is highly efficient, enabling high-fidelity discovery of pruned models on a single GPU in 8.5 hours, yielding a 10$\times$ reduction in GPU-hours compared to training-aware methods. Evaluations on the Llama and Qwen families demonstrate that TraceNAS is competitive with training-aware baselines across commonsense and reasoning benchmarks.

Figures (5)

• Figure 1: Search Efficiency. TraceNAS identifies optimal non-uniform architectures in 8.5 GPU hours, achieving competitive accuracy with 10$\times$ less search overhead than training-aware baselines. The area of the bubble is proportional to the total tokens used for recovery training, highlighting that TraceNAS identified architectures have high recovery potential.
• Figure 2: Visualization of the TraceNAS search framework. TraceNAS uses a gradient-based, training-free proxy to guide structural pruning. Following a one-time initialization of base gradient traces ($g_{\text{base}}$) and importance scores ($I_l$), a population of depth ($\mathbf{d}$) and width ($\boldsymbol{\kappa}$) candidates undergoes iterative evolution via crossover and mutation. Each width configuration is realized using an $O(d^2)$ activation-weighted heuristic. Subsequently, candidates $\mathcal{M}_{sub}$ are ranked by the zero-shot proxy $\Phi$, which measures the gradient trace alignment between the active layers of $\mathcal{M}_{sub}$ relative to $\mathcal{M}_{base}$.
• Figure 3: TraceNAS proxy stability analysis. We report Kendall $\tau$ correlation between ranking scores across search hyperparameters (a) number of samples (N), (b) context length (CL), and (c) LoRA rank ($r$) which define the x-axis. High $\tau$ values demonstrate that TraceNAS consistently ranks models across search settings.
• Figure 4: TraceNAS Evolutionary Dynamics and Search Convergence. Search trajectory for pruning Llama-2-7B to 2.7B across 50 iterations. (a) TraceNAS score $(\Phi)$ evolution: Illustrates the discovery of top scored candidates across the specified parameter budget window. The red star indicates the model with maximal functional inheritance for the given constraint. (b) Attention Width Evolution: Tracks the sparsity ratios for the attention sub-block $(\kappa_{attn})$; the search identifies specific layers where attention heads are critical for maintaining representational flow. (c) MLP Width Evolution: Sparsity ratios for MLP sub-blocks $(\kappa_{mlp})$, revealing the high structural sparsity and exploration across the model depth. (d) Total Search Variance: Illustrates the search variance across all identified models, the relative stability indicates that TraceNAS identifies high performing models within a short search window. In figures (b), (c) and (d), $\kappa = 0.0$ indicates no pruning in that layer and $\kappa = 1$ indicates the layer has been dropped.
• Figure 5: TraceNAS PPL across sparsity levels. WikiText-2 perplexity reported for pruned models identified via TraceNAS evolutionary search and trained using 2.5B tokens of CPT.