JetBrains-Research/learned-transfer-attack

Learned Transfer Membership Inference Attack

A classifier that detects whether a given text was part of a language model's fine-tuning data. It compares the output distributions of a fine-tuned model against its pretrained base, extracting per-token features that a small transformer classifier uses to predict membership. Trained on 10 transformer models × 3 text domains, it generalizes zero-shot to unseen model/dataset combinations, including non-transformer architectures (Mamba, RWKV, RecurrentGemma).

Usage

Install

git clone https://github.com/JetBrains-Research/ltmia.git
cd ltmia
pip install -e .

Inference

import torch
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, AutoModelForCausalLM
from ltmia import extract_per_token_features_both, create_mia_model

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 1. Load your base and fine-tuned models
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

model_ref = AutoModelForCausalLM.from_pretrained("gpt2").to(device)
model_tgt = AutoModelForCausalLM.from_pretrained("./my-finetuned-gpt2").to(device).eval()

# 2. Extract features
texts = ["Text you want to check...", "Another text..."]

feats, masks, _ = extract_per_token_features_both(
    model_tgt, model_ref, tokenizer, texts,
    device=device, batch_size=8, sequence_length=128, k=20,
)

# 3. Load the MIA classifier
ckpt_path = hf_hub_download(
    repo_id="JetBrains-Research/learned-transfer-attack",
    filename="mia_combined_400k.pt",
)
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
mia = create_mia_model(
    architecture=ckpt["architecture"],
    d_in=ckpt["d_in"],
    seq_len=ckpt.get("seq_len", 128),
    **ckpt["mia_hparams"],
)
mia.load_state_dict(ckpt["state_dict"])
mia.to(device).eval()

# 4. Predict membership
with torch.no_grad():
    logits = mia(
        torch.from_numpy(feats).to(device),
        torch.from_numpy(masks).to(device),
    )
    probs = torch.sigmoid(logits)

for text, p in zip(texts, probs):
    prob = p.item()
    label = "MEMBER" if prob > 0.5 else "NON-MEMBER"
    print(f"[{prob:.4f}] {label}  ←  {text[:80]}")

You need black-box query access (full vocabulary logits) to both the fine-tuned model and its pretrained base. sequence_length=128 and k=20 must match this checkpoint. See the GitHub repository for CLI tools, training your own classifier, and evaluation scripts.

Model Details

Architecture: Transformer encoder — 154→112 projection, 3 layers, 4 heads, FFN 224, attention pooling, ~340K parameters.

Input: Per-token features (shape N × 128 × 154) comparing logits, ranks, and losses between target and reference models.

Output: Membership probability per text (sigmoid of scalar logit).

Training data: Features from 10 transformers (DistilGPT-2, GPT-2-XL, Pythia-1.4B, Cerebras-GPT-2.7B, GPT-J-6B, Gemma-2B, Qwen2-1.5B, MPT-7B, Falcon-RW-1B, Falcon-7B) fine-tuned on 3 datasets (News Category, Wikipedia, CNN/DailyMail). 18K samples per combination, 540K total.

Training: AdamW, lr 5e-4, batch 16384, 100 epochs. Checkpoint selected by best validation AUC.

Evaluation (Out-of-Distribution)

Performance on models and datasets never seen during classifier training:

Architecture	Model	Dataset	AUC
Transformer	GPT-2	AG News	0.945
Transformer	Pythia-2.8B	AG News	0.911
Transformer	Mistral-7B	XSum	0.989
Transformer	LLaMA-2-7B	AG News	0.948
Transformer mean	(7 models × 4 datasets)		0.908
State-space	Mamba-2.8B	AG News	0.969
State-space	Mamba-2.8B	WikiText	0.995
Linear attention	RWKV-3B	AG News	0.976
Linear attention	RWKV-3B	XSum	0.998
Gated recurrence	RecurrentGemma-2B	AG News	0.924
Gated recurrence	RecurrentGemma-2B	XSum	0.988
Non-transformer mean	(3 models × 4 datasets)		0.957

Transfer to code (Swallow-Code): 0.865 mean AUC despite training only on natural language.

License

MIT