HGRN-1.3B, 80% semi-structured pruned (EHWS)

This is fla-hub/hgrn-1.3B-100B pruned to 80% weight sparsity with EHWS (Extreme Hierarchical Weight Sparsity), a two-phase ADMM method that trains directly against the true next-token loss under a hard sparsity constraint rather than matching each layer's dense output on calibration data. It extends ELSA (The Unseen Frontier: Pushing the Limits of LLM Sparsity with Surrogate-Free ADMM), splitting ELSA's single global optimization into a cheap per-layer warm start (Phase 1) followed by one global fine-tune (Phase 2).

"Semi-structured" here means every layer is pruned to the same 80% ratio; within a layer, which individual weights survive is still unstructured.

Results

WikiText2 PPL C4 PPL Zero-shot avg (7-task)
Dense 11.84 16.89 0.435
This checkpoint (80% sparsity) 88.23 56.79 0.350

Achieved sparsity: 0.7999. Zero-shot accuracy is the standard 0-shot 7-task average (ARC-Easy/Challenge, BoolQ, HellaSwag, OpenBookQA, RTE, Winogrande) via lm-evaluation-harness.

For reference, HGRN-1.3B isn't one of the models ELSA's own paper covers, so we ran ELSA's official code ourselves on this exact model to get a same-model baseline: 54.08 / 36.92 WikiText2/C4 -- notably better, but under a different setup: unstructured (not semi-structured) sparsity, pure cross-entropy training with no KD term. This checkpoint keeps a CE + knowledge-distillation objective (against the original dense model) active throughout training, which is a real difference in training signal, not just an unlucky hyperparameter gap.

Training setup

Two-phase ADMM: Phase 1 warm-starts each layer sequentially up to 70% sparsity, then Phase 2 optimizes all layers jointly under one shared forward pass up to the final 80% target (128 rounds x 32 steps). Objective is 0.5 * CE + 0.5 * KD against the original dense model, both on unlabeled text. Calibration/training data is 2048 sequences x 2048 tokens from C4.

Key Phase 2 hyperparameters: lr=6e-4, ADMM penalty lambda=0.01 (constant schedule), adam_beta2=0.95, alpha_kd=0.5, bf16. These were tuned by directly auditing ELSA's own reference implementation run against this same model -- notably, the ADMM lambda used in earlier attempts had been about 200x too small for this model, and correcting it alone cut WikiText2 PPL by roughly 40%.

Usage

This architecture needs the flash-linear-attention package (registers HGRNForCausalLM with transformers' Auto* classes) and trust_remote_code=True:

import fla  # noqa: registers HGRN with transformers' Auto* classes
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained(
    "ikimyaii/HGRN-1.3B-semi-structured-EHWS-80pct",
    trust_remote_code=True,
    torch_dtype="bfloat16",
)
tokenizer = AutoTokenizer.from_pretrained("fla-hub/hgrn-1.3B-100B")

Note that "semi-structured" here refers to the sparsity pattern being uniform across layers, not to a hardware-friendly block/N:M pattern -- this checkpoint won't show a runtime speedup without a sparse-aware kernel. It's released for quality evaluation, not as a deployment-ready accelerated checkpoint.

Method code

github.com/KimiaGholami/EHWS

Downloads last month
226
Safetensors
Model size
1B params
Tensor type
BF16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for ikimyaii/HGRN-1.3B-semi-structured-EHWS-80pct

Finetuned
(1)
this model

Paper for ikimyaii/HGRN-1.3B-semi-structured-EHWS-80pct