This repository contains the implementation of SAMA (Subset-Aggregated Membership Attack) and baseline methods for evaluating membership inference vulnerabilities in Diffusion Language Models (DLMs).
Overview of SAMA
. ├── trainer/ # DLM training module │ ├── model/ # Model architectures │ ├── configs/ # Training configurations │ └── train.py # Main training script ├── attack/ # MIA implementation module │ ├── attacks/ # Attack implementations (SAMA, baselines) │ ├── configs/ # Attack configurations │ └── run.py # Main attack script └── dataset/ # Dataset preparation utilities # Create conda environment conda create -n sama python=3.8 conda activate sama # Install dependencies pip install -r trainer/requirements.txt pip install -r attack/requirements.txtBefore running any experiments, set the following environment variables to match your platform. You can export them in your shell, add them to your ~/.bashrc/~/.zshrc, or place them in a .env file and source it.
# Base directory where prepared datasets are stored. export SAMA_DATASET_PATH="/path/to/your/datasets" # Directory to save SAMA attack metadata (step-level scores, subset details). export SAMA_METADATA_DIR="/path/to/metadata/output" # Weights & Biases project and group names for training run tracking. export SAMA_WANDB_PROJECT="Diff_LLM" export SAMA_WANDB_GROUP="SAMA" # HuggingFace token for accessing gated/private models. export HF_TOKEN="hf_..."The table below summarizes every environment variable, where it is read, and its default if unset:
| Variable | File(s) | Default |
|---|---|---|
SAMA_DATASET_PATH | trainer/configs/prep.py, dataset/prep.py, dataset/prep_mimir.py | ./ |
SAMA_METADATA_DIR | attack/attacks/sama.py | ./ |
SAMA_WANDB_PROJECT | trainer/configs/prep.py | Diff_LLM |
SAMA_WANDB_GROUP | trainer/configs/prep.py | SAMA |
HF_TOKEN | attack/attacks/sama.py, attack/attacks/ratio.py | empty |
Both scripts default their output directory to $SAMA_DATASET_PATH (set in Step 0).
# Download and prepare MIMIR benchmark datasets python dataset/prep_mimir.py # Prepare standard NLP datasets (wikitext, ag_news, xsum) python dataset/prep.pyYou can override the output directory or other defaults per-run:
# Custom output directory and sample count python dataset/prep.py --output-dir /tmp/data --num-samples 5000 # Prepare only specific MIMIR subsets python dataset/prep_mimir.py --datasets arxiv pile_cc --split-name ngram_13_0.8Run python dataset/prep.py --help or python dataset/prep_mimir.py --help for all options.
See trainer/README.md for further details.
python python -m attack.run \ -c attack/configs/config_all.yaml \ --output ./attack_results \ --base-dir /path/to/your/trained/modelsKey arguments for attack.run:
| Argument | Description |
|---|---|
-c, --config | Path to attack configuration YAML (required) |
--output | Directory to save results and metadata (required) |
--base-dir | Base directory for resolving relative model/dataset paths in config (default: ./) |
--target-model | Path to target model, overrides config if provided |
--lora-path | Path to LoRA adapter, overrides config if provided |
--seed | Random seed (default: 42) |
Edit attack/configs/config_all.yaml or create your own attack configurations:
sama: steps: 16 # Progressive masking steps min_mask_frac: 0.05 # Starting mask fraction max_mask_frac: 0.50 # Ending mask fraction num_subsets: 128 # Subsets per step subset_size: 10 # Tokens per subset batch_size: 8 # Batch size loss: module: loss mc_num: 4 loss-calibration: module: "ratio" reference_model_path: "GSAI-ML/LLaDA-8B-Base" reference_device: "cuda" ...@article{chen2026membership, title={Membership Inference Attacks Against Fine-tuned Diffusion Language Models}, author={Chen, Yuetian and Zhang, Kaiyuan and Du, Yuntao and Stoppa, Edoardo and Fleming, Charles and Kundu, Ashish and Ribeiro, Bruno and Li, Ninghui}, journal={arXiv preprint arXiv:2601.20125}, year={2026} }