Curvature-Weighted Capacity Allocation: A Minimum Description Length Framework for Layer-Adaptive Large Language Model Optimization

TL;DR

Two convex MDL programs are formulated: a capacity allocation program that distributes expert slots or LoRA rank preferentially to high-curvature layers under diminishing returns, and a pruning program that concentrates sparsity on low-gain layers while protecting high-gain layers from degradation.

Abstract

Layer-wise capacity in large language models is highly non-uniform: some layers contribute disproportionately to loss reduction while others are near-redundant. Existing methods for exploiting this non-uniformity, such as influence-function-based layer scoring, produce sensitivity estimates but offer no principled mechanism for translating them into allocation or pruning decisions under hardware constraints. We address this gap with a unified, curvature-aware framework grounded in the Minimum Description Length (MDL) principle. Our central quantity is the curvature-adjusted layer gain $ζ_k^2 = g_k^\top \widetilde{H}_{kk}^{-1} g_k$, which we show equals twice the maximal second-order reduction in empirical risk achievable by updating layer $k$ alone, and which strictly dominates gradient-norm-based scores by incorporating local curvature. Normalizing these gains into layer quality scores $q_k$, we formulate two convex MDL programs: a capacity allocation program that distributes expert slots or LoRA rank preferentially to high-curvature layers under diminishing returns, and a pruning program that concentrates sparsity on low-gain layers while protecting high-gain layers from degradation. Both programs admit unique closed-form solutions parameterized by a single dual variable, computable in $O(K \log 1/\varepsilon)$ via bisection. We prove an $O(δ^2)$ transfer regret bound showing that source-domain allocations remain near-optimal on target tasks when curvature scores drift by $δ$, with explicit constants tied to the condition number of the target program. Together, these results elevate layer-wise capacity optimization from an empirical heuristic to a theoretically grounded, computationally efficient framework with provable optimality and generalization guarantees.

Figures (4)

• Figure 1: (a) Layer-wise curvature scores $\zeta_k^2$ vary substantially across the transformer stack. High-gain layers (dark bars) hold disproportionate reducible risk and should receive additional capacity; low-gain layers (light bars) are candidates for aggressive pruning. (b) Our framework computes $\zeta_k^2$ from per-layer gradients $g_k$ and regularized Hessian blocks $\widetilde{H}_{kk}$, normalizes them to quality scores $q_k$, and solves two convex programs: a capacity allocation program (Theorem \ref{['thm:alloc']}) that enriches high-gain layers, and a pruning program (Theorem \ref{['thm:prune']}) that concentrates sparsity on low-gain layers. Both programs admit closed-form solutions via $O(K\log 1/\varepsilon)$ bisection (Algorithms \ref{['alg:alloc']}--\ref{['alg:prune']}). Theorem \ref{['thm:transfer']} bounds the cost of transferring source-domain allocations to a target domain.
• Figure 2: Expert allocation accuracy (%) on Mistral-7B-v0.1 (5 epochs)
• Figure 3: Expert allocation accuracy (%) on Gemma-7B, +ve variant only (5 epochs)
• Figure : MDL-Optimal Expert Allocation

Theorems & Definitions (10)

• Lemma 1: Layer-restricted optimum
• Theorem 2: Convexity and closed form for allocation
• Theorem 3: Strong convexity and closed form for pruning
• Remark : Budget flexibility
• proof
• proof
• proof
• Theorem 4: Transfer regret under score drift
• proof
• Remark : Boundary solutions