Skip to content

Latest commit

 

History

4 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Deferred Entity Commitment in Language Model Factual Recall

This repository contains code, processed results, and final figures for the paper:

Relation Before Entity: Deferred Commitment in Language Model Factual Recall Accepted as a virtual poster at the Mechanistic Interpretability Workshop at ICML 2026.

Paper / OpenReview: https://openreview.net/forum?id=gcgoa7EN4v

The project studies whether relation information and entity/argument information become generation-controlling at different stages during factual recall in decoder-only language models. The core finding is that relation information becomes causally active at the final prediction token before entity information does. Entity information is available earlier at the subject/entity token but becomes generation-controlling at the final token only later in the model. We call this deferred entity commitment.


Repository structure

final_paper/
  figures/                        Final paper figures (PDF and PNG)
  model_running_scripts/          Full model-running scripts for four experiments (GPU required)
  results/                        Processed CSV outputs used for paper figures and tables
  scripts/                        Lightweight analysis and figure-generation scripts (no GPU)
  README.md                       This file
  requirements.txt                Minimal Python dependencies

Quickest path to reproducing paper figures

The paper figures can be regenerated from the included processed CSVs without loading any language models:

python scripts/make_paper_figures.py

Output is written to figures/. The main figure files are:

figures/fig2_transfer_curves.pdf
figures/fig3_both_change_competition.pdf
figures/fig4_subject_vs_last_patching.pdf
figures/fig5_steering_asymmetry.pdf

PNG versions are also included for convenience.

Note on fig1: figures/fig1_combined_concept_onset.pdf is included in the repository but does not appear in the submitted paper. It is a conceptual/schematic figure generated during development and kept for reference.


Models

All experiments use four decoder-only language models:

meta-llama/Llama-3.2-3B
meta-llama/Meta-Llama-3-8B
Qwen/Qwen2.5-3B
microsoft/phi-2

Some models require accepting HuggingFace license terms before downloading. Set LOCAL_FILES_ONLY = True in any model-running script to force local cache and skip network access.


Experiments

Experiment 1: Relation vs entity transfer

Script: model_running_scripts/run_exp1_relation_entity_transfer.py Results: results/exp1_relation_entity_transfer/

Final-token activation patching comparing:

  1. Entity transfer: same relation, different argument/entity.
  2. Relation transfer: same argument/entity, different relation.

Also runs unrelated-random diagnostics and wrong-entity relation diagnostics.

Key result files:

entity_transfer_summary.csv
relation_transfer_summary.csv
relation_transfer_by_pair.csv
relation_wrong_entity_summary.csv
unrelated_random_entity_summary.csv
unrelated_random_relation_summary.csv
all_models_ci.csv

Pair/family-balanced onset analysis, generated by scripts/compute_pair_balanced_onsets.py and used for Table 1 and Appendix C:

onset_pair_balanced_check.csv              all thresholds 0.2-0.5
onset_pair_balanced_threshold_04.csv       primary threshold 0.4
relation_pair_balanced_curve.csv
entity_family_balanced_curve.csv
full_pair_family_balanced_aggregate_curve.csv

Experiment 2: Both-change relation/entity competition

Script: model_running_scripts/run_exp2_both_change_and_controls.py Results: results/exp2_both_change/

Both relation and entity change simultaneously. The experiment asks which signal dominates at each layer. Outputs are classified as: original retained, relation wins, entity wins, or mixed/other.

Key result files:

both_change_summary.csv
both_change_by_pair_summary.csv
both_change_crossover_table.csv
both_change_all_models_raw.csv

Control and diagnostic files:

unrelated_ctrl_all_models_raw.csv        overwrite diagnostic
noise_ctrl_all_models_raw.csv            true negative control
self_patch_ctrl_all_models_raw.csv       hook sanity check
alternate_donor_all_models_raw.csv       robustness check

Unrelated-donor overwrite diagnostic summaries, used in Appendix D:

unrelated_donor_by_layer_summary.csv
unrelated_donor_peak_entity_overwrite.csv
unrelated_donor_peak_relation_wins.csv

Note on controls: The unrelated-donor condition is an overwrite diagnostic, not a pure no-effect control. It tests whether arbitrary donor states produce structured wins. The noise-patch condition is the true negative control. The self-patch condition is a hook sanity check. The alternate-donor condition is a robustness check.


Experiment 3: Subject-token vs final-token patching

Script: model_running_scripts/run_exp3_subject_token_patching.py Results: results/exp3_subject_token_patching/

Tests whether entity information is causally effective at the subject/entity token earlier than at the final token, by comparing:

  1. Final-token patch: donor final-token hidden state → recipient final-token position.
  2. Subject/entity-token patch: donor subject/entity-token hidden state → recipient subject/entity-token position.
  3. Subject/entity-average patch: average over donor subject/entity-token positions.

Key result files:

subject_patch_all_models_raw.csv
subject_patch_model_summary.csv
subject_patch_by_family_summary.csv

Experiment 4: Layer-zone steering

Script: model_running_scripts/run_exp4_steering.py Results: results/exp4_steering/

Tests whether relation and entity steering directions are effective in different layer zones, especially mid versus late layers. Each real steering direction is compared against a matched-norm random direction baseline.

Steering conditions:

relation_mid
entity_mid
entity_late
relation_late

Key result files:

steering_all_models_raw.csv
steering_all_models_summary.csv

Analysis scripts

The following scripts do not require GPU access because they operate on the included processed CSV files.

Pair-balanced onset check

python scripts/compute_pair_balanced_onsets.py

Recomputes Experiment 1 onsets from saved CSVs using pair-balanced relation averaging and family-balanced entity averaging. Tests thresholds 0.2, 0.3, 0.4, and 0.5. Produces onset tables used for Table 1 and Appendix C.

Unrelated-donor overwrite diagnostic

python scripts/summarize_unrelated_donor_overwrite.py

Summarizes the Experiment 2 unrelated-donor diagnostic. Tests whether unrelated relation-wins remain near zero while late overwrite sensitivity is explained by the broadly overwrite-sensitive late regime.

Task generation audit

python scripts/audit_task_generation.py

Prints greedy-generation and logit statistics for the task banks. Audit CSVs are included in results/task_generation_audit/ but are not required for reproducing the main paper figures.


Full model reruns

To rerun all full model experiments:

python model_running_scripts/run_exp1_relation_entity_transfer.py
python model_running_scripts/run_exp2_both_change_and_controls.py
python model_running_scripts/run_exp3_subject_token_patching.py
python model_running_scripts/run_exp4_steering.py

Full reruns require GPU access and substantial runtime. Runtime depends heavily on the GPU, model, and experiment.


Installation

Install the minimal Python dependencies:

pip install torch transformers pandas numpy matplotlib tqdm

For CUDA-compatible PyTorch, follow the installation instructions at:

https://pytorch.org/get-started/locally/

Optional 4-bit quantization support:

pip install bitsandbytes

4-bit loading is disabled by default with:

USE_4BIT = False

Enable it in the model-running scripts if GPU memory is limited.


Notes

  • This repository accompanies the accepted Mechanistic Interpretability Workshop at ICML 2026 paper listed above.
  • We use "entity" broadly to mean the prompt argument: a country, verb, noun, adjective, or element depending on the task family.
  • Some scripts use "subject token" in filenames or variable names because the original experiments were written using subject/entity-token terminology during development.
  • Onset threshold 0.4 is used in Table 1 of the main paper. Full sensitivity across thresholds 0.2-0.5 is included in Appendix C and in onset_pair_balanced_check.csv.
  • This is research code intended to make the experimental logic and results transparent. It is not a polished software package.
  • The fastest reproducibility path is python scripts/make_paper_figures.py, which regenerates all paper figures from the included processed CSVs.

Citation

If you use this repository or build on the paper, please cite:

Divyansh Agarwal. Relation Before Entity: Deferred Commitment in Language Model Factual Recall.
Mechanistic Interpretability Workshop at ICML 2026, 2026.
@inproceedings{
agarwal2026relation,
title={Relation Before Entity: Deferred Commitment in Language Model Factual Recall},
author={Divyansh Agarwal},
booktitle={Mechanistic Interpretability Workshop at ICML 2026},
year={2026},
url={https://openreview.net/forum?id=gcgoa7EN4v}
}

About

Code and processed results for deferred entity commitment experiments

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages