Curvature-Weighted Capacity Allocation

August 18, 2026 · View on GitHub

A Minimum Description Length inspired framework for capacity allocation in a Mixture-Of-Experts Large Language Model

UAI 2026 · OpenReview
Arxiv · link

This repository contains the code for two complementary experiments from the paper:

  1. Expert Allocation — uses MDL to assign a non-uniform number of experts per layer for Mixture-of-LoRA (MoLA) fine-tuning.
  2. Layer-wise Pruning — uses MDL to derive per-layer sparsity ratios for unstructured pruning of large language models.

Both experiments share a common MDL core (mdl/) that computes curvature-weighted layer importance scores, which are then consumed by the respective downstream pipelines.

Workflow

workflow

Repository structure

mdl-layerIF/
├── mdl/                            # MDL core: importance scoring, expert allocation, pruning ratios
│   ├── logUtility.py               # MDL → per-layer expert counts (expert allocation)
│   ├── pruning.py                  # MDL → per-layer sparsity ratios (pruning)
│   ├── my_layer_influence.py       # Layer Influence Function computation (kronfluence)
│   ├── util.py / util_mali.py      # Shared MDL utilities
│   ├── get_B.py                    # Curvature (Fisher/Hessian) estimation helpers
│   ├── model_profiling.py          # Model parameter counting utilities
│   ├── data/                       # Pre-computed MDL pruning ratio JSONs (paper results)
│   └── data_dp/                    # Depth-prior variant pruning ratio JSONs
│
├── Expert_Allocation/              # MoLA fine-tuning pipeline
│   ├── LayerIF_Computation/        # Layer Influence Function computation for expert allocation
│   │   ├── compute_IF.py           # Computes per-layer IF values using kronfluence
│   │   ├── lora_model_alpha.py     # LoRA α estimation helper
│   │   └── expert_allocator.ipynb  # Notebook: IF values → expert count assignments
│   ├── src/                        # MoLA model/trainer (adapted from HF Transformers + PEFT)
│   ├── utils/                      # prompter.py (alpaca), callbacks.py (streaming, AGPL-3.0)
│   ├── templates/                  # Alpaca prompt templates
│   ├── configs/                    # Dataset/expert configuration JSONs
│   ├── run_all.sh                  # Train on all six datasets
│   ├── eval_all.sh                 # Evaluate on all six datasets
│   └── expert_number.py            # AlphaLora baseline expert allocation
│
├── LayerIF_Pruning_New/            # Layer-IF pruning pipeline
│   ├── main.py                     # Entry point: prune + evaluate an LLM
│   ├── run_mistral_mdl.sh          # Run MDL-ratio pruning on Mistral-7B
│   ├── run_mistral_mdl_depth_prior.sh  # Depth-prior variant
│   ├── wait_then_run_mistral.sh    # Sequential job launcher
│   ├── data/                       # Pre-computed LayerIF/alpha/OWL metric caches (Mistral, Gemma)
│   └── full-environment.yml        # Conda environment for pruning experiments
│
├── pyproject.toml                  # uv environment for MDL core scripts
├── requirements.txt                # Full pipeline dependencies (pinned lm-eval, kronfluence, …)
├── LICENSE                         # MIT (see NOTICE for third-party licenses)
└── NOTICE                          # Third-party attributions

Setup

There are three environments depending on which part of the pipeline you are running.

MDL core (mdl/) — uv

The MDL scoring scripts (logUtility.py, pruning.py, my_layer_influence.py, etc.) use a lightweight environment managed with uv.

# Install uv if you don't have it
curl -LsSf https://astral.sh/uv/install.sh | sh

# Create the environment and install all dependencies (reads pyproject.toml + uv.lock)
uv venv --python 3.10
uv sync

To activate:

source .venv/bin/activate   # Linux/macOS

Dependencies are declared in pyproject.toml: numpy, scipy, matplotlib, pandas, scikit-learn, weightwatcher, safetensors, powerlaw, tqdm.

Quickstart

Expert Allocation

The pipeline runs in four stages:

1. Compute Layer Influence (IF) values

cd Expert_Allocation/LayerIF_Computation
python compute_IF.py

Outputs per-layer IF scores to outputs/layerIF_values/<model>/.

2. Derive per-layer expert counts from MDL

python mdl/logUtility.py

Reads IF scores and applies MDL budget allocation. Adjust the rho scaling factor (line 188) to target a total expert count (default 160, matching the MoLA convention).

3. Train MoLA on six datasets

Edit Expert_Allocation/run_all.sh to set base_model, number_experts, and top_k, then:

bash Expert_Allocation/run_all.sh

4. Evaluate

bash Expert_Allocation/eval_all.sh

Layer-wise Pruning

Pre-computed MDL pruning ratios are already in mdl/data/. To run pruning directly:

bash LayerIF_Pruning_New/run_mistral_mdl.sh

To re-derive pruning ratios from scratch:

python mdl/pruning.py

Outputs per-layer ratio JSONs consumed by LayerIF_Pruning_New/main.py.

Example pruning command (Mistral-7B-v0.1, magnitude + MDL ratios, 50% sparsity):

CUDA_VISIBLE_DEVICES=0,1,2,3 python LayerIF_Pruning_New/main.py \
    --model mistralai/Mistral-7B-v0.1 \
    --prune_method magnitude_ww \
    --sparsity_ratio 0.5 \
    --ww_metric IF-300-96-smoothed \
    --ww_metric_cache LayerIF_Pruning_New/data/mistral-7b/ \
    --epsilon 0.3 \
    --eval_zero_shot

Citation

@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}
}

License

This repository is released under the MIT License. See LICENSE and NOTICE for details, including third-party attributions. Expert_Allocation/utils/callbacks.py is licensed separately under AGPL-3.0-only (see its header).