Theory explaining how optimality conditions structure SAE (Sparse Autoencoder) dictionaries - hierarchical splitting, absorption, residuals, and dense antipodal features
Install with Codex or Claude Copy this prompt, paste it into Codex, Claude, or another assistant, and let it review the skill page and install it for you.
A direct command skips the review prompt. Inspect the source before running it.
Theory explaining how optimality conditions structure SAE (Sparse Autoencoder) dictionaries - hierarchical splitting, absorption, residuals, and dense antipodal features
SAE Optimality Structures provides a theoretical foundation for understanding what Sparse Autoencoders (SAEs) extract from neural representations. By analyzing local optimality conditions without assuming specific data-generating models, this theory explains observed SAE behaviors: hierarchical splitting, feature absorption, residual structure, and dense antipodal features.
Key Contribution: SAEs have empirical success in parsing neural representations into interpretable concepts, but lack theoretical grounding for what constitutes a "concept". This work derives constraints that any optimal SAE dictionary must satisfy, explaining phenomena through L1 regularization and nonnegativity interactions with data distributions.
Use When:
Understanding SAE feature extraction mechanics
Analyzing hierarchical splitting in learned dictionaries
Definition: Extends Gribonval & Schnass (2010) local optimality conditions to nonnegative joint optimization problem that vanilla SAEs approximate.
Key Insight: Instead of assuming sparse independent feature data models (which poorly approximate LLM representations), directly analyze what properties any optimal dictionary must satisfy.
Mathematical Setup:
Objective: minimize ||x - Dz||² + λ||z||₁
Subject to: z ≥ 0 (nonnegativity)
D ≥ 0 (dictionary nonnegativity)
Where:
x = input representation
D = dictionary matrix (features)
z = sparse code (activations)
λ = L1 regularization strength
2. Optimality Constraints
Constraint 1: Feature Distribution Relationship
Optimal features must satisfy:
D_i ⟂ (X - DZ) Z_i^T ≥ λ|Z_i| (for all features i)
Where:
D_i = ith dictionary feature
X = data matrix
Z_i = ith feature's activation vector
λ = regularization parameter
Interpretation: Features must be "just sparse enough" - activations balance reconstruction error and L1 penalty.
Constraint 2: Nonnegativity Structure
With nonnegativity, optimal solutions have:
D_ij = 0 iff feature i never activates for datapoint j
Z_ji > 0 iff feature i reconstructs j with positive contribution
Phenomenon: Features split into sub-features as dictionary size increases.
Explanation from Optimality:
As |D| increases:
1. Previously optimal single feature D_parent splits into D_child1, D_child2
2. Split condition: ||x - D_parent||² > ||x - D_child1||² + ||x - D_child2||² - λ(||z_child||₁ - ||z_parent||₁)
3. Split occurs when reconstruction gain exceeds sparsity penalty increase
Visual Example:
Parent feature (broad): Child features (specialized):
D_parent = [0.3, 0.3] D_child1 = [0.4, 0.1] (specialized to dim1)
D_child2 = [0.1, 0.4] (specialized to dim2)
Split when data has sub-clusters requiring specialized reconstruction
B. Feature Absorption
Phenomenon: Small features "absorbed" into larger features, disappearing from dictionary.
Explanation:
Absorption occurs when:
1. Feature D_small has low activation rate |Z_small|/n < threshold
2. Its reconstruction contribution is covered by D_large
3. Removing D_small improves objective: λ|Z_small| > ||X - DZ||² contribution
Absorption condition:
||D_small * Z_small||² < λ||Z_small||₁
Interpretation: L1 penalty "kills" rarely-used features, consolidating them into frequently-activated ones.
C. Residual Structure
Phenomenon: Reconstruction residuals have non-random structure, not pure noise.
Explanation:
Residual R = X - DZ satisfies:
1. R ⟂ D (orthogonal to dictionary)
2. ||R||² > λ for non-absorbed features
3. R has interpretable structure: components not captured by current features
Residual analysis reveals:
- Missing feature directions
- Feature interactions not modeled
- Hierarchical organization gaps
Implication: Residuals guide where new features should be added.
D. Dense Antipodal Features
Phenomenon: Some features appear dense (high activation) yet interpretable.
Explanation:
Antipodal pairs: D_a ≈ -D_b (but nonnegativity constraint forces both positive)
For data x = αD_a + βD_b + noise:
Optimal encoding: z_a = α, z_b = β (both dense)
Why interpretable?
1. D_a and D_b represent opposite directions in semantic space
2. High activation = strong presence of that semantic direction
3. Dense ≠ uninterpretable if direction is meaningful
Example in LLM representations:
D_positive_sentiment: activates for happy, joy, good (dense)
D_negative_sentiment: activates for sad, anger, bad (dense)
Both interpretable despite high activation frequency
Implementation Methodology
Phase 1: Optimality Analysis
Step 1: Compute Local Optimality Conditions
defcheck_local_optimality(D, Z, X, lambda_reg):
"""
Verify if current (D, Z) satisfies local optimality conditions
Returns:
- is_optimal: bool
- violation_details: dict of constraint violations
"""# Constraint 1: Feature distribution
reconstruction_error = X - D @ Z
feature_constraints = {}
for i inrange(D.shape[1]):
# Check: D_i ⟂ (X - DZ) Z_i^T ≥ λ|Z_i|
gradient_Di = reconstruction_error @ Z[i]
penalty_term = lambda_reg * np.abs(Z[i]).sum()
feature_constraints[i] = {
'gradient_norm': np.linalg.norm(gradient_Di),
'penalty': penalty_term,
'satisfied': np.linalg.norm(gradient_Di) >= penalty_term
}
# Constraint 2: Nonnegativity
nonneg_D = (D >= 0).all()
nonneg_Z = (Z >= 0).all()
is_optimal = all(c['satisfied'] for c in feature_constraints.values()) and nonneg_D and nonneg_Z
return {
'is_optimal': is_optimal,
'feature_constraints': feature_constraints,
'nonnegativity_satisfied': {'D': nonneg_D, 'Z': nonneg_Z}
}
Step 2: Detect Hierarchical Splitting
defdetect_feature_splitting(D_large, D_small, Z_large, Z_small, X, lambda_reg):
"""
Analyze if features in D_small are splits of D_large features
Split detection criteria:
1. Child features reconstruct better than parent
2. Sparsity penalty increase is compensated
3. Child features specialize on sub-clusters
"""
splits_detected = []
for parent_idx inrange(D_large.shape[1]):
D_parent = D_large[:, parent_idx]
# Find potential children (similar direction, more specialized)
potential_children = []
for child_idx inrange(D_small.shape[1]):
D_child = D_small[:, child_idx]
# Similarity measure
similarity = cosine_similarity(D_parent, D_child)
if similarity > 0.7: # High similarity threshold
potential_children.append(child_idx)
iflen(potential_children) >= 2:
# Check split condition
children_features = D_small[:, potential_children]
children_codes = Z_small[potential_children]
# Reconstruction comparison
parent_reconstruction = D_parent @ Z_large[parent_idx]
children_reconstruction = children_features @ children_codes
parent_error = np.linalg.norm(X - parent_reconstruction)
children_error = np.linalg.norm(X - children_reconstruction)
# Sparsity penalty comparison
parent_penalty = lambda_reg * np.abs(Z_large[parent_idx]).sum()
children_penalty = lambda_reg * np.abs(children_codes).sum()
# Split beneficial?
split_beneficial = (parent_error - children_error) > (children_penalty - parent_penalty)
if split_beneficial:
splits_detected.append({
'parent_idx': parent_idx,
'child_indices': potential_children,
'reconstruction_gain': parent_error - children_error,
'sparsity_cost': children_penalty - parent_penalty,
'net_benefit': (parent_error - children_error) - (children_penalty - parent_penalty)
})
return splits_detected
Step 3: Analyze Residual Structure
defanalyze_residual_structure(D, Z, X):
"""
Analyze reconstruction residual structure to find missing features
Key analyses:
1. Residual magnitude per dimension
2. Residual clustering (potential new feature directions)
3. Residual-feature orthogonality verification
"""
residual = X - D @ Z
# 1. Magnitude per dimension
residual_magnitude = np.linalg.norm(residual, axis=0)
high_residual_dims = np.where(residual_magnitude > residual_magnitude.mean())[0]
# 2. Clustering analysis (PCA on residuals)from sklearn.decomposition import PCA
pca = PCA(n_components=10)
residual_components = pca.fit_transform(residual.T)
# Check if residual components are orthogonal to D
residual_directions = pca.components_
orthogonal_to_D = []
for direction in residual_directions:
orthogonality = np.abs(D.T @ direction).max()
if orthogonality < 0.1: # Nearly orthogonal
orthogonal_to_D.append(direction)
# 3. Potential new features
potential_new_features = {
'residual_magnitude': residual_magnitude,
'high_residual_dims': high_residual_dims,
'residual_components': residual_components,
'orthogonal_directions': orthogonal_to_D,
'suggested_features': len(orthogonal_to_D)
}
return potential_new_features
Phase 2: Large-Dictionary Convex Problem
Step 1: Construct Convex Relaxation
defconstruct_large_dictionary_problem(X, n_atoms, lambda_reg):
"""
Construct convex problem for large dictionary limit
Key insight: As |D| → n_datapoints, problem becomes convex
(atom-per-datapoint limit)
"""
n_samples, n_features = X.shape
# Convex relaxation: each datapoint gets dedicated atom# Objective becomes: minimize ||X - DZ||² + λ||Z||₁# with Z diagonal (one atom per point)# This simplifies to per-point optimization:
convex_solution = {}
for i inrange(n_samples):
x_i = X[i]
# Optimal atom for x_i: D_i = x_i (reconstruction perfect)# Optimal code: z_i = 1 if ||x_i||² > λ, else 0if np.linalg.norm(x_i) ** 2 > lambda_reg:
convex_solution[i] = {
'atom': x_i,
'code': 1.0,
'active': True
}
else:
convex_solution[i] = {
'atom': np.zeros(n_features),
'code': 0.0,
'active': False
}
# Active atoms count
n_active = sum(s['active'] for s in convex_solution.values())
return {
'solution': convex_solution,
'n_active_atoms': n_active,
'sparsity': n_active / n_samples,
'convex_limit': True
}
Step 2: Explore Wide Atom Limit
defexplore_wide_atom_limit(X, lambda_reg_values):
"""
Explore behavior as n_atoms → ∞ (wide dictionary limit)
Phenomena to observe:
1. Sparsity saturation
2. Feature specialization
3. Hierarchical depth
"""
results = {}
for lambda_reg in lambda_reg_values:
# Large dictionary convex solution
convex_sol = construct_large_dictionary_problem(X, X.shape[0], lambda_reg)
# Compute properties
results[lambda_reg] = {
'active_ratio': convex_sol['sparsity'],
'avg_feature_norm': np.mean([
np.linalg.norm(s['atom']) for s in convex_sol['solution'].values()
if s['active']
]),
'feature_specialization': compute_specialization_metric(convex_sol),
'hierarchical_depth': estimate_hierarchy_depth(convex_sol)
}
return results
defcompute_specialization_metric(convex_solution):
"""
Measure how specialized features are (vs. general/broad)
Specialization = low overlap between features
"""
active_atoms = [s['atom'] for s in convex_solution['solution'].values() if s['active']]
iflen(active_atoms) < 2:
return0.0# Compute pairwise cosine similarityfrom sklearn.metrics.pairwise import cosine_similarity
similarity_matrix = cosine_similarity(active_atoms)
# Specialization = 1 - mean similarity (excluding self-similarity)
n = len(active_atoms)
mean_similarity = (similarity_matrix.sum() - n) / (n * (n - 1))
return1 - mean_similarity
defestimate_hierarchy_depth(convex_solution):
"""
Estimate hierarchical tree depth from feature structure
Approximation: clustering depth of active atoms
"""
active_atoms = [s['atom'] for s in convex_solution['solution'].values() if s['active']]
iflen(active_atoms) < 10:
return1# Hierarchical clusteringfrom sklearn.cluster import AgglomerativeClustering
clustering = AgglomerativeClustering(n_clusters=None, distance_threshold=0.3)
labels = clustering.fit_predict(active_atoms)
# Depth = max cluster levelreturn clustering.n_clusters_
Phase 3: Validation & Interpretation
Step 1: Validate SAE Behaviors
defvalidate_sae_phenomena(D, Z, X, lambda_reg):
"""
Validate if observed SAE behaviors match optimality theory predictions
Checks:
1. Hierarchical splitting matches data sub-clusters
2. Absorbed features were low-activation
3. Residuals have interpretable structure
4. Dense features are antipodal pairs
"""
validation_results = {}
# 1. Splitting validation
splits = detect_feature_splitting(D[:, :D.shape[1]//2], D, Z[:D.shape[1]//2], Z, X, lambda_reg)
validation_results['splitting'] = {
'n_splits': len(splits),
'avg_reconstruction_gain': np.mean([s['reconstruction_gain'] for s in splits]) if splits else0,
'matches_subclusters': verify_subcluster_match(splits, X)
}
# 2. Absorption validation
activation_rates = np.abs(Z).sum(axis=1) / X.shape[0]
low_activation_features = np.where(activation_rates < 0.05)[0]
validation_results['absorption'] = {
'n_low_activation': len(low_activation_features),
'absorption_candidates': low_activation_features,
'would_be_absorbed': [
idx for idx in low_activation_features
if np.linalg.norm(D[:, idx] @ Z[idx]) ** 2 < lambda_reg * np.abs(Z[idx]).sum()
]
}
# 3. Residual structure validation
residual_analysis = analyze_residual_structure(D, Z, X)
validation_results['residual'] = {
'has_structure': len(residual_analysis['orthogonal_directions']) > 0,
'n_potential_features': residual_analysis['suggested_features'],
'interpretable': check_interpretability(residual_analysis['orthogonal_directions'])
}
# 4. Dense antipodal features
dense_features = np.where(activation_rates > 0.3)[0]
antipodal_pairs = find_antipodal_pairs(D[:, dense_features])
validation_results['antipodal'] = {
'n_dense': len(dense_features),
'n_antipodal_pairs': len(antipodal_pairs),
'pairs': antipodal_pairs
}
return validation_results
deffind_antipodal_pairs(D_subset, similarity_threshold=-0.7):
"""
Find features that are antipodal (opposite directions)
"""
pairs = []
n = D_subset.shape[1]
for i inrange(n):
for j inrange(i+1, n):
similarity = cosine_similarity(D_subset[:, i], D_subset[:, j])
if similarity < similarity_threshold: # Negative = antipodal
pairs.append((i, j, similarity))
return pairs
defcheck_interpretability(directions):
"""
Heuristic: interpretable directions have low entropy activation patterns
"""ifnot directions:
returnFalse# Simplified: assume interpretable if directions cluster wellfrom sklearn.cluster import KMeans
try:
kmeans = KMeans(n_clusters=min(5, len(directions)))
labels = kmeans.fit_predict(directions)
# Low cluster variance = interpretable
silhouette = silhouette_score(directions, labels)
return silhouette > 0.5except:
returnFalse
Technical Pitfalls
Pitfall 1: Nonnegativity Violations
Problem: Standard SAE implementations allow negative activations
Solution: Enforce strict nonnegativity via projection
# Correct approachdefenforce_nonnegativity(Z):
"""Project activations to nonnegative space"""return np.maximum(Z, 0)
# In training loop
Z = enforce_nonnegativity(Z)
D = enforce_nonnegativity(D)
Pitfall 2: L1 Regularization Scaling
Problem: λ not scaled to data magnitude → incorrect sparsity
Solution: Normalize λ relative to ||X||²
