Paper Reading AI Learner

TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

2026-02-02 22:49:39
Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

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.

Abstract (translated)

结构化剪枝对于大型语言模型(LLM)的高效部署至关重要。由于LLM子模块对剪枝敏感度不同,因此需要识别最优的非均匀剪枝模型。现有的方法分别评估层、注意力头或权重通道的重要性,这种局部关注忽略了模型内部复杂的全局结构性依赖关系。训练感知结构化剪枝可以解决这一问题,但其计算成本可能与后剪枝训练一样昂贵。为了减轻训练感知剪枝的计算负担并捕获全局结构性依赖关系,我们提出了TraceNAS,这是一种无需训练的神经架构搜索(NAS)框架,它同时探索LLM深度和宽度的结构化剪枝。TraceNAS通过使用比例不变的零样本代理来识别与预训练模型损失景观高度对齐的剪枝模型,有效地选择在后剪枝训练中表现出最大性能潜力的模型。TraceNAS效率极高,在单个GPU上仅需8.5小时即可发现高质量的剪枝模型,相比训练感知方法减少了10倍的GPU时间消耗。在Llama和Qwen家族上的评估表明,TraceNAS在常识和推理基准测试方面与训练感知基线具有竞争力。

URL

https://arxiv.org/abs/2602.02891

PDF

https://arxiv.org/pdf/2602.02891.pdf


Tags
3D Action Action_Localization Action_Recognition Activity Adversarial Agent Attention Autonomous Bert Boundary_Detection Caption Chat Classification CNN Compressive_Sensing Contour Contrastive_Learning Deep_Learning Denoising Detection Dialog Diffusion Drone Dynamic_Memory_Network Edge_Detection Embedding Embodied Emotion Enhancement Face Face_Detection Face_Recognition Facial_Landmark Few-Shot Gait_Recognition GAN Gaze_Estimation Gesture Gradient_Descent Handwriting Human_Parsing Image_Caption Image_Classification Image_Compression Image_Enhancement Image_Generation Image_Matting Image_Retrieval Inference Inpainting Intelligent_Chip Knowledge Knowledge_Graph Language_Model LLM Matching Medical Memory_Networks Multi_Modal Multi_Task NAS NMT Object_Detection Object_Tracking OCR Ontology Optical_Character Optical_Flow Optimization Person_Re-identification Point_Cloud Portrait_Generation Pose Pose_Estimation Prediction QA Quantitative Quantitative_Finance Quantization Re-identification Recognition Recommendation Reconstruction Regularization Reinforcement_Learning Relation Relation_Extraction Represenation Represenation_Learning Restoration Review RNN Robot Salient Scene_Classification Scene_Generation Scene_Parsing Scene_Text Segmentation Self-Supervised Semantic_Instance_Segmentation Semantic_Segmentation Semi_Global Semi_Supervised Sence_graph Sentiment Sentiment_Classification Sketch SLAM Sparse Speech Speech_Recognition Style_Transfer Summarization Super_Resolution Surveillance Survey Text_Classification Text_Generation Time_Series Tracking Transfer_Learning Transformer Unsupervised Video_Caption Video_Classification Video_Indexing Video_Prediction Video_Retrieval Visual_Relation VQA Weakly_Supervised Zero-Shot