jrsh/sd-control-and-representation

Stable Diffusion Control and Representation

This repository contains Sparse Autoencoder (SAE) models trained on Stable Diffusion representations. These models enable interpretable analysis and control of the diffusion process by learning sparse feature decompositions of the model's internal activations.

Purpose

Sparse Autoencoders decompose neural network activations into interpretable, sparse feature representations. These SAE models allow you to:

• Analyze which features activate during image generation
• Control the generation process by steering specific features
• Interpret what the diffusion model has learned at different timesteps
• Debug and understand model behavior

Repository Structure

The repository is organized by model version and experiment type:

README.md
sd-v1-5/
└── sae/
    ├── exp32_topk16/
    │   ├── config.json
    │   ├── model.pt
    │   └── merged_feature_sums.pt
    ├── exp36_topk32/
    │   ├── config.json
    │   ├── model.pt
    │   └── merged_feature_sums.pt
    └── exp36_topk64/
        ├── config.json
        ├── model.pt
        └── merged_feature_sums.pt

Additional model versions and architectures will be added in separate directories.

Available Models

sd-v1-5/sae/

• exp32_topk16: Top-K=16, experiment 32 configuration
• exp36_topk32: Top-K=32, experiment 36 configuration
• exp36_topk64: Top-K=64, experiment 36 configuration

Files in Each Model

• config.json: SAE architecture configuration (dimensions, hyperparameters)
• model.pt: Trained SAE weights (encoder and decoder)
• merged_feature_sums.pt: Aggregated feature activation statistics across training data

Usage

Installation

pip install huggingface_hub torch

Loading a Model

from huggingface_hub import hf_hub_download
import torch
import json

# Choose your model
REPO_ID = "jrsh/sd-control-and-representation"
SUBFOLDER = "sd-v1-5/sae/exp36_topk32"  # Choose: exp32_topk16, exp36_topk32, or exp36_topk64

# 1. Load configuration
config_path = hf_hub_download(
    repo_id=REPO_ID,
    subfolder=SUBFOLDER,
    filename="config.json"
)
with open(config_path, 'r') as f:
    config = json.load(f)

# 2. Load model weights
model_path = hf_hub_download(
    repo_id=REPO_ID,
    subfolder=SUBFOLDER,
    filename="model.pt"
)
sae_state_dict = torch.load(model_path, map_location='cpu')

# 3. Load feature statistics 
feature_sums_path = hf_hub_download(
    repo_id=REPO_ID,
    subfolder=SUBFOLDER,
    filename="merged_feature_sums.pt"
)
feature_sums = torch.load(feature_sums_path, map_location='cpu')

print(f"Successfully loaded SAE from {SUBFOLDER}")
print(f"Config: {config}")

Related Resources

• Code Repository: https://github.com/jrsh44/sd-control-and-representation