Task Two: Model Tracer

The implementation identifies which generative AI model created a given image. It can distinguish between 8 different model variants from two families (RAR and VAR) and detect outlier images from unknown sources.

Overview

The initial approach is in two steps:

1. Family Attribution: Determine if an image belongs to the RAR family, VAR family, or is an outlier
2. Variant Attribution: If from a known family, identify the specific model variant (e.g., var16, var20, rarxl, etc.)

Extraly, a one-step approach, PRADA, is also implemented.

Architecture

Core Components

1. Inverse Decoder Training (inv_decoder.py)

Fine-tune the encoder for each model family that maps generated images back to their quantized feature maps. Augmentation is defined by (augmentation.py)

Training Process:

• Uses images generated by all variants within a family (e.g., var16, var20, var24, var30)
• Freezes the decoder and quantizer (only trains the encoder)
• Loss: MSE between predicted latent codes and ground truth latent codes
• Output: rar_inverse_decoder.pth and var_inverse_decoder.pth

2. Provenance Detection (provenance.py)

The main inference pipeline that classifies test images.

Two-Stage:

Stage 1 - Family Attribution (Provenance Scores):

• QuantLoss: Measures how well the image fits the model's codebook (Equation 5)
  • Uses optimized quantization (quantizer.py) to find best-fit tokens
  • Lower loss = better fit to the codebook
• EncLoss: Measures encoding stability (Equation 10)
  • Ratio of reconstruction error vs. re-encoding error
  • Unstable encoding suggests image wasn't generated by this model
• Provenance Score: QuantLoss × EncLoss (Equation 11)
  • Three thresholding strategies
  • If score exceeds threshold → likely not from this family
  • If both RAR and VAR thresholds are passed (not very correct):
    • Compare QuantLoss values and assign to family with lower QuantLoss (better codebook fit)

Stage 2 - Variant Attribution:

• Once family is determined, compare perplexity across model variants
• Several feature engineering (feature_extractor.py) are explored, including feature_extractor/get_perplexity() and variant_classifier.py.

3. Optimized Quantization (quantizer.py)

Implements Algorithm 3 (optimized/differentiable quantization for VAR models using gradient-based optimization) in Appendix A.

Challenge: VAR uses multi-scale residual quantization, which is discrete (non-differentiable)

Solution:

• Converts discrete token selection into soft assignment using logits
• Optimizes logits to minimize reconstruction error

4. PRADA Framework (flat_prada.py)

Model Families

VAR (Visual Autoregressive)

• Multi-scale residual quantization (1×1, 2×2, ..., 16×16)
• 4 variants: var16, var20, var24, var30 (different transformer depths)
• Special handling: First scale treated as conditioning token
• Input construction requires idxBl_to_var_input() for proper residual accumulation
• Input images requires normalization to [-1,1]

RAR (Randomized Autoregressive)

• Single-scale quantization (16×16 = 256 tokens)
• 4 variants: rarb, rarl, rarxl, rarxxl (different model sizes)
• Uses permutation-based training (tokens shuffled during training)
• Prepends condition token, uses causal attention
• Condition requires preprocessing

File Structure

.
├── inv_decoder.py              # Train inverse decoders
├── augmentation.py             # Augmentations for training 
├── submission.py               # Main inference pipeline
├── provenance.py               # Two-step approach
├── feature_extractor.py        # Feature engineering
├── feature_plotter.py          # Diagnosis of feature engineering
├── flat_prada.py               # One-step approach (PRADA framework)
├── variant_classifier.py       # Variant attribution
├── quantizer.py                # VAR optimized quantization
├── run_inverse_decoder.sh      # SLURM script for training
├── run_submission.sh           # SLURM script for inference
├── plot_inv_dec_loss_train.py  # Plot training curves
├── Dataset/
│   ├── train/                  # Training images by model
│   │   ├── rarb/, rarl/, rarxl/, rarxxl/
│   │   └── var16/, var20/, var24/, var30/
│   ├── val/                    # Validation images
│   └── test/                   # Test images (unlabeled)
├── RAR/                        # RAR model code
│   └── tokenizer_1d/
├── VAR/                        # VAR model code
├── logs/                       # Training/inference logs
└── plots/                      # Saved plots
    └── prada_diagnostics/

Submissions

Two versions of submission.csv are provided. One is produced by variant attribution with get_perplexity, and the other is with MIA features. Both are supposed to achieve 75% - 80% accuracy.

Notes

• In the implementation of VAR models, there is one more operation quant_conv between encoder and quantize. Similarly, there is one more operation post_quant_conv between quantize and decoder.