trained-lm-heads

Trained LM Heads for Layer-Pruning Examination

Per-layer language modeling heads used to examine intermediate layers of 7B/8B decoder models before pruning. Each checkpoint is an lm_head trained to read out the hidden state at one specific layer, so you can measure how much predictive signal that layer already carries and use the resulting curve to decide where to cut.

These are diagnostic artifacts, not pruned model weights. Nothing here is sparsified or pruned — the tensors are dense and full-size. Do not use them as drop-in replacements for a base model's lm_head.

Contents

Two sets, one per base model, covering layers 16–31 (the back half of a 32-layer stack — the region where early exit is plausible):

lm_head_checkpoints_L/lm_head_prune16.pt  ...  lm_head_prune31.pt   (16 files)
lm_head_checkpoints_M/lm_head_prune16.pt  ...  lm_head_prune31.pt   (16 files)

The number in each filename is the layer index the head attaches to, not a training step and not a pruning ratio.

Vocabulary sizes are padded up from the tokenizers' native values (Llama 3: 128256; Mistral: 32000 → 32768).

Format

Each .pt file is an OrderedDict with a single key:

{"weight": Tensor(vocab_size, 4096)}

No bias, no optimizer state, no metadata.

Usage

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B")

LAYER = 24
sd = torch.load(f"lm_head_checkpoints_L/lm_head_prune{LAYER}.pt", map_location="cpu")

head = torch.nn.Linear(4096, sd["weight"].shape[0], bias=False)
head.load_state_dict(sd)

# Read out logits from layer LAYER's hidden state
out = model(input_ids, output_hidden_states=True)
logits = head(out.hidden_states[LAYER])

Sweep LAYER from 16 to 31 and compare perplexity or task accuracy against the full model to find the depth at which quality degrades unacceptably.

Match the set to the base model — the heads are not interchangeable between L and M, and the vocab dimensions differ.