repro.py

import os
import sys
import math
import random
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from transformers import GPT2Config, GPT2LMHeadModel
import numpy as np
import matplotlib.pyplot as plt
from sklearn.metrics import f1_score
from tqdm import tqdm

# Set seeds
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}", flush=True)

# Helper to count parameters
def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

class SyntheticDataset(Dataset):
    def __init__(self, num_samples, seq_len, vocab_size):
        self.data = torch.randint(0, vocab_size, (num_samples, seq_len))
        
    def __len__(self):
        return len(self.data)
        
    def __getitem__(self, idx):
        return self.data[idx]

def train_model(model, train_loader, epochs=150, lr=2e-3, desc=""):
    model.train()
    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    
    for epoch in range(epochs):
        total_loss = 0
        for batch in train_loader:
            if isinstance(batch, (list, tuple)):
                batch = batch[0]
            batch = batch.to(device)
            outputs = model(batch, labels=batch)
            loss = outputs.loss
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            
        avg_loss = total_loss / len(train_loader)
        if (epoch + 1) % 25 == 0 or epoch == 0:
            print(f"  [{desc}] Epoch {epoch+1:3d}/{epochs}: Loss = {avg_loss:.6f}", flush=True)
            
        # Early stopping if fully overfitted (loss < 1e-4)
        if avg_loss < 1e-4:
            print(f"  [{desc}] Early stopping at epoch {epoch+1} (Loss = {avg_loss:.6f})", flush=True)
            break
            
    return model

def compute_dataset_entropy(N, S, V):
    return N * S * math.log2(V)

def compute_model_description_length(model, data_loader):
    model.eval()
    total_nll_bits = 0.0
    with torch.no_grad():
        for batch in data_loader:
            if isinstance(batch, (list, tuple)):
                batch = batch[0]
            batch = batch.to(device)
            outputs = model(batch, labels=batch)
            num_elements = batch.numel()
            nll_bits = (outputs.loss.item() * num_elements) / math.log(2.0)
            total_nll_bits += nll_bits
    return total_nll_bits

def run_synthetic_experiments():
    print("\n--- Running Claim 1 (Synthetic Capacity) ---", flush=True)
    vocab_size = 2048
    seq_len = 64
    
    # 2 layers, 2 heads, 64 embedding size
    config = GPT2Config(
        vocab_size=vocab_size,
        n_positions=seq_len,
        n_ctx=seq_len,
        n_embd=64,
        n_layer=2,
        n_head=2,
        bos_token_id=0,
        eos_token_id=0
    )
    
    model = GPT2LMHeadModel(config)
    num_params = count_parameters(model)
    print(f"Model parameters: {num_params}", flush=True)
    
    dataset_sizes = [50, 200, 500, 1000, 1500, 2000, 2500]
    memorization_results = []
    
    for N in dataset_sizes:
        print(f"\n[Claim 1] Training on dataset size N={N}...", flush=True)
        dataset = SyntheticDataset(N, seq_len, vocab_size)
        loader = DataLoader(dataset, batch_size=64, shuffle=True)
        
        m = GPT2LMHeadModel(config).to(device)
        m = train_model(m, loader, epochs=150, lr=2e-3, desc=f"Synth N={N}")
        
        H_X = compute_dataset_entropy(N, seq_len, vocab_size)
        eval_loader = DataLoader(dataset, batch_size=64, shuffle=False)
        H_K_X_given_theta = compute_model_description_length(m, eval_loader)
        
        mem = H_X - H_K_X_given_theta
        mem = max(0.0, mem)
        bpp = mem / num_params
        memorization_results.append((N, mem, bpp))
        print(f"  [Claim 1] N={N} | Entropy H(X): {H_X:.1f} bits | NLL H^K: {H_K_X_given_theta:.1f} bits", flush=True)
        print(f"  [Claim 1] N={N} | Memorized: {mem:.1f} bits ({bpp:.4f} bits/param)", flush=True)
        
    max_mem_entry = max(memorization_results, key=lambda x: x[1])
    capacity = max_mem_entry[1]
    alpha = max_mem_entry[2]
    print(f"\n[Claim 1] Estimated Model Capacity: {capacity:.1f} bits", flush=True)
    print(f"[Claim 1] Estimated alpha (bits per parameter): {alpha:.4f}", flush=True)
    
    Ns = [r[0] for r in memorization_results]
    mems = [r[1] for r in memorization_results]
    bpps = [r[2] for r in memorization_results]
    
    plt.figure(figsize=(10, 4))
    plt.subplot(1, 2, 1)
    plt.plot(Ns, mems, 'o-')
    plt.xlabel('Dataset Size (N)')
    plt.ylabel('Memorized Bits')
    plt.title('Total Memorization')
    plt.grid(True)
    
    plt.subplot(1, 2, 2)
    plt.plot(Ns, bpps, 'o-', color='orange')
    plt.axhline(y=alpha, color='r', linestyle='--', label=f'Alpha: {alpha:.3f}')
    plt.xlabel('Dataset Size (N)')
    plt.ylabel('Bits Per Parameter (bpp)')
    plt.title('Memorization Capacity (bpp)')
    plt.legend()
    plt.grid(True)
    
    os.makedirs('outputs', exist_ok=True)
    plt.savefig('outputs/synthetic_capacity.png')
    plt.close()
    
    return num_params, capacity, alpha

def run_text_experiments(num_params, capacity):
    print("\n--- Running Claim 2 (Double Descent) and Claim 3 (Membership Inference) ---", flush=True)
    words = ["the", "quick", "brown", "fox", "jumps", "over", "lazy", "dog", "model", "parameter", "memorize",
             "generalize", "capacity", "data", "training", "entropy", "compress", "information", "theory",
             "double", "descent", "grokking", "sigmoid", "scaling", "law", "inference", "privacy", "copyright"]
    for i in range(972):
        words.append(f"word_{i}")
        
    vocab_size = len(words)
    seq_len = 64
    
    def generate_text_sample(length):
        sentence = []
        for _ in range(length):
            if random.random() < 0.2:
                sentence.append(random.choice([0, 1, 2, 3]))
            else:
                sentence.append(random.randint(0, vocab_size-1))
        return torch.tensor(sentence)
        
    N_sizes = [50, 200, 500, 1000, 1500, 2000, 3000]
    test_size = 200
    
    test_data = torch.stack([generate_text_sample(seq_len) for _ in range(test_size)])
    test_dataset = torch.utils.data.TensorDataset(test_data)
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
    
    config = GPT2Config(
        vocab_size=vocab_size,
        n_positions=seq_len,
        n_ctx=seq_len,
        n_embd=64,
        n_layer=2,
        n_head=2,
        bos_token_id=0,
        eos_token_id=0
    )
    
    train_losses = []
    test_losses = []
    f1_scores = []
    
    for N in N_sizes:
        print(f"\n[Claim 2&3] Training on text dataset size N={N}...", flush=True)
        train_data = torch.stack([generate_text_sample(seq_len) for _ in range(N)])
        train_dataset = torch.utils.data.TensorDataset(train_data)
        train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
        
        m = GPT2LMHeadModel(config).to(device)
        m = train_model(m, train_loader, epochs=150, lr=2e-3, desc=f"Text N={N}")
        
        m.eval()
        train_loss = 0.0
        with torch.no_grad():
            for batch in train_loader:
                batch = batch[0].to(device)
                outputs = m(batch, labels=batch)
                train_loss += outputs.loss.item()
        train_loss /= len(train_loader)
        train_losses.append(train_loss)
        
        test_loss = 0.0
        with torch.no_grad():
            for batch in test_loader:
                batch = batch[0].to(device)
                outputs = m(batch, labels=batch)
                test_loss += outputs.loss.item()
        test_loss /= len(test_loader)
        test_losses.append(test_loss)
        
        train_losses_individual = []
        with torch.no_grad():
            for i in range(N):
                sample = train_data[i:i+1].to(device)
                loss = m(sample, labels=sample).loss.item()
                train_losses_individual.append(loss)
                
        test_losses_individual = []
        with torch.no_grad():
            for i in range(test_size):
                sample = test_data[i:i+1].to(device)
                loss = m(sample, labels=sample).loss.item()
                test_losses_individual.append(loss)
                
        all_losses = np.array(train_losses_individual + test_losses_individual)
        labels = np.array([1]*N + [0]*test_size)
        
        best_f1 = 0.0
        thresholds = np.percentile(all_losses, np.linspace(0, 100, 100))
        for t in thresholds:
            preds = (all_losses < t).astype(int)
            f1 = f1_score(labels, preds)
            if f1 > best_f1:
                best_f1 = f1
                
        f1_scores.append(best_f1)
        print(f"  [Claim 2&3] N={N} | Train Loss: {train_loss:.4f} | Test Loss: {test_loss:.4f} | Membership F1: {best_f1:.4f}", flush=True)
        
    plt.figure(figsize=(10, 4))
    plt.subplot(1, 2, 1)
    plt.plot(N_sizes, train_losses, 'o-', label='Train Loss')
    plt.plot(N_sizes, test_losses, 'o-', label='Test Loss')
    plt.axvline(x=capacity / (seq_len * math.log2(vocab_size)), color='r', linestyle='--', label='Est. Capacity')
    plt.xlabel('Dataset Size (N)')
    plt.ylabel('Loss')
    plt.title('Train/Test Losses vs Dataset Size')
    plt.legend()
    plt.grid(True)
    
    c1, c2, c3 = 1.34, -0.034, -33.14
    pred_f1s = []
    for N in N_sizes:
        ratio = capacity / N
        sig = 1.0 / (1.0 + math.exp(-(c2 * (ratio + c3))))
        pred = 0.5 * (1.0 + c1 * sig)
        pred_f1s.append(pred)
        
    plt.subplot(1, 2, 2)
    plt.plot(N_sizes, f1_scores, 'o-', label='Empirical F1')
    plt.plot(N_sizes, pred_f1s, 'x--', label='Paper Scaling Law')
    plt.xlabel('Dataset Size (N)')
    plt.ylabel('Membership F1')
    plt.title('Membership Inference F1')
    plt.legend()
    plt.grid(True)
    
    plt.savefig('outputs/text_results.png')
    plt.close()
    
    print("\nScaling Law verification results:", flush=True)
    for N, emp, pred in zip(N_sizes, f1_scores, pred_f1s):
        print(f"  N={N:4d} | Empirical F1: {emp:.4f} | Predicted F1: {pred:.4f} | Diff: {abs(emp - pred):.4f}", flush=True)

if __name__ == "__main__":
    num_params, capacity, alpha = run_synthetic_experiments()
    run_text_experiments(num_params, capacity)