Hugging Face's logo

JetBrains-Research
/
learned-transfer-attack

Other
PyTorch
ltmia
membership-inference-attack
privacy
security
language-models
Model card Files
xet
Community
learned-transfer-attack
File size: 4,357 Bytes 
14666fc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
 
---
license: mit
tags:
  - membership-inference-attack
  - privacy
  - security
  - language-models
  - pytorch
pipeline_tag: other
library_name: ltmia
---

# 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

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

### Inference

```python
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](https://github.com/JetBrains-Research/ltmia) 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