DecompressionLM

Zero-shot concept graph extraction from LLMs using deterministic arithmetic sampling,
to discover what language models encode in the form of concept graphs

Test Your Model

How to

prompt = """Generate United States bar exam legal concepts as keywords.

Please output ONE concept per line.
Each concept can be multiple words if needed.
Do not include explanations or extra text.
Please begin from any random concept.
Please use English.
"""

results = estimate_prefix_mass(
    model=model,  # by default, you can load the model in torch.bfloat16
    tokenizer=tokenizer,
    prefix=prompt,
    prefix_len=16,  # please leave this as-is; it is not used since it is disabled by prob_threshold=1.0
    max_len=16,
    max_samples=8192,
    prob_threshold=1.0,  # disabled, as it is not used in this paper
    use_chat_template=True,
    batch_size=256,
    display_interval=256,  # please set this to batch_size or a multiple of it if sampling is fast
    save_path=(
        "results/api_call_example/"
        "Qwen2.5-7B-Instruct__BF16_BASE__us_law_bar_meta__L16__T1.0.delm.parquet"
    ),
    model_name="Qwen2.5-7B-Instruct__BF16_BASE",  # used for logging only; set as you see fit
    enable_graph_analysis=True,
)

Please keep in mind that, because of our custom implementation of sampling, some models with different KV cache specifications might not work out of the box and may throw confusing errors. In the worst case, you might need to modify the code in src/arithmetic.py around line 96, where the variable is_gemma can be used to disable the KV cache when set to True. This overrides KV cache issues but may be slower.

What you will get

================================================================================
GRAPH ANALYSIS
================================================================================
Concept Extraction:
Total concepts extracted   : 25933
Unique before merging      : 1621
Unique after merging       : 985
Total raw edges            : 17741
Graph Statistics:
Nodes (concepts)           : 809
Edges (relations)          : 3273
Graph density              : 0.005007098361257634
Average degree             : 8.091470951792337
Node Connectivity:
Orphan nodes (degree=0)    : 0 (0.0%)
Weakly connected (deg≤1)   : 198 (24.5%)
Well connected (deg>1)     : 611 (75.5%)
Graph Structure:
Connected components       : 6
Largest component          : 788 nodes (97.4%)
Largest component density  : 0.005251904359548243
================================================================================
📄 Loaded from file (not recomputed)
================================================================================
Interactive viewer ready. Type 'help' for commands.
>>> concepts
================================================================================
FINAL GRAPH CONCEPTS (after filtering & merging)
================================================================================
Total concepts extracted   : 25933
Unique before merging      : 1621
Unique after merging       : 985
Merge reduction            : 39.2%
Final graph nodes          : 809
Top 64 most connected concepts:

(degree=247) contractlaw
(degree=190) realproperty
(degree=189) intellectualproperty
(degree=166) evidence
(degree=152) criminallaw
(degree=132) constitutionallaw
(degree=129) criminalprocedur
(degree=121) familylaw
(degree=116) legal ethic
(degree=113) habeas corpu
(degree=104) tort
(degree= 95) professionalresponsibility
(degree= 92) civ procedure
(degree= 90) statutesoflimitation
(degree= 85) equity
(degree= 81) eminentdomain
(degree= 79) tortlaw
(degree= 71) evidencelaw
(degree= 65) propertylaw
(degree= 60) contract
(degree= 60) liability
(degree= 56) mens rea
(degree= 54) negligence
(degree= 53) appellate procedure
(degree= 49) mergers and acquisition
(degree= 49) habitual resident
(degree= 49) tortfeasor
(degree= 49) strict liability
(degree= 48) contract formation
(degree= 47) real estate
(degree= 46) equitable estoppel
(degree= 45) statutory interpretation
(degree= 45) standing
(degree= 44) realestatelaw
(degree= 44) attorneyclientprivilege
(degree= 43) equitable remedy
(degree= 42) due proces
(degree= 41) trustsandestate
(degree= 41) conflict of law
(degree= 40) evidentiary rule
(degree= 40) lien
(degree= 39) ethic
(degree= 39) equal protection
(degree= 36) causation
(degree= 35) willsandtrust
(degree= 33) administrative law
(degree= 32) realpropertylaw
(degree= 32) mortgage
(degree= 31) moot court
(degree= 31) habitualoffender
(degree= 30) estoppel
(degree= 30) res judicata
(degree= 28) hearsay
(degree= 28) corporatelaw
(degree= 28) legal malpractice
(degree= 27) property
(degree= 25) constitution
(degree= 25) business organization
(degree= 25) jurisdiction
(degree= 24) mers system
(degree= 24) estate planning
(degree= 24) mental capacity
(degree= 23) evidencerule
(degree= 23) statute of fraud
================================================================================

This is an example of what you could get when testing a model on US bar exam concepts, with the interactive viewer mentioned below.

Experiments

To run the experiments as in the paper:

python -m exp.01_quantization

To view a saved sampling file:

python -m tools.view_results results/qwen2p5_7b_it_quant_ladder__us_law_16_newlines/Qwen2.5-7B-Instruct__AWQ_4BIT__us_law_bar_meta__L16__T1.0.delm.parquet

To reproduce the experiments reported in our paper, the scripts you will need are in the exp/ directory:

1. 01_quantization_{16,32}.py – Quantization ladder (BF16→AWQ-4bit→GPTQ-Int4)
2. 02_quantization_offset.py – VdC offset robustness (8 runs per variant)
3. 03_perplexity_16.py – Concept understanding via perplexity
4. 04_mmlu.py – Hallucination verification (22 models ranked by MMLU-Pro Law)

Other scripts in this exp/ directory helps to print out the results in a more readable, more compact, form of tables.

Each individual experiment will be saved in the results/ directory once it finishes. It is therefore fine if a script is interrupted midway; when rerun, it will skip experiments that have already completed.

Some models may be gated, so please remember to log in to your Hugging Face CLI (hf auth).

If you pipe the output with tee, many ANSI characters that are rewritten over time (for plotting) will be saved as well. You can use the sed command and regex rules below to clean it up.

On Logging

The printout is very long, so it is a good idea to pipe stdout into a file. If you pipe the output into a text file, it will contain ANSI sequences, including plots that are drawn over time. To debloat the log file:

sed -r 's/\x1B\[[0-9;]*[A-Za-z]//g' exp-01.log > exp-01.clean.txt

This reduces the file size, but the plot colors will be lost.

To remove ANSI sequences entirely (plots will be removed) and repeated sampling statistics:

sed -r '/^\x1B\[/d; /^\[SAMPLING\]/,+3d' exp-02.log > exp-02.clean.txt

Alternatively, you can rerun the script after all experiments finish to obtain only the final results.

Installation

pip install -r requirements.txt
pip install gptqmodel>=1.0.0

• Because of a dependency mismatch between autoawq (support for newer PyTorch pipelines was recently dropped) and gptqmodel (actively maintained), specifically regarding transformers, please run pip install -r requirements.txt first and then pip install gptqmodel. You may see a dependency warning indicating that transformers is outdated, but gptqmodel will still be installed and should work for now.

Citation

@misc{hong2026decompressionlmdeterministicdiagnosticzeroshot,
      title={DecompressionLM: Deterministic, Diagnostic, and Zero-Shot Concept Graph Extraction from Language Models}, 
      author={Zhaochen Hong and Jiaxuan You},
      year={2026},
      eprint={2602.00377},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2602.00377}
}