WaLeM: Weight-Aware Learnable Merging for Large Language Model Depth Compression

The residual connections in modern Large Language Models (LLMs) can be formulated as forward-Euler discretizations of continuous-depth dynamical systems. From this Neural ODE perspective, heuristic depth-compression techniques—such as layer pruning and merging—mathematically correspond to step-size coarsening operations over discrete trajectories. However, existing methods implicitly apply uniform coarsening schemes that ignore the heterogeneous smoothness of distinct Transformer operators, thereby incurring unmitigated local truncation errors. To address this, we introduce WaLeM, a principled, smoothness-aware operator-split coarsening framework for LLM depth compression. Specifically, WaLeM estimates local trajectory smoothness via Centered Kernel Alignment (CKA) and employs dynamic programming to construct an adaptive temporal grid that minimizes a global truncation-error surrogate. Furthermore, it learns component-wise quadrature weights to strictly eliminate operator-coupling errors. Theoretically, we establish an end-to-end approximation error bound of $O(L{M}_f((1-\gamma )B+L_f{B}^2)e^{L_{f}L})$ and demonstrate that our scheduling strategy optimally minimizes this bound. Across five mainstream LLM families, WaLeM outperforms state-of-the-art pruning and merging baselines with or without post-compression recovery, isolating the gain to the coarsening itself rather than to subsequent fine-tuning. Notably, it exhibits exceptional robustness on long-chain reasoning tasks (e.g., GSM8K), where early-stage discretization errors typically compound exponentially. Our work formally bridges heuristic layer reduction with rigorous numerical analysis, establishing a unified design space for continuous-depth LLM compression.