Natively Unlearnable LLMs

Official code for NULLs — an architecture for training natively unlearnable language models.

The core idea is to route memorization of individual training sequences into sequence-specific MLP neurons (the "memory sinks"), while keeping a shared pool of "generalization" neurons that learn transferable capabilities. Because each fact is memorized into a deterministic, sequence-keyed subset of neurons, a fact can later be unlearned simply by masking out the neurons tied to the sequence that taught it — without retraining and without degrading general performance.

The model is a litgpt GPT variant (GPTSeqTD) trained with PyTorch Lightning Fabric. Experiments target SmolLM2-style configurations.

---

How it works

A standard MLP is replaced by LLaMAMLPSeqTD. Its intermediate neurons are partitioned into two groups (see src/src/SeqTDModel.py):

• Generalization neurons — a fraction p_gen of the intermediate width. Always active, for every token of every sequence.
• Memory neurons — the remaining 1 - p_gen. For a given training sequence, only a deterministic pseudo-random subset of them is unmasked. Each memory neuron is active with probability p_mem, and which neurons are active is a function of the sequence's integer ID (seq_id), computed by a hash (batch_seqtied_mask_mult). Two different documents therefore memorize into largely disjoint neuron subsets.

Every training document is assigned a unique seq_id that flows through the forward pass alongside the tokens. Attention also respects document boundaries via a seq_ids-derived block-diagonal mask, so packed sequences don't attend across documents.

At evaluation time the MLP supports several eval_modes:

eval_mode	Behavior
all	All memory neurons on (scaled by p_mem) — the model with all knowledge available.
dropout	All memory neurons off — generalization-only behavior.
activate_seq	Activate only the memory neurons for given seq_ids, optionally excluding the neurons of exclude_seq_ids.

Passing a fact's sequence id to exclude_seq_ids is the mechanism for native unlearning: the neurons that memorized that fact are switched off while everything else is left untouched.

---

Repository layout

.
├── environment.yml          # Conda environment (name: semdedup)
├── package-list.txt         # Full pinned conda/pip package list
└── MemSinks/
    ├── src/                 # Installable `src` package (the model + data + utils)
    │   ├── setup.py
    │   └── src/
    │       ├── SeqTDModel.py            # GPTSeqTD, LLaMAMLPSeqTD, sequence-tied masking
    │       ├── SeqTDConfig.py           # Model config (SmolLM2 presets + SeqTD fields)
    │       ├── SeqTDTrainingArguments.py
    │       ├── SeqTDDataLoader.py
    │       ├── seqtd_single_ds.py       # SingleSeqTDID streaming dataset (per-seq ids)
    │       ├── seq_td_with_id.py        # Multi-split / weighted-mixture data module
    │       ├── seqtd_data.py
    │       ├── seqtd_multi_eval.py      # Cluster-wise / multi-split validation
    │       ├── SeqTDGenerate.py         # Generation with seq-tied masking
    │       └── seqtd_single_ds_with_rep.py
    ├── train/               # Training entry points + YAML configs
    │   ├── train_memsinks.py            # Pretrain a MemSinks model
    │   ├── train_memsinks_wiki.py       # MemSinks pretraining on the Wikipedia setup
    │   ├── train_mixture.py             # Train on a weighted mixture of splits
    │   ├── train_standard_with_memsinks_data.py  # Standard (no-sink) baseline
    │   ├── count_tokens.py
    │   └── configs/
    │       ├── memsinks/                # MemSinks configs (mlp_class_name: LLaMAMLPSeqTD)
    │       └── standard/                # Baseline configs (mlp_class_name: LLaMAMLP)
    ├── data_preparation/    # Build + tokenize the training/eval corpora
    │   ├── generate_wikipedia_qa.py     # LLM-generated QA pairs from Wikipedia
    │   └── tokenize/
    │       └── tokenize_wikipedia_topic.py
    ├── eval/wiki_facts/      # Factual-knowledge evaluation
    │   ├── wiki_generate_cloze_probes.py     # Cloze probes w/ rephrasings + distractors
    │   ├── generate_cloze_from_sentences*.py
    │   ├── compute_truth_ratio*.py           # Truth-ratio metric over probes
    │   ├── better_cluster/                   # Duplicate / outlier analysis (semhash)
    │   └── launchers/                        # SLURM launchers for the eval pipeline
    └── new_launcher/         # Experiment orchestration (SLURM + optional GCS)
        ├── setup_env.py                      # Central env/path config -> Project.config
        └── wiki_launchers/                   # Wikipedia tokenize + MemSinks launchers

---

Setup

Requires Python 3.11 and a CUDA-capable GPU.

# 1. Create the environment
conda env create -f environment.yml
conda activate semdedup        # the environment is named `semdedup`

# 2. Install the model package (editable)
cd MemSinks/src
pip install -e .

Key pinned dependencies (see package-list.txt for the complete list): litgpt==0.5.10, torch==2.6.0, lightning, transformers==4.51.3, datasets==3.5.0, litdata, jsonargparse, sentence-transformers, semhash==0.4.1, wandb.

---

Configuration & paths

This code is environment-agnostic: no absolute paths are hardcoded. Two mechanisms supply environment-specific locations.

1. Training/eval YAML configs (MemSinks/train/configs/**) use /path/to/... placeholders for data, general_val_data, out_dir, and tokenizer_dir. Edit them to point at your tokenized data and tokenizer before running a config directly.

2. Launchers (MemSinks/new_launcher, MemSinks/eval/.../launchers) read all paths from Project.config, populated by new_launcher/setup_env.py. Every value has an environment-variable override and a sensible default:

Env var	Purpose	Default
MEMSINKS_DATA_BASE_DIR	Base directory for data	$HOME/<project>/data
MEMSINKS_CODE_DIR	Path to this MemSinks checkout	$HOME/<project>/MemSinks
MEMSINKS_ARTIFACT_PATH	Checkpoint directory	$HOME/<project>/ckpt
MEMSINKS_TOKENIZER_PATH	Tokenizer directory	$HOME/tokenizers
MEMSINKS_CONDA_PATH / _ENV	Conda activation for SLURM jobs	$HOME/miniconda3/.../conda.sh
MEMSINKS_CLUSTER / _PARTITION / _QOS / _SLURM_ACCOUNT	SLURM submission defaults	babel / general / normal_qos
MEMSINKS_GCS_BUCKET, MEMSINKS_USE_GCS	Optional GCS sync of data/ckpts	(disabled)

Populate the project config once:

cd MemSinks/new_launcher
# export any MEMSINKS_* overrides you need first
python setup_env.py

The launchers (new_launcher, the eval launchers/) depend on an external experiments orchestration library (from experiments import Project, ...). The model, training scripts, and eval scripts under src/, train/, and eval/wiki_facts/*.py run without it.

---

Usage

Train

Run a training script against a YAML config (after filling in its paths):

cd MemSinks/train

# MemSinks model (memory sinks enabled)
python train_memsinks.py --config configs/memsinks/memsinks_initial.yaml

# Standard baseline (no memory sinks)
python train_standard_with_memsinks_data.py --config configs/standard/standard_360M.yaml

The MemSinks knobs live under train: in the config — p_gen (fraction of generalization neurons) and p_mem (per-memory-neuron activation probability); the example config uses p_gen: 0.3, p_mem: 0.7. The architecture is selected via model_config.mlp_class_name (LLaMAMLPSeqTD for MemSinks vs. LLaMAMLP for the baseline).

Prepare data

cd MemSinks/data_preparation
python generate_wikipedia_qa.py --help            # build QA pairs from Wikipedia
python tokenize/tokenize_wikipedia_topic.py --help # tokenize into litdata chunks

Evaluate factual knowledge

cd MemSinks/eval/wiki_facts
python wiki_generate_cloze_probes.py --help   # generate cloze probes (+ rephrasings)
python compute_truth_ratio.py --help          # truth-ratio metric for a checkpoint

To measure unlearning, evaluate a checkpoint with eval_mode="activate_seq" and pass the target fact's sequence id(s) via exclude_seq_ids, then compare the truth ratio against the all baseline.

---

Citation

If you use this code, please cite the Memorization Sinks work.

@misc{ghosal2026nativelyunlearnablelargelanguage,
      title={Natively Unlearnable Large Language Models}, 
      author={Gaurav R. Ghosal and Pratyush Maini and Aditi Raghunathan},
      year={2026},
      eprint={2606.13873},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2606.13873}, 
}