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

Paper: 2603.00910 Authors: Theophilus Amaefuna, Hitesh Vaidya, Anshuman Chhabra, Ankur Mali Categories: cs.IT, cs.AI, cs.LG

The Gap

We know LLM layers aren’t created equal—some do heavy lifting while others coast. Methods like influence functions can score layer importance, but they stop at measurement. The frontier problem: you have sensitivity estimates for each layer, but no principled way to convert those numbers into actual decisions when you’re constrained by memory, compute, or parameter budgets. Should you give layer 12 twice the LoRA rank of layer 8? Should you prune layer 15 to 90% sparsity or 95%? Existing approaches rely on ad-hoc thresholds or grid search, leaving a chasm between “this layer matters” and “allocate exactly this much capacity here.”

The Increment

Before: Layer importance scores → manual tuning → hope it works. After: Layer importance scores → convex optimization → provably optimal allocation with closed-form solutions.

The core mechanism is a two-stage pipeline. First, compute a curvature-adjusted gain $\zeta_k^2 = g_k^\top \widetilde{H}_{kk}^{-1} g_k$ for each layer $k$, where $g_k$ is the gradient and $\widetilde{H}_{kk}$ is the layer’s Hessian block. This quantity captures not just how steep the loss landscape is (gradient norm) but also how curved it is—layers in flat regions get downweighted even if their gradients are large. Second, normalize these gains into quality scores $q_k$ and feed them into one of two convex programs: an allocation program that distributes capacity (expert slots, LoRA rank) with diminishing returns, or a pruning program that concentrates sparsity on low-gain layers while shielding high-gain ones.

Think of it like portfolio rebalancing with risk-adjusted returns. Each layer is an asset. The gradient norm is like raw return, but the curvature is volatility—high curvature means small updates yield big changes (high risk-adjusted return). The quality score $q_k$ is your Sharpe ratio. The allocation program is your portfolio optimizer: given a fixed budget (total capacity), it distributes capital (parameters, rank) to maximize expected return under diminishing marginal utility (concave objective). The pruning program is the inverse: given a sparsity target, it sells off low-Sharpe assets first while protecting high-performers. The dual variable in both programs is the Lagrange multiplier—the “shadow price” of capacity or sparsity.

Both programs reduce to solving a single scalar equation via bisection in $O(K \log 1/\varepsilon)$ time, where $K$ is the number of layers. The solutions are unique and globally optimal because the objectives are strictly convex. The framework also comes with a transfer regret bound: if you compute allocations on a source task and apply them to a target task where curvature scores drift by $\delta$, your regret scales as $O(\delta^2)$—meaning allocations are robust to moderate distribution shift.

Key Concepts

Curvature-Adjusted Gain ($\zeta_k^2$)

Start with gradient descent. When you update a layer’s parameters by $-\eta g_k$, the loss drops by approximately $\eta \|g_k\|^2$ (first-order). But this ignores geometry—if the loss surface is highly curved, a small step might overshoot or undershoot. The Hessian $H_{kk}$ encodes this curvature. The Newton decrement $g_k^\top H_{kk}^{-1} g_k$ is the maximal second-order loss reduction achievable by updating layer $k$ alone, accounting for curvature. The authors use a damped version $\widetilde{H}_{kk} = H_{kk} + \lambda I$ for numerical stability, yielding $\zeta_k^2 = g_k^\top \widetilde{H}_{kk}^{-1} g_k$.

Concrete example: Imagine two layers with identical gradient norms $\|g_1\| = \|g_2\| = 10$. Layer 1 sits in a narrow valley (high curvature, $H_{11}$ has large eigenvalues), so $H_{11}^{-1}$ is small, and $\zeta_1^2$ is modest. Layer 2 sits on a gentle slope (low curvature, $H_{22}$ has small eigenvalues), so $H_{22}^{-1}$ is large, and $\zeta_2^2$ is huge. Layer 2 gets more capacity because updates there are more “efficient”—you get more bang per parameter buck.

Minimum Description Length (MDL) Principle

MDL is a formalization of Occam’s razor from information theory. The idea: the best model is the one that compresses the data most. Formally, the description length is $L(\theta) + L(D|\theta)$, where $L(\theta)$ is the cost of encoding the model parameters and $L(D|\theta)$ is the cost of encoding the data given the model (i.e., the loss). In this paper, MDL justifies the allocation objective: giving more capacity to high-gain layers reduces $L(D|\theta)$ (better fit), but capacity itself has a cost (larger $L(\theta)$). The convex program balances these via a concave utility function that encodes diminishing returns—the first bit of capacity to a layer helps a lot, the hundredth bit helps less.

Example: You have 100 LoRA rank slots to distribute across 10 layers. Layer 1 has $q_1 = 0.5$ (high quality), layer 10 has $q_{10} = 0.01$ (low quality). A naive uniform allocation gives each layer 10 slots. MDL says: give layer 1 maybe 30 slots and layer 10 maybe 2 slots, because the marginal utility of the 11th slot to layer 1 (high $q_1$) exceeds the marginal utility of the 3rd slot to layer 10 (low $q_{10}$). The exact allocation comes from solving $\sum_k \log(1 + \alpha q_k r_k) - \mu (\sum_k r_k - R) = 0$ for the dual variable $\mu$, where $r_k$ is the rank assigned to layer $k$ and $R$ is the total budget.

Transfer Regret Bound

You train on task A, compute curvature scores $q_k^A$, and derive an allocation $r_k^A$. Now you switch to task B with scores $q_k^B$. How suboptimal is $r_k^A$ on task B? The regret is the difference between the optimal objective on B using $r_k^B$ and the objective using $r_k^A$. The paper proves this regret is $O(\delta^2)$, where $\delta = \max_k |q_k^A - q_k^B|$ is the maximum drift in quality scores. The constant depends on the condition number of the allocation program on task B.

Why $\delta^2$? Because the objective is smooth (twice differentiable). By Taylor expansion, perturbing the input by $\delta$ changes the output by $O(\delta)$ (first-order) and the objective by $O(\delta^2)$ (second-order). This means if curvature scores are stable across tasks (small $\delta$), you can reuse allocations without recomputing—a huge win for multi-task or continual learning scenarios.

Expert Assessment

Problem significance: High. Parameter-efficient fine-tuning (PEFT) methods like LoRA and MoE are ubiquitous, but practitioners tune hyperparameters (rank, sparsity) via expensive grid search. A principled allocation framework could save compute and improve performance across the entire LLM ecosystem—millions of fine-tuning runs annually.

Method maturity: Proof-of-concept leaning toward practical. The convex programs are elegant and the closed-form solutions are a major plus. However, the paper doesn’t provide wall-clock timing comparisons for computing $\zeta_k^2$ (Hessian-vector products are expensive) or validate the $O(K \log 1/\varepsilon)$ bisection claim empirically. The transfer regret bound assumes $\delta$ is small, but the paper doesn’t characterize when this holds in practice—do curvature scores actually stay stable across tasks, or do they drift wildly? The damping parameter $\lambda$ in $\widetilde{H}_{kk}$ is also a hyperparameter, and the paper doesn’t discuss how to set it.

Experimental rigor: The abstract promises “provable optimality and generalization guarantees,” but the paper itself (based on the abstract) doesn’t describe experiments—only theory. This is a red flag. We need to see: (1) Does the curvature-adjusted gain $\zeta_k^2$ actually outperform gradient norm or influence functions on real LLMs? (2) Do the MDL allocations beat uniform or heuristic allocations on downstream tasks? (3) How sensitive are results to $\lambda$? Without empirical validation, this is a theory paper masquerading as a systems paper.

Verdict: Weak accept — the theoretical framework is solid and the problem is important, but the lack of experimental evidence and unaddressed practical concerns (Hessian computation cost, hyperparameter sensitivity) prevent a strong accept.

Takeaways

Curvature as a first-class citizen: Don’t just look at gradients—look at curvature. In any optimization problem where you’re allocating resources (compute, memory, attention), the curvature-adjusted gain $g^\top H^{-1} g$ is a better signal than $\|g\|^2$ because it accounts for local geometry. This applies beyond LLMs: neural architecture search, hyperparameter optimization, even A/B testing (allocate traffic to high-curvature variants).

Convex relaxations with closed-form solutions: When facing a discrete allocation problem (how many experts per layer? how much sparsity?), try formulating a continuous convex relaxation. If the objective is strictly convex, you get uniqueness and can solve via bisection on the dual variable. This pattern shows up in resource allocation, scheduling, and portfolio optimization.

Transfer regret as a design principle: If you’re building a system that needs to adapt across tasks (multi-task learning, domain adaptation), design your allocation mechanism so that the regret scales as $O(\delta^2)$ in the drift $\delta$. This means small changes in the environment don’t require full recomputation—you can reuse allocations with bounded suboptimality. The key is smoothness: make your objective twice differentiable and control the condition number.