SelfIE Labels

Label Generator

The label generator is a neural network that maps SAE decoder vectors to natural language descriptions using a novel soft prompt approach:

# Simplified forward pass
def generate_label(sae_decoder_vector):
    # 1. Project SAE vector to soft token
    soft_token = projection_layer(sae_decoder_vector)

    # 2. Inject into template
    template = "This pattern activates when: <soft_token>"

    # 3. Generate label using base LLM
    label = base_llm.generate(template_with_soft_token)

    return label

The projection layer is the only trainable component - the base LLM remains frozen. This makes training efficient and leverages the LLM's pre-existing language capabilities.

Reflective Coherence Reward

The reward function measures how well a generated label activates its target feature:

def compute_reward(generated_label, target_feature_id, sae, base_model):
    # Generate text using the label as a prompt
    prompts = [
        f"Write text that heavily features {generated_label}",
        f"Generate content with lots of {generated_label}",
        f"Create text emphasizing {generated_label}"
    ]

    activations = []
    for prompt in prompts:
        # Generate text
        text = base_model.generate(prompt)

        # Get SAE activations on generated text
        hidden_states = base_model.get_hidden_states(text)
        feature_acts = sae.encode(hidden_states)

        # Extract target feature activation
        activations.append(feature_acts[:, target_feature_id].mean())

    # Reward is mean activation of target feature
    reward = torch.stack(activations).mean()
    return reward

High rewards indicate the label successfully captures what makes the feature activate.

Projection Layer

The projection layer transforms SAE decoder vectors into soft prompt embeddings:

s = W · d + b

Where:

• d ∈ ℝ^d_model = SAE decoder vector (e.g., 8192-dim for Llama 3.1 8B)
• W ∈ ℝ^{d_model × d_model} = Learnable projection matrix
• b ∈ ℝ^d_model = Learnable bias vector (the "universal interpretation direction")
• s ∈ ℝ^d_model = Soft prompt token embedding

Bias Vector: Universal Interpretation

The bias vector b is particularly important. It represents a universal interpretation direction learned across all features:

• Feature-specific: α · d captures what makes this feature unique
• Universal: b captures common patterns across all features (e.g., "this activates when...")
• Transfer learning: The bias vector can transfer between models and SAEs

GRPO Training Objective

The reinforcement learning phase uses Group Relative Policy Optimization (GRPO):

Maximize: E[reward(label)] - β * KL(π || π_ref)

Where:
- reward(label) = mean activation of target feature
- π = current policy (label generator)
- π_ref = reference policy (pretrained checkpoint)
- β = KL penalty coefficient (balances exploration)

This encourages generating labels that activate features while staying close to the pretrained distribution.

Cross-Model Label Generation

SelfIE can use a larger model to generate labels for a smaller model's SAE:

• SAE Model: Llama 3.1 8B (whose features we want to label)
• Label Generator: Llama 3.3 70B (generates better descriptions)
• Projection: 8192-dim → 8192-dim (non-square matrices supported)

The larger model's superior language abilities often produce clearer, more nuanced labels.

DiffMean Bias Training

Recent work trains the bias vector using steering vectors from contrastive text pairs:

1. 1. Generate diverse conversational properties (e.g., "formal language", "asking for clarification")
2. 2. Create contrastive text pairs with/without each property
3. 3. Compute DiffMean steering vectors from activation differences
4. 4. Train bias vector on (steering vector, description) dataset

This allows training SelfIE without needing SAE features at all, using pure steering vectors as the training data.

Log-Scale Parameterization

The scale parameter α uses logarithmic parameterization for better training dynamics:

# Instead of: scale = α (standard)
# Use: scale = exp(log_scale) (log parameterization)

Δscale ≈ scale × learning_rate × gradient

# This makes updates proportional to current scale
# Same LR works across scale ranges: 0.1 to 100+

This is particularly important when training across multiple layers, as different layers may need vastly different optimal scales.