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

Theophilus Amaefuna, Hitesh Ulhas Vaidya, Anshuman Chhabra, Ankur Mali
Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, PMLR 337:62-86, 2026.

Abstract

Layer-wise capacity in large language models is highly non-uniform: some layers contribute disproportionately to loss reduction, whereas others are nearly redundant. Existing layer-scoring methods provide sensitivity estimates but do not give a principled rule for converting those estimates into allocation or pruning decisions under a global hardware budget. We introduce a curvature-aware, MDL-inspired framework built around the layer gain $\zeta_k^2=g_k^\top\widetilde H_{kk}^{-1}g_k$. This quantity equals twice the maximal decrease predicted by the regularized layer-restricted quadratic model and incorporates inverse local curvature; it is therefore a local surrogate for reducible risk, not a universal dominance claim over gradient-norm scores. After normalizing the gains into scores $q_k$, we formulate two convex programs: one allocates expert slots under diminishing returns, and the other assigns layer-wise pruning ratios while protecting high-score layers. Both continuous programs have unique globally optimal solutions characterized by one dual variable and computable in $O(K\log(1/\varepsilon))$ time by bisection. We also prove a quadratic transfer-regret bound: when source and target score vectors differ by at most $\delta$, the target surrogate cost of the transferred decision is within $O(\delta^2)$ of the target optimum. Experiments on {Mistral-7B} and Gemma-7B show clear allocation gains in some settings and competitive, though mixed, pruning performance. The framework therefore replaces an empirical score-to-decision heuristic with a budget-feasible optimization procedure whose guarantees apply to the stated continuous surrogates. Code is available on github repo - \href{https://github.com/TKAI-LAB-Mali/Curvature-Weighted-Capacity-Allocation.git}{TKAI-LAB-Mali/Curvature-Weighted-Capacity-Allocation}

Cite this Paper


BibTeX
@InProceedings{pmlr-v337-amaefuna26a, title = {Curvature-Weighted Capacity Allocation: A Minimum Description Length Framework for Layer-Adaptive Large Language Model Optimization}, author = {Amaefuna, Theophilus and Vaidya, Hitesh Ulhas and Chhabra, Anshuman and Mali, Ankur}, booktitle = {Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence}, pages = {62--86}, year = {2026}, editor = {Perković, Emilija and Malinsky, Daniel}, volume = {337}, series = {Proceedings of Machine Learning Research}, month = {17--21 Aug}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v337/main/assets/amaefuna26a/amaefuna26a.pdf}, url = {https://proceedings.mlr.press/v337/amaefuna26a.html}, abstract = {Layer-wise capacity in large language models is highly non-uniform: some layers contribute disproportionately to loss reduction, whereas others are nearly redundant. Existing layer-scoring methods provide sensitivity estimates but do not give a principled rule for converting those estimates into allocation or pruning decisions under a global hardware budget. We introduce a curvature-aware, MDL-inspired framework built around the layer gain $\zeta_k^2=g_k^\top\widetilde H_{kk}^{-1}g_k$. This quantity equals twice the maximal decrease predicted by the regularized layer-restricted quadratic model and incorporates inverse local curvature; it is therefore a local surrogate for reducible risk, not a universal dominance claim over gradient-norm scores. After normalizing the gains into scores $q_k$, we formulate two convex programs: one allocates expert slots under diminishing returns, and the other assigns layer-wise pruning ratios while protecting high-score layers. Both continuous programs have unique globally optimal solutions characterized by one dual variable and computable in $O(K\log(1/\varepsilon))$ time by bisection. We also prove a quadratic transfer-regret bound: when source and target score vectors differ by at most $\delta$, the target surrogate cost of the transferred decision is within $O(\delta^2)$ of the target optimum. Experiments on {Mistral-7B} and Gemma-7B show clear allocation gains in some settings and competitive, though mixed, pruning performance. The framework therefore replaces an empirical score-to-decision heuristic with a budget-feasible optimization procedure whose guarantees apply to the stated continuous surrogates. Code is available on github repo - \href{https://github.com/TKAI-LAB-Mali/Curvature-Weighted-Capacity-Allocation.git}{TKAI-LAB-Mali/Curvature-Weighted-Capacity-Allocation}} }
Endnote
%0 Conference Paper %T Curvature-Weighted Capacity Allocation: A Minimum Description Length Framework for Layer-Adaptive Large Language Model Optimization %A Theophilus Amaefuna %A Hitesh Ulhas Vaidya %A Anshuman Chhabra %A Ankur Mali %B Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence %C Proceedings of Machine Learning Research %D 2026 %E Emilija Perković %E Daniel Malinsky %F pmlr-v337-amaefuna26a %I PMLR %P 62--86 %U https://proceedings.mlr.press/v337/amaefuna26a.html %V 337 %X Layer-wise capacity in large language models is highly non-uniform: some layers contribute disproportionately to loss reduction, whereas others are nearly redundant. Existing layer-scoring methods provide sensitivity estimates but do not give a principled rule for converting those estimates into allocation or pruning decisions under a global hardware budget. We introduce a curvature-aware, MDL-inspired framework built around the layer gain $\zeta_k^2=g_k^\top\widetilde H_{kk}^{-1}g_k$. This quantity equals twice the maximal decrease predicted by the regularized layer-restricted quadratic model and incorporates inverse local curvature; it is therefore a local surrogate for reducible risk, not a universal dominance claim over gradient-norm scores. After normalizing the gains into scores $q_k$, we formulate two convex programs: one allocates expert slots under diminishing returns, and the other assigns layer-wise pruning ratios while protecting high-score layers. Both continuous programs have unique globally optimal solutions characterized by one dual variable and computable in $O(K\log(1/\varepsilon))$ time by bisection. We also prove a quadratic transfer-regret bound: when source and target score vectors differ by at most $\delta$, the target surrogate cost of the transferred decision is within $O(\delta^2)$ of the target optimum. Experiments on {Mistral-7B} and Gemma-7B show clear allocation gains in some settings and competitive, though mixed, pruning performance. The framework therefore replaces an empirical score-to-decision heuristic with a budget-feasible optimization procedure whose guarantees apply to the stated continuous surrogates. Code is available on github repo - \href{https://github.com/TKAI-LAB-Mali/Curvature-Weighted-Capacity-Allocation.git}{TKAI-LAB-Mali/Curvature-Weighted-Capacity-Allocation}
APA
Amaefuna, T., Vaidya, H.U., Chhabra, A. & Mali, A.. (2026). Curvature-Weighted Capacity Allocation: A Minimum Description Length Framework for Layer-Adaptive Large Language Model Optimization. Proceedings of the 42nd Conference on Uncertainty in Artificial Intelligence, in Proceedings of Machine Learning Research 337:62-86 Available from https://proceedings.mlr.press/v337/amaefuna26a.html.

Related Material