Pruning_LayerCollapse

An experimental implementation of LaCo: Large Language Model Pruning via Layer Collapse.

The basic idea is pretty neat: instead of simply deleting transformer blocks, we merge a layer with some of the layers after it, then remove those layers. The hope is that the model keeps enough of its behaviour while becoming shallower and cheaper to run.

This repository is still very much an experiment. It is useful for poking at the method and seeing where layers can be collapsed, but it is not yet a reproduction of the paper's full Llama-2/Baichuan2 evaluation.

What is here?

src/Layercollapse/laco.py
    Main LaCo/RDSC merge experiment on Qwen/Qwen2.5-0.5B.

src/Layercollapse/lesser_constraints.py
    A more relaxed experiment that can run against different Hugging Face
    causal language models.

src/pruning/
    Earlier structured-pruning experiments and quantization code.

The current merge follows the RDSC-style update:

theta_merged = theta_base + sum(theta_next - theta_base)

After a merge, the surviving decoder layers are reindexed. This matters for newer versions of Transformers: Llama-style attention layers keep a layer_idx used by the KV cache, and leaving the old indices in place causes an IndexError after pruning.

Setup

Create or activate a Python environment, then install the dependencies:

pip install -r requirements.txt

The models are downloaded from Hugging Face on the first run, so the download size can be significant.

Running it

From the repository root:

python src/Layercollapse/laco.py

For the more relaxed, model-selectable experiment:

python src/Layercollapse/lesser_constraints.py

You can pass another Hugging Face causal language model:

python src/Layercollapse/lesser_constraints.py \
  --model Qwen/Qwen2.5-0.5B

From inside src/Layercollapse, use python lesser_constraints.py instead. Do not prefix the path with src/Layercollapse/ from that directory, or the path will be duplicated.

The relaxed experiment currently uses:

C = 8       # maximum merge window
L = 0       # allow checking from the first layer
I = 1       # move through layers one at a time
T = 0.30    # cosine-similarity acceptance threshold

These settings are intentionally permissive. They are useful for discovering possible collapse locations, not for claiming that the resulting model still works well.

Findings so far

SmolLM2-135M

The relaxed experiment completed successfully:

Layers before: 30
Layers after : 1
Layers removed: 29

The first accepted large merge was around layer 21. After that, many single-layer merges were accepted by the low threshold. This shows that the code can perform layer collapse, but collapsing 29 of 30 layers is clearly an overly aggressive experiment. The similarity score is also based on only five short calibration prompts, so it is not a substitute for perplexity or task evaluation.

TinyLlama and cache indexing

The first TinyLlama run failed after a successful merge with:

IndexError: list index out of range

The cause was not the merge itself. The remaining attention blocks still had their original layer indices, while Transformers created a cache for the new, shorter model. Reindexing the remaining layers fixed this particular error.

Memory limits

The experiments keep a modified candidate model while evaluating it against a reference. On this laptop, the larger TinyLlama run was eventually killed with exit code 137 because of memory pressure. The baseline comparison was then changed to cache the baseline logits once, which avoids keeping an additional full reference model in memory.

Llama 2-7B was also tried, but its roughly 13.5 GB checkpoint downloaded too slowly to finish and would still need considerably more memory for repeated candidate evaluation. Llama 2-13B was not run for the same reason. The paper's Llama-2-7B/13B and Baichuan2-7B/13B experiments need a much larger machine, ideally with a more memory-efficient evaluation strategy.

Important caveats

• T=0.30 is deliberately loose. A merge being accepted does not mean the model quality is preserved.
• The current score compares next-token logits with cosine similarity. It is a quick signal, not a proper benchmark.
• The calibration set is tiny and hand-written.
• The current experiments use CPU because repeated model copies caused MPS out-of-memory failures.
• The paper evaluates base Llama-2 and Baichuan2 models. Chat/instruction checkpoints such as TinyLlama/TinyLlama-1.1B-Chat-v1.0 are useful for debugging, but they are not an exact match for that setup.

What should happen next?

The next useful step is to add a real evaluation loop: perplexity on a held out text set, plus a small generation comparison before and after collapse. Then the threshold can be tuned to a target removal ratio, such as 10%, 20%, or 25%, instead of allowing the relaxed experiment to collapse nearly the entire network.