Technical Summary: TraceNAS

Problem Definition

The deployment of Large Language Models (LLMs) is constrained by massive parameter counts and computational requirements. While structured pruning offers a pathway to efficient deployment by removing architectural units (e.g., attention heads, transformer blocks), existing methods face significant limitations:

1. Local vs. Global Dependencies: Most pruning methods evaluate layer, head, or channel importance in isolation using local heuristics (e.g., activation-weighted magnitudes or Hessian-based reconstruction). These approaches ignore complex global structural dependencies and the propagation of pruning-induced errors across the model.
2. Computational Cost of Training-Aware Search: Methods that attempt to capture global dependencies often rely on training-aware Neural Architecture Search (NAS). These approaches require iterative gradient updates and continuous streams of unseen tokens to prevent overfitting, making the search process computationally as expensive as the post-pruning recovery training itself.
3. Inadequacy of Existing Zero-Shot Proxies: Current training-free proxies are designed to identify trainable, randomly initialized architectures by focusing on trainability (e.g., gradient stability, loss landscape smoothness). They fail to evaluate "functional inheritance"—the ability of a pruned model to recover performance by preserving the pretrained loss landscape.

Methodology

TraceNAS is proposed as a training-free NAS framework designed to jointly optimize the depth and width of LLMs without search-time training.

1. Search Space and Encoding

The framework treats the pretrained LLM as a super-network. The search space $G$ is defined as a joint distribution over:

• Depth ($d$): A discrete mask over $L$ transformer blocks, allowing for non-uniform depth reduction while maintaining residual stream continuity.
• Width ($\kappa$): Parameter retention ratios for attention and MLP sub-blocks, enabling heterogeneous sparsity across the model.

2. In-Place Architectural Realization

To avoid the memory overhead of instantiating thousands of candidate models, TraceNAS employs an in-place masking strategy:

• Width Masking: Uses an activation-weighted product (similar to Wanda) to generate binary masks for weight channels. This $O(d^2)$ heuristic identifies functionally critical channels while preserving outlier features.
• Depth Masking: Bypasses entire blocks if deactivated, routing activations through residual connections.
• Execution: The framework modifies the base model's weight matrices temporarily to compute gradients on a minimal calibration set, then restores the original weights, preserving the base model for subsequent iterations.

3. The Zero-Shot Proxy: Gradient Trace Correlation

The core innovation is a scale-invariant proxy $\Phi$ that measures functional alignment between a pruned candidate ($M_{sub}$) and the pretrained base ($M_{base}$).

• Gradient Traces: Gradients are computed on a calibration set. To ensure scalability, these are projected into a Low-Rank (LoRA) subspace, leveraging the intrinsic low-dimensionality of LLMs to reduce memory overhead by orders of magnitude.
• Alignment Metric: The proxy calculates the Pearson Correlation Coefficient between the standardized gradient traces of the pruned and base models for each layer. This ensures scale invariance, focusing on directional alignment rather than magnitude shifts.
• Sparsity-Weighted Aggregation: The final score $\Phi$ is a weighted sum of layer-wise correlations, where weights correspond to the parameter retention ratios. This anchors the global score in high-capacity regions, preventing noise from aggressively pruned sub-blocks from skewing the ranking.

4. Evolutionary Search

TraceNAS employs an evolutionary strategy to navigate the discrete search space:

• Initialization: A population is generated based on global importance priors.
• Evolution: The population undergoes crossover and mutation. Depth evolution uses bit-flip mutations and contiguous-range crossover, while width evolution uses interpolation and Gaussian jitter.
• Selection: Candidates are ranked solely by the proxy $\Phi$, selecting models with the highest potential for performance recovery during post-pruning continued pretraining (CPT).

Key Contributions

1. Scalable Training-Free NAS: TraceNAS provides a unified framework for joint depth and width pruning, enabling the discovery of non-uniform architectures without search-time training.
2. Functional Inheritance Proxy: Introduces a novel, scale-invariant gradient-based proxy that evaluates the alignment of gradient traces between pruned and pretrained models. It achieves high correlation with downstream accuracy (Spearman $\rho = 0.94$, Kendall $\tau = 0.82$).
3. Efficiency: The framework reduces search overhead by utilizing low-rank gradient subspaces and in-place masking, enabling high-fidelity model discovery on a single GPU.

Experimental Results

TraceNAS was evaluated on Llama-2-7B, Llama-3.1-8B, and Qwen-2.5-14B, pruning them to 2.7B, 4.6B, and 8.4B parameters respectively.

• Search Efficiency: TraceNAS identified optimal architectures in 8.5 GPU hours (on an A100), representing a 10× reduction in search time compared to training-aware baselines like DarwinLM. It required only 98M search tokens, a 10× reduction compared to DarwinLM.
• Performance:
  • Llama-2-7B: TraceNAS achieved the highest average accuracy (62.81%) across eight reasoning benchmarks, outperforming DarwinLM (62.57%) and ShearedLLaMA (62.63%). Notably, ShearedLLaMA's performance dropped significantly (to 60.86%) when restricted to the same token budget used by TraceNAS.
  • Llama-3.1-8B & Qwen-2.5-14B: TraceNAS outperformed uniform pruning baselines and training-aware methods (DarwinLM, E3-Pruner) on complex reasoning tasks (e.g., ARC-Challenge, BoolQ) while requiring significantly fewer recovery tokens (20B vs. 94B for Minitron).
• Proxy Validation: The proxy demonstrated high stability across varying hyperparameters (sample size, context length, and LoRA rank), confirming that gradient trace alignment effectively captures model inheritance even under constrained settings.

Significance and Claims

The paper claims that TraceNAS addresses the critical bottleneck of computational cost in structured pruning. By shifting the focus from "trainability" (learning from scratch) to "inheritance" (preserving the pretrained loss landscape), TraceNAS enables the discovery of high-performing, non-uniform pruned models with minimal resource expenditure.

The authors emphasize that their method:

• Eliminates the prohibitive cost of search-time training, making large-scale model compression accessible.
• Captures global structural dependencies that local heuristics miss, ensuring the preservation of complex reasoning capabilities.
• Reduces the carbon footprint of model compression by drastically lowering the GPU-hours required for architecture search.

The work positions TraceNAS as a scalable foundation for deploying efficient LLMs in resource-constrained environments while maintaining the functional integrity of dense pretrained models.