Search papers, labs, and topics across Lattice.
This paper introduces SHIFT-LLM, a training-free framework that employs Linear Residual Adapters (LRAs) to correct the distributional mismatch caused by depth pruning in large language models. By preserving the identity pathway of pruned Transformer blocks and adding lightweight affine corrections, SHIFT-LLM effectively mitigates accuracy loss without the need for gradient computation. Experimental results demonstrate that this approach can recover significant accuracy, achieving up to a 15.7-point improvement on Llama-3.1-8B-Instruct across various configurations and benchmarks.
Depth pruning doesn't have to mean sacrificing accuracy鈥擲HIFT-LLM recovers lost performance with minimal overhead.
Depth pruning removes entire Transformer blocks to reduce the inference cost of large language models, but disrupts the hidden-state distributions expected by downstream layers, leading to significant accuracy loss. We introduce SHIFT-LLM, a training-free post-pruning correction framework that inserts a Linear Residual Adapter (LRA) at each pruning site. Each LRA preserves the identity pathway of the original residual block and adds a lightweight affine residual correction. This correction is calibrated via closed-form least-squares regression on a small held-out set, without gradient computation, to approximate the missing residual update produced by the pruned block. Together with the preserved identity pathway, the resulting LRA output approximates the hidden state produced by the original block, thereby mitigating the distributional mismatch introduced by layer removal while avoiding the expensive attention and feed-forward computations of the removed blocks. The resulting LRAs support low-rank factorization and exact merging across consecutive pruned layers for additional compression, and combine naturally with parameter-efficient fine-tuning for further recovery beyond fine-tuning the pruned model alone. Experiments on five model families, six layer-selection criteria, and seven zero-shot benchmarks show that SHIFT-LLM consistently recovers accuracy lost to depth pruning across most configurations, achieving gains up to +15.7 points on Llama-3.1-8B-Instruct while requiring only a few hundred calibration samples and no gradient computation.
