chore: update TransformerLens submodule and remove MixCoder references

- Updated the TransformerLens submodule to the latest commit.
- Removed all references to `MixCoder` from the codebase, including imports and related configurations.
- Adjusted `BaseSAEConfig` and `SAEConfig` to eliminate `mixcoder` from the `sae_type` options.
- Cleaned up the `FeatureAnalyzer` class by removing modality-specific metrics related to `MixCoder`.

Author: Frankstein73

TransformerLens

@@ -13,7 +13,6 @@
     InitializerConfig,
     LanguageModelConfig,
     LLaDAConfig,
-    MixCoderConfig,
     MongoDBConfig,
     SAEConfig,
     TrainerConfig,
@@ -22,7 +21,6 @@
 from .crosscoder import CrossCoder
 from .database import MongoClient
 from .evaluator import EvalConfig, Evaluator
-from .mixcoder import MixCoder
 from .resource_loaders import load_dataset, load_model
 from .runners import (
     AnalyzeCrossCoderSettings,
@@ -83,8 +81,6 @@
     "FeatureAnalyzerConfig",
     "MongoDBConfig",
     "MongoClient",
-    "MixCoderConfig",
-    "MixCoder",
     "AnalyzeCrossCoderSettings",
     "analyze_crosscoder",
     "AutoInterpSettings",

src/lm_saes/__init__.py

@@ -10,7 +10,6 @@
 from lm_saes.abstract_sae import AbstractSparseAutoEncoder
 from lm_saes.config import FeatureAnalyzerConfig
 from lm_saes.crosscoder import CrossCoder
-from lm_saes.mixcoder import MixCoder
 from lm_saes.utils.discrete import KeyedDiscreteMapper
 from lm_saes.utils.distributed import DimMap
 from lm_saes.utils.misc import is_primary_rank
@@ -245,18 +244,6 @@ def analyze_chunk(
         max_feature_acts = torch.zeros((d_sae_local,), dtype=sae.cfg.dtype, device=sae.cfg.device)
         mapper = KeyedDiscreteMapper()
 
-        if isinstance(sae, MixCoder):
-            act_times_modalities = {
-                k: torch.zeros((d_sae_local,), dtype=torch.long, device=sae.cfg.device) for k in sae.cfg.modality_names
-            }
-            max_feature_acts_modalities = {
-                k: torch.zeros((d_sae_local,), dtype=sae.cfg.dtype, device=sae.cfg.device)
-                for k in sae.cfg.modality_names
-            }
-        else:
-            act_times_modalities = None
-            max_feature_acts_modalities = None
-
         # Process activation batches
         for batch in activation_stream:
             # Reshape meta to zip outer dimensions to inner
@@ -295,15 +282,6 @@ def analyze_chunk(
             act_times += feature_acts.gt(0.0).sum(dim=[0, 1])
             max_feature_acts = torch.max(max_feature_acts, feature_acts.max(dim=0).values.max(dim=0).values)
 
-            if isinstance(sae, MixCoder):
-                assert act_times_modalities is not None and max_feature_acts_modalities is not None
-                for i, k in enumerate(sae.cfg.modality_names):
-                    feature_acts_modality = feature_acts * (batch["modalities"] == i).long().unsqueeze(-1)
-                    act_times_modalities[k] += feature_acts_modality.gt(0.0).sum(dim=[0, 1])
-                    max_feature_acts_modalities[k] = torch.max(
-                        max_feature_acts_modalities[k], feature_acts_modality.max(dim=0).values.max(dim=0).values
-                    )
-
             # TODO: Filter out meta that is not string
             discrete_meta = {
                 k: torch.tensor(mapper.encode(k, v), device=sae.cfg.device, dtype=torch.int32) for k, v in meta.items()
@@ -339,8 +317,6 @@ def analyze_chunk(
             max_feature_acts=max_feature_acts,
             sample_result=sample_result,
             mapper=mapper,
-            act_times_modalities=act_times_modalities,
-            max_feature_acts_modalities=max_feature_acts_modalities,
             device_mesh=device_mesh,
         )
 
@@ -352,8 +328,6 @@ def _format_analysis_results(
         max_feature_acts: torch.Tensor,
         sample_result: dict[str, dict[str, torch.Tensor]],
         mapper: KeyedDiscreteMapper,
-        act_times_modalities: dict[str, torch.Tensor] | None = None,
-        max_feature_acts_modalities: dict[str, torch.Tensor] | None = None,
         device_mesh: DeviceMesh | None = None,
     ) -> list[dict[str, Any]]:
         """Format the analysis results into the final per-feature format.
@@ -365,8 +339,6 @@ def _format_analysis_results(
             max_feature_acts: Tensor of maximum activation values for each feature
             sample_result: Dictionary of sampling results
             mapper: MetaMapper for encoding/decoding metadata
-            act_times_modalities: Optional dictionary of activation times per modality (for MixCoder)
-            max_feature_acts_modalities: Optional dictionary of maximum activation values per modality (for MixCoder)
 
         Returns:
             List of dictionaries containing per-feature analysis results
@@ -466,16 +438,6 @@ def _format_analysis_results(
                 feature_result["decoder_similarity_matrix"] = decoder_similarity_matrices[i, :, :].tolist()
                 feature_result["decoder_inner_product_matrix"] = decoder_inner_product_matrices[i, :, :].tolist()
 
-            # Add modality-specific metrics for MixCoder
-            if (
-                isinstance(sae, MixCoder)
-                and act_times_modalities is not None
-                and max_feature_acts_modalities is not None
-            ):
-                feature_result["act_times_modalities"] = {k: v[i].item() for k, v in act_times_modalities.items()}
-                feature_result["max_feature_acts_modalities"] = {
-                    k: v[i].item() for k, v in max_feature_acts_modalities.items()
-                }
 
             results.append(feature_result)
 

src/lm_saes/analysis/feature_analyzer.py

@@ -48,10 +48,10 @@ class BaseSAEConfig(BaseModelConfig, ABC):
     """
     Base class for SAE configs.
     Initializer will initialize SAE based on config type.
-    So this class should not be used directly but only as a base config class for other SAE variants like SAEConfig, MixCoderConfig, CrossCoderConfig, etc.
+    So this class should not be used directly but only as a base config class for other SAE variants like SAEConfig, CrossCoderConfig, etc.
     """
 
-    sae_type: Literal["sae", "crosscoder", "mixcoder"]
+    sae_type: Literal["sae", "crosscoder"]
     d_model: int
     expansion_factor: int
     use_decoder_bias: bool = True
@@ -114,7 +114,7 @@ def associated_hook_points(self) -> list[str]:
 
 
 class SAEConfig(BaseSAEConfig):
-    sae_type: Literal["sae", "crosscoder", "mixcoder"] = "sae"
+    sae_type: Literal["sae", "crosscoder"] = "sae"
     hook_point_in: str
     hook_point_out: str = Field(default_factory=lambda validated_model: validated_model["hook_point_in"])
     use_glu_encoder: bool = False
@@ -125,7 +125,7 @@ def associated_hook_points(self) -> list[str]:
 
 
 class CrossCoderConfig(BaseSAEConfig):
-    sae_type: Literal["sae", "crosscoder", "mixcoder"] = "crosscoder"
+    sae_type: Literal["sae", "crosscoder"] = "crosscoder"
     hook_points: list[str]
 
     @property

src/lm_saes/config.py

@@ -42,9 +42,6 @@ class AnalyzeSAESettings(BaseSettings):
     model: Optional[LanguageModelConfig] = None
     """Configuration for the language model. Required if using dataset sources."""
 
-    model_name: Optional[str] = None
-    """Name of the tokenizer to load. Mixcoder requires a tokenizer to get the modality indices."""
-
     analyzer: FeatureAnalyzerConfig
     """Configuration for feature analysis"""
 

src/lm_saes/mixcoder.py

@@ -43,9 +43,6 @@ class EvaluateSAESettings(BaseSettings):
     model: Optional[LanguageModelConfig] = None
     """Configuration for the language model. Required if using dataset sources."""
 
-    model_name: Optional[str] = None
-    """Name of the tokenizer to load. Mixcoder requires a tokenizer to get the modality indices."""
-
     eval: EvalConfig
     """Configuration for evaluation"""