Core API

deeplens.extractor

ExtractSingleSample

Extract MLP activations from individual text samples for analysis and intervention.

This class provides functionality to extract activations from single text inputs, useful for interactive analysis, debugging, and testing feature interventions on specific examples.

Source code in deeplens\extractor.py

class ExtractSingleSample():
    """Extract MLP activations from individual text samples for analysis and intervention.

    This class provides functionality to extract activations from single text inputs,
    useful for interactive analysis, debugging, and testing feature interventions on
    specific examples.
    """
    def __init__(
            self, 
            hf_model: str = "gpt2", 
            layer: int = 3, 
            max_length: int = 1024, 
            device: str = "auto",
            cache_dir: str = 'cache'
        ) -> None:
        """Initialize the single sample extractor with model configuration.

        Loads the specified model and tokenizer, and configures extraction parameters.
        The model is set to evaluation mode and moved to the appropriate device.

        Args:
            hf_model (str, optional): Name or path of the HuggingFace model to load.
                Should match the model used for sparse autoencoder training for consistency.
                Defaults to "gpt2".
            layer (int, optional): Index of the transformer layer to extract activations from.
                Should match the layer used for SAE training. 0-indexed. Defaults to 3.
            max_length (int, optional): Maximum sequence length for tokenization. Longer
                sequences will be truncated. Defaults to 1024.
            device (str, optional): Device for model inference. Can be "auto" for automatic
                selection, "cuda", "mps", or "cpu". Defaults to "auto".
            cache_dir (str, optional): Directory to cache downloaded models and datasets.
                Defaults to 'cache'.
        """
        os.makedirs(cache_dir, exist_ok=True)
        self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir)
        self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
        self.layer = layer
        self.max_length = max_length

        self.device = get_device(device)
        print(f"Running on device: {self.device}")

        self.model.to(self.device)
        self.model.eval()

    @torch.no_grad()
    def get_mlp_acts(self, sample: str) -> torch.Tensor:
        """Extract MLP activations for a single text sample.

        Processes the input text through the model and captures the MLP activations
        from the configured layer. The hook is automatically removed after extraction.

        Args:
            sample (str): Input text to process. Can be a word, phrase, or full sentence.
                Will be tokenized according to the model's tokenizer.

        Returns:
            torch.Tensor: Activation tensor with shape (sequence_length, hidden_dim),
                where sequence_length depends on the tokenized length of the input.
                The batch dimension is squeezed out.

        Note:
            The activations are automatically moved to CPU to save GPU memory.
        """
        hook, activations = self.set_forward_hook_and_return_activations(self.layer)
        tokens = self.tokenize(sample)
        _ = self.model(**tokens)
        acts = activations[-1].squeeze()
        hook.remove()
        return acts

    def tokenize(self, sample: str) -> dict:
        """Tokenize a single text sample without padding.

        Converts the input text into token IDs suitable for model input. No padding is
        applied since this is for single sample processing.

        Args:
            sample (str): Text string to tokenize.

        Returns:
            dict: Dictionary containing tokenized outputs with 'input_ids', 'attention_mask',
                and other tokenizer-specific keys as tensors on the configured device.
                Shape is (1, actual_length) where actual_length ≤ max_length.
        """
        return self.tokenizer(
            sample,
            truncation=True,
            padding=False,
            max_length=self.max_length,
            return_tensors='pt'
        ).to(self.device)

    def set_forward_hook_and_return_activations(self, layer_idx) -> tuple:
        """Register a forward hook to capture MLP activations from a specific layer.

        Creates a hook function that captures the output of the MLP activation function
        at the specified layer during forward passes. Activations are detached and moved
        to CPU to save GPU memory.

        Args:
            layer_idx (int): Index of the transformer layer to hook (0-indexed).

        Returns:
            tuple: A tuple containing:
                - hook (torch.utils.hooks.RemovableHandle): Handle to remove the hook later
                - activations (list): List that will be populated with activation tensors
                    during forward passes
        """
        activations = []
        def hook_fn(module, input, output):
            activations.append(output.detach().cpu())

        if isinstance(self.model, (
            transformers.GPT2LMHeadModel, 
            transformers.FalconForCausalLM
        )):
            hook = self.model.transformer.h[layer_idx].mlp.act.register_forward_hook(hook_fn)
        elif isinstance(self.model, (
            transformers.LlamaForCausalLM, 
            transformers.MistralForCausalLM, 
            transformers.Gemma3ForCausalLM, 
            transformers.GemmaForCausalLM, 
            transformers.Qwen2ForCausalLM,
            transformers.Qwen3ForCausalLM
        )):
            hook = self.model.model.layers[layer_idx].mlp.act_fn.register_forward_hook(hook_fn)
        elif isinstance(self.model, (
            transformers.PhiForCausalLM, 
            transformers.Phi3ForCausalLM
        )):
            hook = self.model.model.layers[layer_idx].mlp.activation_fn.register_forward_hook(hook_fn)
        else:
            raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

        return hook, activations

__init__(hf_model='gpt2', layer=3, max_length=1024, device='auto', cache_dir='cache')

Initialize the single sample extractor with model configuration.

Loads the specified model and tokenizer, and configures extraction parameters. The model is set to evaluation mode and moved to the appropriate device.

Parameters:

Name	Type	Description	Default
hf_model	str	Name or path of the HuggingFace model to load. Should match the model used for sparse autoencoder training for consistency. Defaults to "gpt2".	'gpt2'
layer	int	Index of the transformer layer to extract activations from. Should match the layer used for SAE training. 0-indexed. Defaults to 3.	3
max_length	int	Maximum sequence length for tokenization. Longer sequences will be truncated. Defaults to 1024.	1024
device	str	Device for model inference. Can be "auto" for automatic selection, "cuda", "mps", or "cpu". Defaults to "auto".	'auto'
cache_dir	str	Directory to cache downloaded models and datasets. Defaults to 'cache'.	'cache'

Source code in deeplens\extractor.py

def __init__(
        self, 
        hf_model: str = "gpt2", 
        layer: int = 3, 
        max_length: int = 1024, 
        device: str = "auto",
        cache_dir: str = 'cache'
    ) -> None:
    """Initialize the single sample extractor with model configuration.

    Loads the specified model and tokenizer, and configures extraction parameters.
    The model is set to evaluation mode and moved to the appropriate device.

    Args:
        hf_model (str, optional): Name or path of the HuggingFace model to load.
            Should match the model used for sparse autoencoder training for consistency.
            Defaults to "gpt2".
        layer (int, optional): Index of the transformer layer to extract activations from.
            Should match the layer used for SAE training. 0-indexed. Defaults to 3.
        max_length (int, optional): Maximum sequence length for tokenization. Longer
            sequences will be truncated. Defaults to 1024.
        device (str, optional): Device for model inference. Can be "auto" for automatic
            selection, "cuda", "mps", or "cpu". Defaults to "auto".
        cache_dir (str, optional): Directory to cache downloaded models and datasets.
            Defaults to 'cache'.
    """
    os.makedirs(cache_dir, exist_ok=True)
    self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir)
    self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
    self.layer = layer
    self.max_length = max_length

    self.device = get_device(device)
    print(f"Running on device: {self.device}")

    self.model.to(self.device)
    self.model.eval()

get_mlp_acts(sample)

Extract MLP activations for a single text sample.

Processes the input text through the model and captures the MLP activations from the configured layer. The hook is automatically removed after extraction.

Parameters:

Name	Type	Description	Default
sample	str	Input text to process. Can be a word, phrase, or full sentence. Will be tokenized according to the model's tokenizer.	required

Returns:

Type	Description
Tensor	torch.Tensor: Activation tensor with shape (sequence_length, hidden_dim), where sequence_length depends on the tokenized length of the input. The batch dimension is squeezed out.

Note

The activations are automatically moved to CPU to save GPU memory.

Source code in deeplens\extractor.py

@torch.no_grad()
def get_mlp_acts(self, sample: str) -> torch.Tensor:
    """Extract MLP activations for a single text sample.

    Processes the input text through the model and captures the MLP activations
    from the configured layer. The hook is automatically removed after extraction.

    Args:
        sample (str): Input text to process. Can be a word, phrase, or full sentence.
            Will be tokenized according to the model's tokenizer.

    Returns:
        torch.Tensor: Activation tensor with shape (sequence_length, hidden_dim),
            where sequence_length depends on the tokenized length of the input.
            The batch dimension is squeezed out.

    Note:
        The activations are automatically moved to CPU to save GPU memory.
    """
    hook, activations = self.set_forward_hook_and_return_activations(self.layer)
    tokens = self.tokenize(sample)
    _ = self.model(**tokens)
    acts = activations[-1].squeeze()
    hook.remove()
    return acts

set_forward_hook_and_return_activations(layer_idx)

Register a forward hook to capture MLP activations from a specific layer.

Creates a hook function that captures the output of the MLP activation function at the specified layer during forward passes. Activations are detached and moved to CPU to save GPU memory.

Parameters:

Name	Type	Description	Default
layer_idx	int	Index of the transformer layer to hook (0-indexed).	required

Returns:

Name	Type	Description
tuple	tuple	A tuple containing: - hook (torch.utils.hooks.RemovableHandle): Handle to remove the hook later - activations (list): List that will be populated with activation tensors during forward passes

Source code in deeplens\extractor.py

def set_forward_hook_and_return_activations(self, layer_idx) -> tuple:
    """Register a forward hook to capture MLP activations from a specific layer.

    Creates a hook function that captures the output of the MLP activation function
    at the specified layer during forward passes. Activations are detached and moved
    to CPU to save GPU memory.

    Args:
        layer_idx (int): Index of the transformer layer to hook (0-indexed).

    Returns:
        tuple: A tuple containing:
            - hook (torch.utils.hooks.RemovableHandle): Handle to remove the hook later
            - activations (list): List that will be populated with activation tensors
                during forward passes
    """
    activations = []
    def hook_fn(module, input, output):
        activations.append(output.detach().cpu())

    if isinstance(self.model, (
        transformers.GPT2LMHeadModel, 
        transformers.FalconForCausalLM
    )):
        hook = self.model.transformer.h[layer_idx].mlp.act.register_forward_hook(hook_fn)
    elif isinstance(self.model, (
        transformers.LlamaForCausalLM, 
        transformers.MistralForCausalLM, 
        transformers.Gemma3ForCausalLM, 
        transformers.GemmaForCausalLM, 
        transformers.Qwen2ForCausalLM,
        transformers.Qwen3ForCausalLM
    )):
        hook = self.model.model.layers[layer_idx].mlp.act_fn.register_forward_hook(hook_fn)
    elif isinstance(self.model, (
        transformers.PhiForCausalLM, 
        transformers.Phi3ForCausalLM
    )):
        hook = self.model.model.layers[layer_idx].mlp.activation_fn.register_forward_hook(hook_fn)
    else:
        raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

    return hook, activations

tokenize(sample)

Tokenize a single text sample without padding.

Converts the input text into token IDs suitable for model input. No padding is applied since this is for single sample processing.

Parameters:

Name	Type	Description	Default
sample	str	Text string to tokenize.	required

Returns:

Name	Type	Description
dict	dict	Dictionary containing tokenized outputs with 'input_ids', 'attention_mask', and other tokenizer-specific keys as tensors on the configured device. Shape is (1, actual_length) where actual_length ≤ max_length.

Source code in deeplens\extractor.py

def tokenize(self, sample: str) -> dict:
    """Tokenize a single text sample without padding.

    Converts the input text into token IDs suitable for model input. No padding is
    applied since this is for single sample processing.

    Args:
        sample (str): Text string to tokenize.

    Returns:
        dict: Dictionary containing tokenized outputs with 'input_ids', 'attention_mask',
            and other tokenizer-specific keys as tensors on the configured device.
            Shape is (1, actual_length) where actual_length ≤ max_length.
    """
    return self.tokenizer(
        sample,
        truncation=True,
        padding=False,
        max_length=self.max_length,
        return_tensors='pt'
    ).to(self.device)

FromHuggingFace

Extract MLP activations from transformer models using HuggingFace datasets.

This class loads a pre-trained transformer model and processes samples from a streaming dataset to extract and save intermediate layer activations. Designed for collecting training data for sparse autoencoders.

Source code in deeplens\extractor.py

class FromHuggingFace():
    """Extract MLP activations from transformer models using HuggingFace datasets.

    This class loads a pre-trained transformer model and processes samples from a streaming
    dataset to extract and save intermediate layer activations. Designed for collecting
    training data for sparse autoencoders.
    """
    def __init__(
            self, 
            hf_model: str = "gpt2", 
            layer: int = 6,
            dataset_name: str = "HuggingFaceFW/fineweb",
            num_samples: int = 100000,
            seq_length: int = 128,
            inference_batch_size: int = 16, 
            device: str = "auto",
            save_features: bool = True,
            cache_dir: str = 'cache'
        ) -> None:
        """Initialize the activation extractor with model and dataset configuration.

        Loads the specified model and tokenizer, sets up dataset streaming, and configures
        extraction parameters. The model is set to evaluation mode and moved to the
        appropriate device.

        Args:
            hf_model (str, optional): Name or path of the HuggingFace model to load.
                Should be a valid model identifier (e.g., "gpt2", "meta-llama/Llama-2-7b").
                Defaults to "gpt2".
            layer (int, optional): Index of the transformer layer to extract activations from.
                0-indexed. Defaults to 6.
            dataset_name (str, optional): Name of the HuggingFace dataset to stream.
                Must be a valid dataset identifier. Defaults to "HuggingFaceFW/fineweb".
            num_samples (int, optional): Number of samples to extract from the dataset.
                Defaults to 100000.
            seq_length (int, optional): Maximum sequence length for tokenization. Sequences
                will be truncated or padded to this length. Defaults to 128.
            inference_batch_size (int, optional): Batch size for processing samples through
                the model. Higher values increase memory usage but improve speed.
                Defaults to 16.
            device (str, optional): Device for model inference. Can be "auto" for automatic
                selection, "cuda", "mps", or "cpu". Defaults to "auto".
            save_features (bool, optional): Whether to save extracted features to disk in
                the 'saved_features' directory. Defaults to True.
            cache_dir (str, optional): Directory to cache downloaded models and datasets.
                Defaults to 'cache'.
        """
        os.makedirs(cache_dir, exist_ok=True)
        self.model_name = hf_model.split('/')[-1]
        self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir)
        self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
        self.tokenizer.pad_token = self.tokenizer.eos_token

        self.layer = layer
        self.batch_size = inference_batch_size
        self.save_features = save_features
        self.seq_length = seq_length
        self.num_samples = num_samples
        self.dataset = load_dataset(
            dataset_name, 
            split='train',
            streaming=True,
            cache_dir=cache_dir
        ).take(num_samples)

        self.device = get_device(device)
        print(f"Running on device: {self.device}")

        self.model.to(self.device)
        self.model.eval()

    def tokenize(self, examples) -> dict:
        """Tokenize text examples with padding and truncation.

        Converts raw text into token IDs suitable for model input, applying padding to
        seq_length and truncation as needed.

        Args:
            examples (dict): Dictionary containing a 'text' key with a list of text strings
                to tokenize.

        Returns:
            dict: Dictionary with tokenized outputs including 'input_ids', 'attention_mask',
                and other tokenizer-specific keys. All tensors have shape (batch_size, seq_length).
        """
        return self.tokenizer(
            examples['text'],
            truncation=True,
            padding='max_length',
            max_length=self.seq_length,
            return_tensors='pt'
        )

    def set_forward_hook_and_return_activations(self, layer_idx) -> tuple:
        """Register a forward hook to capture MLP activations from a specific layer.

        Creates a hook function that captures the output of the MLP activation function
        at the specified layer during forward passes. Activations are detached and moved
        to CPU to save GPU memory.

        Args:
            layer_idx (int): Index of the transformer layer to hook (0-indexed).

        Returns:
            tuple: A tuple containing:
                - hook (torch.utils.hooks.RemovableHandle): Handle to remove the hook later
                - activations (list): List that will be populated with activation tensors
                    during forward passes
        """
        activations = []
        def hook_fn(module, input, output):
            activations.append(output.detach().cpu())

        if isinstance(self.model, (
            transformers.GPT2LMHeadModel, 
            transformers.FalconForCausalLM
        )):
            hook = self.model.transformer.h[layer_idx].mlp.act.register_forward_hook(hook_fn)
        elif isinstance(self.model, (
            transformers.LlamaForCausalLM, 
            transformers.MistralForCausalLM, 
            transformers.Gemma3ForCausalLM, 
            transformers.GemmaForCausalLM, 
            transformers.Qwen2ForCausalLM,
            transformers.Qwen3ForCausalLM
        )):
            hook = self.model.model.layers[layer_idx].mlp.act_fn.register_forward_hook(hook_fn)
        elif isinstance(self.model, (
            transformers.PhiForCausalLM, 
            transformers.Phi3ForCausalLM
        )):
            hook = self.model.model.layers[layer_idx].mlp.activation_fn.register_forward_hook(hook_fn)
        else:
            raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

        return hook, activations

    @torch.no_grad()
    def extract_features(self) -> torch.Tensor:
        """Extract MLP activations from the specified layer across the entire dataset.

        Processes the dataset in batches, extracting activations from the configured layer.
        Filters out padding tokens to ensure only valid activations are collected. Optionally
        saves the extracted features to disk.

        The extraction process:
        1. Batches text samples for efficient processing
        2. Tokenizes and pads/truncates to seq_length
        3. Runs forward pass and captures activations via hook
        4. Filters out activations from padding tokens using attention mask
        5. Concatenates all valid activations into a single tensor

        Returns:
            torch.Tensor: Concatenated activation tensor with shape (total_tokens, hidden_dim),
                where total_tokens is the sum of all non-padding tokens across all samples.
                The tensor is saved to 'saved_features/features_layer_{layer}_{num_tokens}.pt'
                if save_features=True.

        Note:
            The hook is automatically removed after extraction to prevent memory leaks.
            Progress is displayed via tqdm progress bar.
        """
        hook, activations = self.set_forward_hook_and_return_activations(self.layer)
        all_activations = []
        batch_texts = []     
        for example in tqdm(self.dataset, desc=f"Extracting from L{self.layer}", total=self.num_samples):
            batch_texts.append(example['text'])
            if len(batch_texts) == self.batch_size:
                tokens = self.tokenize({'text': batch_texts})
                tokens = {k: v.to(self.device) for k, v in tokens.items()}
                _ = self.model(**tokens)
                batch_acts = activations[-1]
                attention_mask = tokens["attention_mask"].cpu()
                for i in range(batch_acts.shape[0]):
                    non_pad_mask = attention_mask[i].bool()
                    valid_acts = batch_acts[i][non_pad_mask]
                    all_activations.append(valid_acts)
                batch_texts = []

        # for residual text not batched
        if batch_texts:
            tokens = self.tokenize({'text': batch_texts})
            tokens = {k: v.to(self.device) for k, v in tokens.items()}
            _ = self.model(**tokens)
            batch_acts = activations[-1]
            attention_mask = tokens["attention_mask"].cpu()
            for i in range(batch_acts.shape[0]):
                non_pad_mask = attention_mask[i].bool()
                valid_acts = batch_acts[i][non_pad_mask]
                all_activations.append(valid_acts)

        hook.remove()

        features = torch.cat(all_activations, dim=0)
        print(f"Extracting features... (shape: {features.shape})")

        if self.save_features:
            os.makedirs(f'saved_features/{self.model_name}', exist_ok=True)
            save_path = f"saved_features/{self.model_name}/features_layer_{self.layer}_{features.shape[0]}.pt"
            torch.save(features, save_path)
            print(f"Features saved to {save_path}")

        return features

__init__(hf_model='gpt2', layer=6, dataset_name='HuggingFaceFW/fineweb', num_samples=100000, seq_length=128, inference_batch_size=16, device='auto', save_features=True, cache_dir='cache')

Initialize the activation extractor with model and dataset configuration.

Loads the specified model and tokenizer, sets up dataset streaming, and configures extraction parameters. The model is set to evaluation mode and moved to the appropriate device.

Parameters:

Name	Type	Description	Default
hf_model	str	Name or path of the HuggingFace model to load. Should be a valid model identifier (e.g., "gpt2", "meta-llama/Llama-2-7b"). Defaults to "gpt2".	'gpt2'
layer	int	Index of the transformer layer to extract activations from. 0-indexed. Defaults to 6.	6
dataset_name	str	Name of the HuggingFace dataset to stream. Must be a valid dataset identifier. Defaults to "HuggingFaceFW/fineweb".	'HuggingFaceFW/fineweb'
num_samples	int	Number of samples to extract from the dataset. Defaults to 100000.	100000
seq_length	int	Maximum sequence length for tokenization. Sequences will be truncated or padded to this length. Defaults to 128.	128
inference_batch_size	int	Batch size for processing samples through the model. Higher values increase memory usage but improve speed. Defaults to 16.	16
device	str	Device for model inference. Can be "auto" for automatic selection, "cuda", "mps", or "cpu". Defaults to "auto".	'auto'
save_features	bool	Whether to save extracted features to disk in the 'saved_features' directory. Defaults to True.	True
cache_dir	str	Directory to cache downloaded models and datasets. Defaults to 'cache'.	'cache'

Source code in deeplens\extractor.py

def __init__(
        self, 
        hf_model: str = "gpt2", 
        layer: int = 6,
        dataset_name: str = "HuggingFaceFW/fineweb",
        num_samples: int = 100000,
        seq_length: int = 128,
        inference_batch_size: int = 16, 
        device: str = "auto",
        save_features: bool = True,
        cache_dir: str = 'cache'
    ) -> None:
    """Initialize the activation extractor with model and dataset configuration.

    Loads the specified model and tokenizer, sets up dataset streaming, and configures
    extraction parameters. The model is set to evaluation mode and moved to the
    appropriate device.

    Args:
        hf_model (str, optional): Name or path of the HuggingFace model to load.
            Should be a valid model identifier (e.g., "gpt2", "meta-llama/Llama-2-7b").
            Defaults to "gpt2".
        layer (int, optional): Index of the transformer layer to extract activations from.
            0-indexed. Defaults to 6.
        dataset_name (str, optional): Name of the HuggingFace dataset to stream.
            Must be a valid dataset identifier. Defaults to "HuggingFaceFW/fineweb".
        num_samples (int, optional): Number of samples to extract from the dataset.
            Defaults to 100000.
        seq_length (int, optional): Maximum sequence length for tokenization. Sequences
            will be truncated or padded to this length. Defaults to 128.
        inference_batch_size (int, optional): Batch size for processing samples through
            the model. Higher values increase memory usage but improve speed.
            Defaults to 16.
        device (str, optional): Device for model inference. Can be "auto" for automatic
            selection, "cuda", "mps", or "cpu". Defaults to "auto".
        save_features (bool, optional): Whether to save extracted features to disk in
            the 'saved_features' directory. Defaults to True.
        cache_dir (str, optional): Directory to cache downloaded models and datasets.
            Defaults to 'cache'.
    """
    os.makedirs(cache_dir, exist_ok=True)
    self.model_name = hf_model.split('/')[-1]
    self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir)
    self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
    self.tokenizer.pad_token = self.tokenizer.eos_token

    self.layer = layer
    self.batch_size = inference_batch_size
    self.save_features = save_features
    self.seq_length = seq_length
    self.num_samples = num_samples
    self.dataset = load_dataset(
        dataset_name, 
        split='train',
        streaming=True,
        cache_dir=cache_dir
    ).take(num_samples)

    self.device = get_device(device)
    print(f"Running on device: {self.device}")

    self.model.to(self.device)
    self.model.eval()

extract_features()

Extract MLP activations from the specified layer across the entire dataset.

Processes the dataset in batches, extracting activations from the configured layer. Filters out padding tokens to ensure only valid activations are collected. Optionally saves the extracted features to disk.

The extraction process: 1. Batches text samples for efficient processing 2. Tokenizes and pads/truncates to seq_length 3. Runs forward pass and captures activations via hook 4. Filters out activations from padding tokens using attention mask 5. Concatenates all valid activations into a single tensor

Returns:

Type	Description
Tensor	torch.Tensor: Concatenated activation tensor with shape (total_tokens, hidden_dim), where total_tokens is the sum of all non-padding tokens across all samples. The tensor is saved to 'saved_features/features_layer_{layer}_{num_tokens}.pt' if save_features=True.

Note

The hook is automatically removed after extraction to prevent memory leaks. Progress is displayed via tqdm progress bar.

Source code in deeplens\extractor.py

@torch.no_grad()
def extract_features(self) -> torch.Tensor:
    """Extract MLP activations from the specified layer across the entire dataset.

    Processes the dataset in batches, extracting activations from the configured layer.
    Filters out padding tokens to ensure only valid activations are collected. Optionally
    saves the extracted features to disk.

    The extraction process:
    1. Batches text samples for efficient processing
    2. Tokenizes and pads/truncates to seq_length
    3. Runs forward pass and captures activations via hook
    4. Filters out activations from padding tokens using attention mask
    5. Concatenates all valid activations into a single tensor

    Returns:
        torch.Tensor: Concatenated activation tensor with shape (total_tokens, hidden_dim),
            where total_tokens is the sum of all non-padding tokens across all samples.
            The tensor is saved to 'saved_features/features_layer_{layer}_{num_tokens}.pt'
            if save_features=True.

    Note:
        The hook is automatically removed after extraction to prevent memory leaks.
        Progress is displayed via tqdm progress bar.
    """
    hook, activations = self.set_forward_hook_and_return_activations(self.layer)
    all_activations = []
    batch_texts = []     
    for example in tqdm(self.dataset, desc=f"Extracting from L{self.layer}", total=self.num_samples):
        batch_texts.append(example['text'])
        if len(batch_texts) == self.batch_size:
            tokens = self.tokenize({'text': batch_texts})
            tokens = {k: v.to(self.device) for k, v in tokens.items()}
            _ = self.model(**tokens)
            batch_acts = activations[-1]
            attention_mask = tokens["attention_mask"].cpu()
            for i in range(batch_acts.shape[0]):
                non_pad_mask = attention_mask[i].bool()
                valid_acts = batch_acts[i][non_pad_mask]
                all_activations.append(valid_acts)
            batch_texts = []

    # for residual text not batched
    if batch_texts:
        tokens = self.tokenize({'text': batch_texts})
        tokens = {k: v.to(self.device) for k, v in tokens.items()}
        _ = self.model(**tokens)
        batch_acts = activations[-1]
        attention_mask = tokens["attention_mask"].cpu()
        for i in range(batch_acts.shape[0]):
            non_pad_mask = attention_mask[i].bool()
            valid_acts = batch_acts[i][non_pad_mask]
            all_activations.append(valid_acts)

    hook.remove()

    features = torch.cat(all_activations, dim=0)
    print(f"Extracting features... (shape: {features.shape})")

    if self.save_features:
        os.makedirs(f'saved_features/{self.model_name}', exist_ok=True)
        save_path = f"saved_features/{self.model_name}/features_layer_{self.layer}_{features.shape[0]}.pt"
        torch.save(features, save_path)
        print(f"Features saved to {save_path}")

    return features

set_forward_hook_and_return_activations(layer_idx)

Register a forward hook to capture MLP activations from a specific layer.

Creates a hook function that captures the output of the MLP activation function at the specified layer during forward passes. Activations are detached and moved to CPU to save GPU memory.

Parameters:

Name	Type	Description	Default
layer_idx	int	Index of the transformer layer to hook (0-indexed).	required

Returns:

Name	Type	Description
tuple	tuple	A tuple containing: - hook (torch.utils.hooks.RemovableHandle): Handle to remove the hook later - activations (list): List that will be populated with activation tensors during forward passes

Source code in deeplens\extractor.py

def set_forward_hook_and_return_activations(self, layer_idx) -> tuple:
    """Register a forward hook to capture MLP activations from a specific layer.

    Creates a hook function that captures the output of the MLP activation function
    at the specified layer during forward passes. Activations are detached and moved
    to CPU to save GPU memory.

    Args:
        layer_idx (int): Index of the transformer layer to hook (0-indexed).

    Returns:
        tuple: A tuple containing:
            - hook (torch.utils.hooks.RemovableHandle): Handle to remove the hook later
            - activations (list): List that will be populated with activation tensors
                during forward passes
    """
    activations = []
    def hook_fn(module, input, output):
        activations.append(output.detach().cpu())

    if isinstance(self.model, (
        transformers.GPT2LMHeadModel, 
        transformers.FalconForCausalLM
    )):
        hook = self.model.transformer.h[layer_idx].mlp.act.register_forward_hook(hook_fn)
    elif isinstance(self.model, (
        transformers.LlamaForCausalLM, 
        transformers.MistralForCausalLM, 
        transformers.Gemma3ForCausalLM, 
        transformers.GemmaForCausalLM, 
        transformers.Qwen2ForCausalLM,
        transformers.Qwen3ForCausalLM
    )):
        hook = self.model.model.layers[layer_idx].mlp.act_fn.register_forward_hook(hook_fn)
    elif isinstance(self.model, (
        transformers.PhiForCausalLM, 
        transformers.Phi3ForCausalLM
    )):
        hook = self.model.model.layers[layer_idx].mlp.activation_fn.register_forward_hook(hook_fn)
    else:
        raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

    return hook, activations

tokenize(examples)

Tokenize text examples with padding and truncation.

Converts raw text into token IDs suitable for model input, applying padding to seq_length and truncation as needed.

Parameters:

Name	Type	Description	Default
examples	dict	Dictionary containing a 'text' key with a list of text strings to tokenize.	required

Returns:

Name	Type	Description
dict	dict	Dictionary with tokenized outputs including 'input_ids', 'attention_mask', and other tokenizer-specific keys. All tensors have shape (batch_size, seq_length).

Source code in deeplens\extractor.py

def tokenize(self, examples) -> dict:
    """Tokenize text examples with padding and truncation.

    Converts raw text into token IDs suitable for model input, applying padding to
    seq_length and truncation as needed.

    Args:
        examples (dict): Dictionary containing a 'text' key with a list of text strings
            to tokenize.

    Returns:
        dict: Dictionary with tokenized outputs including 'input_ids', 'attention_mask',
            and other tokenizer-specific keys. All tensors have shape (batch_size, seq_length).
    """
    return self.tokenizer(
        examples['text'],
        truncation=True,
        padding='max_length',
        max_length=self.seq_length,
        return_tensors='pt'
    )

deeplens.intervene

InterveneFeatures

Manipulate and intervene on sparse autoencoder latent features.

This class loads a trained sparse autoencoder and provides methods to analyze and modify its latent feature space, enabling causal analysis of feature effects on model behavior.

Source code in deeplens\intervene.py

class InterveneFeatures():
    """Manipulate and intervene on sparse autoencoder latent features.

    This class loads a trained sparse autoencoder and provides methods to analyze
    and modify its latent feature space, enabling causal analysis of feature effects
    on model behavior.
    """
    def __init__(
            self,
            sae_model: str = None,
            sae_config: str | dict = None,
            device: str = "auto"
        ):
        """Initialize the InterveneFeatures class for manipulating sparse autoencoder latent features.

        This class provides functionality to load a trained sparse autoencoder model and
        intervene on its latent feature space to analyze and modify activations.

        Args:
            sae_model (str, optional): Path to the trained sparse autoencoder model weights file.
                Should be a .pt or .pth file containing the model state dict. Defaults to None.
            sae_config (str | dict, optional): Configuration for the sparse autoencoder.
                Can be either a dictionary containing model hyperparameters or a path to a
                YAML configuration file. Defaults to None.
            device (str, optional): Device to run computations on. Can be "auto" for automatic
                selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
        """
        self.model_dir = sae_model

        self.device = get_device(device)
        print(f"Running on device: {self.device}")

        if isinstance(sae_config, dict):
            self.model_config = sae_config
        elif isinstance(sae_config, str) and sae_config.endswith(".yaml"):
            self.model_config = self.config_from_yaml(sae_config)
        else:
            raise ValueError("sae_config must be dict or path to .yaml file.")

        self.model = self.load_model()

    @torch.no_grad()
    def get_decoded(self, activations) -> torch.Tensor:
        """Encode input activations through the sparse autoencoder to get latent features.

        Passes the input activations through the sparse autoencoder's forward pass and
        returns the latent feature representation (z) from the encoded space.

        Args:
            activations (torch.Tensor | array-like): Input activations to encode. Can be a
                PyTorch tensor or any array-like structure that can be converted to a tensor.

        Returns:
            torch.Tensor: The latent feature representation (z) from the sparse autoencoder's
                encoded space.
        """
        if not isinstance(activations, torch.Tensor):
            activations = torch.Tensor(activations)
        activations = activations.to(self.device)
        _, z, _ = self.model(activations)
        return z

    @torch.no_grad()
    def get_alive_features(
            self, 
            activations: torch.Tensor, 
            token_position: int = -1, 
            k: int | None = None,
            return_values: bool = False
        ) -> torch.Tensor:
        """Get indices of non-zero (active) features in the latent space for a specific token.

        Encodes the input activations through the sparse autoencoder and identifies which
        latent features are active (non-zero) at the specified token position.

        Args:
            activations (torch.Tensor | array-like): Input activations to encode. Can be a
                PyTorch tensor or any array-like structure that can be converted to a tensor.
            token_position (int, optional): Position of the token in the sequence to analyze.
                Use -1 for the last token. Defaults to -1.
            k (int, optional): If provided, returns only the top-k most active features
                instead of all non-zero features. Defaults to None.
            return_values (bool, optional): If True, returns both indices and values.
                Defaults to False.

        Returns:
            torch.Tensor | tuple[torch.Tensor, torch.Tensor]: If return_values is False,
                returns a 1D tensor containing the indices of non-zero features. If True,
                returns a tuple of (indices, values).
        """
        z = self.get_decoded(activations)
        z_token = z[token_position]
        if k is not None:
            topk_result = torch.topk(z_token, k=k)
            feature_idxs = topk_result.indices
            feature_vals = topk_result.values
        else:
            feature_idxs = torch.nonzero(z_token != 0, as_tuple=False).squeeze(-1)
            feature_vals = z_token[feature_idxs]
        if return_values:
            return feature_idxs, feature_vals
        return feature_idxs

    @torch.no_grad()
    def intervene_feature(
            self, 
            activations, 
            feature: int, 
            alpha: float = 2.0,
            token_positions: int | list[int] | None = None
        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Intervene on a specific latent feature by scaling its activation.

        Encodes the input activations, multiplies the specified feature by alpha at the
        given token positions, and returns both the original and modified decoded outputs
        for comparison.

        Args:
            activations (torch.Tensor | array-like): Input activations to encode and modify.
                Can be a PyTorch tensor or any array-like structure.
            feature (int): Index of the latent feature to intervene on.
            alpha (float, optional): Scaling factor to apply to the feature. Values > 1
                amplify the feature, values < 1 suppress it. Defaults to 2.0.
            token_positions (int | list[int] | None, optional): Token position(s) at which
                to apply the intervention. If None, applies to all tokens. If int, applies
                to a single position. If list, applies to multiple positions. Defaults to None.

        Returns:
            tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
                - activations: The original input activations
                - original_decoded: Decoded output without intervention
                - modified_decoded: Decoded output with the feature intervention applied
        """
        if not isinstance(activations, torch.Tensor):
            activations = torch.Tensor(activations).unsqueeze(0)

        activations = activations.to(self.device)
        _, z, _ = self.model(activations)
        modified = z.clone()

        if token_positions is None:
            modified[:, feature] *= alpha
        elif isinstance(token_positions, int):
            modified[token_positions, feature] *= alpha
        else:
            for pos in token_positions:
                modified[pos, feature] *= alpha

        original_decoded = self.model.decode(z)
        modified_decoded = self.model.decode(modified)
        return activations, original_decoded, modified_decoded

    def load_model(self) -> torch.nn.Module:
        """Load the sparse autoencoder model from disk.

        Loads the model weights from the specified path and initializes a
        SparseAutoencoder instance with the provided configuration.

        Returns:
            torch.nn.Module: The loaded sparse autoencoder model moved to the
                appropriate device.
        """
        weights = torch.load(self.model_dir, map_location=self.device)
        model = SparseAutoencoder(**self.model_config)
        model.load_state_dict(state_dict=weights)
        return model.to(self.device)

    def config_from_yaml(self, file) -> dict:
        """Load sparse autoencoder configuration from a YAML file.

        Reads and parses a YAML configuration file containing the hyperparameters
        for the sparse autoencoder model.

        Args:
            file (str): Path to the YAML configuration file.

        Returns:
            dict: Dictionary containing the model configuration parameters.
        """
        try:
            with open(file, "r") as f:
                config = yaml.safe_load(f)
            return config
        except FileNotFoundError:
            raise FileNotFoundError(f"Config file not found: {file}")
        except yaml.YAMLError as e:
            raise ValueError(f"Invalid YAML in {file}: {e}")

__init__(sae_model=None, sae_config=None, device='auto')

Initialize the InterveneFeatures class for manipulating sparse autoencoder latent features.

This class provides functionality to load a trained sparse autoencoder model and intervene on its latent feature space to analyze and modify activations.

Parameters:

Name	Type	Description	Default
sae_model	str	Path to the trained sparse autoencoder model weights file. Should be a .pt or .pth file containing the model state dict. Defaults to None.	None
sae_config	str | dict	Configuration for the sparse autoencoder. Can be either a dictionary containing model hyperparameters or a path to a YAML configuration file. Defaults to None.	None
device	str	Device to run computations on. Can be "auto" for automatic selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".	'auto'

Source code in deeplens\intervene.py

def __init__(
        self,
        sae_model: str = None,
        sae_config: str | dict = None,
        device: str = "auto"
    ):
    """Initialize the InterveneFeatures class for manipulating sparse autoencoder latent features.

    This class provides functionality to load a trained sparse autoencoder model and
    intervene on its latent feature space to analyze and modify activations.

    Args:
        sae_model (str, optional): Path to the trained sparse autoencoder model weights file.
            Should be a .pt or .pth file containing the model state dict. Defaults to None.
        sae_config (str | dict, optional): Configuration for the sparse autoencoder.
            Can be either a dictionary containing model hyperparameters or a path to a
            YAML configuration file. Defaults to None.
        device (str, optional): Device to run computations on. Can be "auto" for automatic
            selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
    """
    self.model_dir = sae_model

    self.device = get_device(device)
    print(f"Running on device: {self.device}")

    if isinstance(sae_config, dict):
        self.model_config = sae_config
    elif isinstance(sae_config, str) and sae_config.endswith(".yaml"):
        self.model_config = self.config_from_yaml(sae_config)
    else:
        raise ValueError("sae_config must be dict or path to .yaml file.")

    self.model = self.load_model()

config_from_yaml(file)

Load sparse autoencoder configuration from a YAML file.

Reads and parses a YAML configuration file containing the hyperparameters for the sparse autoencoder model.

Parameters:

Name	Type	Description	Default
file	str	Path to the YAML configuration file.	required

Returns:

Name	Type	Description
dict	dict	Dictionary containing the model configuration parameters.

Source code in deeplens\intervene.py

def config_from_yaml(self, file) -> dict:
    """Load sparse autoencoder configuration from a YAML file.

    Reads and parses a YAML configuration file containing the hyperparameters
    for the sparse autoencoder model.

    Args:
        file (str): Path to the YAML configuration file.

    Returns:
        dict: Dictionary containing the model configuration parameters.
    """
    try:
        with open(file, "r") as f:
            config = yaml.safe_load(f)
        return config
    except FileNotFoundError:
        raise FileNotFoundError(f"Config file not found: {file}")
    except yaml.YAMLError as e:
        raise ValueError(f"Invalid YAML in {file}: {e}")

get_alive_features(activations, token_position=-1, k=None, return_values=False)

Get indices of non-zero (active) features in the latent space for a specific token.

Encodes the input activations through the sparse autoencoder and identifies which latent features are active (non-zero) at the specified token position.

Parameters:

Name	Type	Description	Default
activations	Tensor | array - like	Input activations to encode. Can be a PyTorch tensor or any array-like structure that can be converted to a tensor.	required
token_position	int	Position of the token in the sequence to analyze. Use -1 for the last token. Defaults to -1.	-1
k	int	If provided, returns only the top-k most active features instead of all non-zero features. Defaults to None.	None
return_values	bool	If True, returns both indices and values. Defaults to False.	False

Returns:

Type	Description
Tensor	torch.Tensor | tuple[torch.Tensor, torch.Tensor]: If return_values is False, returns a 1D tensor containing the indices of non-zero features. If True, returns a tuple of (indices, values).

Source code in deeplens\intervene.py

@torch.no_grad()
def get_alive_features(
        self, 
        activations: torch.Tensor, 
        token_position: int = -1, 
        k: int | None = None,
        return_values: bool = False
    ) -> torch.Tensor:
    """Get indices of non-zero (active) features in the latent space for a specific token.

    Encodes the input activations through the sparse autoencoder and identifies which
    latent features are active (non-zero) at the specified token position.

    Args:
        activations (torch.Tensor | array-like): Input activations to encode. Can be a
            PyTorch tensor or any array-like structure that can be converted to a tensor.
        token_position (int, optional): Position of the token in the sequence to analyze.
            Use -1 for the last token. Defaults to -1.
        k (int, optional): If provided, returns only the top-k most active features
            instead of all non-zero features. Defaults to None.
        return_values (bool, optional): If True, returns both indices and values.
            Defaults to False.

    Returns:
        torch.Tensor | tuple[torch.Tensor, torch.Tensor]: If return_values is False,
            returns a 1D tensor containing the indices of non-zero features. If True,
            returns a tuple of (indices, values).
    """
    z = self.get_decoded(activations)
    z_token = z[token_position]
    if k is not None:
        topk_result = torch.topk(z_token, k=k)
        feature_idxs = topk_result.indices
        feature_vals = topk_result.values
    else:
        feature_idxs = torch.nonzero(z_token != 0, as_tuple=False).squeeze(-1)
        feature_vals = z_token[feature_idxs]
    if return_values:
        return feature_idxs, feature_vals
    return feature_idxs

get_decoded(activations)

Encode input activations through the sparse autoencoder to get latent features.

Passes the input activations through the sparse autoencoder's forward pass and returns the latent feature representation (z) from the encoded space.

Parameters:

Name	Type	Description	Default
activations	Tensor | array - like	Input activations to encode. Can be a PyTorch tensor or any array-like structure that can be converted to a tensor.	required

Returns:

Type	Description
Tensor	torch.Tensor: The latent feature representation (z) from the sparse autoencoder's encoded space.

Source code in deeplens\intervene.py

@torch.no_grad()
def get_decoded(self, activations) -> torch.Tensor:
    """Encode input activations through the sparse autoencoder to get latent features.

    Passes the input activations through the sparse autoencoder's forward pass and
    returns the latent feature representation (z) from the encoded space.

    Args:
        activations (torch.Tensor | array-like): Input activations to encode. Can be a
            PyTorch tensor or any array-like structure that can be converted to a tensor.

    Returns:
        torch.Tensor: The latent feature representation (z) from the sparse autoencoder's
            encoded space.
    """
    if not isinstance(activations, torch.Tensor):
        activations = torch.Tensor(activations)
    activations = activations.to(self.device)
    _, z, _ = self.model(activations)
    return z

intervene_feature(activations, feature, alpha=2.0, token_positions=None)

Intervene on a specific latent feature by scaling its activation.

Encodes the input activations, multiplies the specified feature by alpha at the given token positions, and returns both the original and modified decoded outputs for comparison.

Parameters:

Name	Type	Description	Default
activations	Tensor | array - like	Input activations to encode and modify. Can be a PyTorch tensor or any array-like structure.	required
feature	int	Index of the latent feature to intervene on.	required
alpha	float	Scaling factor to apply to the feature. Values > 1 amplify the feature, values < 1 suppress it. Defaults to 2.0.	2.0
token_positions	int | list[int] | None	Token position(s) at which to apply the intervention. If None, applies to all tokens. If int, applies to a single position. If list, applies to multiple positions. Defaults to None.	None

Returns:

Type	Description
tuple[Tensor, Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing: - activations: The original input activations - original_decoded: Decoded output without intervention - modified_decoded: Decoded output with the feature intervention applied

Source code in deeplens\intervene.py

@torch.no_grad()
def intervene_feature(
        self, 
        activations, 
        feature: int, 
        alpha: float = 2.0,
        token_positions: int | list[int] | None = None
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Intervene on a specific latent feature by scaling its activation.

    Encodes the input activations, multiplies the specified feature by alpha at the
    given token positions, and returns both the original and modified decoded outputs
    for comparison.

    Args:
        activations (torch.Tensor | array-like): Input activations to encode and modify.
            Can be a PyTorch tensor or any array-like structure.
        feature (int): Index of the latent feature to intervene on.
        alpha (float, optional): Scaling factor to apply to the feature. Values > 1
            amplify the feature, values < 1 suppress it. Defaults to 2.0.
        token_positions (int | list[int] | None, optional): Token position(s) at which
            to apply the intervention. If None, applies to all tokens. If int, applies
            to a single position. If list, applies to multiple positions. Defaults to None.

    Returns:
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
            - activations: The original input activations
            - original_decoded: Decoded output without intervention
            - modified_decoded: Decoded output with the feature intervention applied
    """
    if not isinstance(activations, torch.Tensor):
        activations = torch.Tensor(activations).unsqueeze(0)

    activations = activations.to(self.device)
    _, z, _ = self.model(activations)
    modified = z.clone()

    if token_positions is None:
        modified[:, feature] *= alpha
    elif isinstance(token_positions, int):
        modified[token_positions, feature] *= alpha
    else:
        for pos in token_positions:
            modified[pos, feature] *= alpha

    original_decoded = self.model.decode(z)
    modified_decoded = self.model.decode(modified)
    return activations, original_decoded, modified_decoded

load_model()

Load the sparse autoencoder model from disk.

Loads the model weights from the specified path and initializes a SparseAutoencoder instance with the provided configuration.

Returns:

Type	Description
Module	torch.nn.Module: The loaded sparse autoencoder model moved to the appropriate device.

Source code in deeplens\intervene.py

def load_model(self) -> torch.nn.Module:
    """Load the sparse autoencoder model from disk.

    Loads the model weights from the specified path and initializes a
    SparseAutoencoder instance with the provided configuration.

    Returns:
        torch.nn.Module: The loaded sparse autoencoder model moved to the
            appropriate device.
    """
    weights = torch.load(self.model_dir, map_location=self.device)
    model = SparseAutoencoder(**self.model_config)
    model.load_state_dict(state_dict=weights)
    return model.to(self.device)

ReinjectSingleSample

Reinject modified activations into a language model for causal inference.

This class enables injecting modified activations back into a transformer model's forward pass to observe the causal effects on model predictions and generated text. Useful for validating feature interventions and conducting mechanistic interpretability experiments.

Source code in deeplens\intervene.py

class ReinjectSingleSample():
    """Reinject modified activations into a language model for causal inference.

    This class enables injecting modified activations back into a transformer model's
    forward pass to observe the causal effects on model predictions and generated text.
    Useful for validating feature interventions and conducting mechanistic interpretability
    experiments.
    """
    def __init__(
            self, 
            hf_model: str, 
            device: str = "auto", 
            cache_dir: str = 'cache'
        ):
        """Initialize the ReinjectSingleSample class for causal inference with modified activations.

        Loads a HuggingFace causal language model and tokenizer to enable reinjection of
        modified activations into the model's forward pass for text generation and analysis.

        Args:
            hf_model (str): Name or path of the HuggingFace model to load.
                Should be a valid model identifier from the HuggingFace model hub
                (e.g., "gpt2", "meta-llama/Llama-2-7b").
            device (str, optional): Device to run computations on. Can be "auto" for automatic
                selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
            cache_dir (str, optional): Directory to cache downloaded models.
                Defaults to 'cache'.
        """
        self.device = get_device(device)
        print(f"Running on device: {self.device}")

        os.makedirs(cache_dir, exist_ok=True)
        self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
        self.model.eval()

    @torch.no_grad()
    def reinject_and_generate(
            self, 
            text, 
            modified_activations, 
            layer: int = 3, 
            generate: bool = False, 
            max_new_tokens: int = 25, 
            temperature: float = 1.0
        ) -> torch.Tensor | str:
        """Reinject modified activations into a model layer and optionally generate text.

        Replaces the activations at the specified layer with the provided modified activations
        during the forward pass. Can either return logits for the input text or generate
        new tokens autoregressively.

        Args:
            text (str): Input text to tokenize and process through the model.
            modified_activations (torch.Tensor): The modified activations to inject at the
                specified layer. Should have the appropriate shape for the layer's output.
            layer (int, optional): Index of the transformer layer where activations should
                be replaced. Defaults to 3.
            generate (bool, optional): If True, generates new tokens autoregressively.
                If False, only returns logits for the input. Defaults to False.
            max_new_tokens (int, optional): Maximum number of new tokens to generate when
                generate=True. Defaults to 25.
            temperature (float, optional): Sampling temperature for generation. Higher values
                (>1.0) make output more random, lower values (<1.0) more deterministic.
                Set to 0 for greedy decoding. Defaults to 1.0.

        Returns:
            torch.Tensor | str: If generate=False, returns the model's logits as a tensor.
                If generate=True, returns the generated text as a string.

        Note:
            The hook is automatically removed after execution to prevent interference
            with subsequent forward passes. For generation mode, the hook only affects
            the first forward pass to avoid applying the intervention to newly generated tokens.
        """
        modified_activations = modified_activations.to(self.device)
        call_count = [0]
        def replacement_hook(module, input, output):
            if generate and call_count[0] > 0:
                return output
            call_count[0] += 1
            return modified_activations

        mlp_module = self.get_module_for_replacement_hook(layer_idx=layer)
        hook = mlp_module.register_forward_hook(replacement_hook)
        tokens = self.tokenizer(text, return_tensors='pt').to(self.device)
        try:
            if generate:
                generated_ids = self.model.generate(
                    **tokens,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    do_sample=temperature > 0
                )
                return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
            else:
                out = self.model(**tokens)
                return out.logits
        finally:
            hook.remove()

    def get_module_for_replacement_hook(self, layer_idx) -> torch.nn.Module:
        """Get the MLP activation module for a specific layer.

        Retrieves the MLP activation function module at the specified layer,
        which can be used to register forward hooks for activation replacement.

        Args:
            layer_idx (int): Index of the transformer layer (0-indexed).

        Returns:
            torch.nn.Module: The MLP activation module at the specified layer.
        """
        if isinstance(self.model, (
            transformers.GPT2LMHeadModel, 
            transformers.FalconForCausalLM
        )):
            module = self.model.transformer.h[layer_idx].mlp.act
        elif isinstance(self.model, (
            transformers.LlamaForCausalLM, 
            transformers.MistralForCausalLM, 
            transformers.Gemma3ForCausalLM, 
            transformers.GemmaForCausalLM, 
            transformers.Qwen2ForCausalLM,
            transformers.Qwen3ForCausalLM
        )):
            module = self.model.model.layers[layer_idx].mlp.act_fn
        elif isinstance(self.model, (
            transformers.PhiForCausalLM, 
            transformers.Phi3ForCausalLM
        )):
            module = self.model.model.layers[layer_idx].mlp.activation_fn
        else:
            raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

        return module

__init__(hf_model, device='auto', cache_dir='cache')

Initialize the ReinjectSingleSample class for causal inference with modified activations.

Loads a HuggingFace causal language model and tokenizer to enable reinjection of modified activations into the model's forward pass for text generation and analysis.

Parameters:

Name	Type	Description	Default
hf_model	str	Name or path of the HuggingFace model to load. Should be a valid model identifier from the HuggingFace model hub (e.g., "gpt2", "meta-llama/Llama-2-7b").	required
device	str	Device to run computations on. Can be "auto" for automatic selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".	'auto'
cache_dir	str	Directory to cache downloaded models. Defaults to 'cache'.	'cache'

Source code in deeplens\intervene.py

def __init__(
        self, 
        hf_model: str, 
        device: str = "auto", 
        cache_dir: str = 'cache'
    ):
    """Initialize the ReinjectSingleSample class for causal inference with modified activations.

    Loads a HuggingFace causal language model and tokenizer to enable reinjection of
    modified activations into the model's forward pass for text generation and analysis.

    Args:
        hf_model (str): Name or path of the HuggingFace model to load.
            Should be a valid model identifier from the HuggingFace model hub
            (e.g., "gpt2", "meta-llama/Llama-2-7b").
        device (str, optional): Device to run computations on. Can be "auto" for automatic
            selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
        cache_dir (str, optional): Directory to cache downloaded models.
            Defaults to 'cache'.
    """
    self.device = get_device(device)
    print(f"Running on device: {self.device}")

    os.makedirs(cache_dir, exist_ok=True)
    self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir).to(self.device)
    self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
    self.model.eval()

get_module_for_replacement_hook(layer_idx)

Get the MLP activation module for a specific layer.

Retrieves the MLP activation function module at the specified layer, which can be used to register forward hooks for activation replacement.

Parameters:

Name	Type	Description	Default
layer_idx	int	Index of the transformer layer (0-indexed).	required

Returns:

Type	Description
Module	torch.nn.Module: The MLP activation module at the specified layer.

Source code in deeplens\intervene.py

def get_module_for_replacement_hook(self, layer_idx) -> torch.nn.Module:
    """Get the MLP activation module for a specific layer.

    Retrieves the MLP activation function module at the specified layer,
    which can be used to register forward hooks for activation replacement.

    Args:
        layer_idx (int): Index of the transformer layer (0-indexed).

    Returns:
        torch.nn.Module: The MLP activation module at the specified layer.
    """
    if isinstance(self.model, (
        transformers.GPT2LMHeadModel, 
        transformers.FalconForCausalLM
    )):
        module = self.model.transformer.h[layer_idx].mlp.act
    elif isinstance(self.model, (
        transformers.LlamaForCausalLM, 
        transformers.MistralForCausalLM, 
        transformers.Gemma3ForCausalLM, 
        transformers.GemmaForCausalLM, 
        transformers.Qwen2ForCausalLM,
        transformers.Qwen3ForCausalLM
    )):
        module = self.model.model.layers[layer_idx].mlp.act_fn
    elif isinstance(self.model, (
        transformers.PhiForCausalLM, 
        transformers.Phi3ForCausalLM
    )):
        module = self.model.model.layers[layer_idx].mlp.activation_fn
    else:
        raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

    return module

reinject_and_generate(text, modified_activations, layer=3, generate=False, max_new_tokens=25, temperature=1.0)

Reinject modified activations into a model layer and optionally generate text.

Replaces the activations at the specified layer with the provided modified activations during the forward pass. Can either return logits for the input text or generate new tokens autoregressively.

Parameters:

Name	Type	Description	Default
text	str	Input text to tokenize and process through the model.	required
modified_activations	Tensor	The modified activations to inject at the specified layer. Should have the appropriate shape for the layer's output.	required
layer	int	Index of the transformer layer where activations should be replaced. Defaults to 3.	3
generate	bool	If True, generates new tokens autoregressively. If False, only returns logits for the input. Defaults to False.	False
max_new_tokens	int	Maximum number of new tokens to generate when generate=True. Defaults to 25.	25
temperature	float	Sampling temperature for generation. Higher values (>1.0) make output more random, lower values (<1.0) more deterministic. Set to 0 for greedy decoding. Defaults to 1.0.	1.0

Returns:

Type	Description
Tensor | str	torch.Tensor | str: If generate=False, returns the model's logits as a tensor. If generate=True, returns the generated text as a string.

Note

The hook is automatically removed after execution to prevent interference with subsequent forward passes. For generation mode, the hook only affects the first forward pass to avoid applying the intervention to newly generated tokens.

Source code in deeplens\intervene.py

@torch.no_grad()
def reinject_and_generate(
        self, 
        text, 
        modified_activations, 
        layer: int = 3, 
        generate: bool = False, 
        max_new_tokens: int = 25, 
        temperature: float = 1.0
    ) -> torch.Tensor | str:
    """Reinject modified activations into a model layer and optionally generate text.

    Replaces the activations at the specified layer with the provided modified activations
    during the forward pass. Can either return logits for the input text or generate
    new tokens autoregressively.

    Args:
        text (str): Input text to tokenize and process through the model.
        modified_activations (torch.Tensor): The modified activations to inject at the
            specified layer. Should have the appropriate shape for the layer's output.
        layer (int, optional): Index of the transformer layer where activations should
            be replaced. Defaults to 3.
        generate (bool, optional): If True, generates new tokens autoregressively.
            If False, only returns logits for the input. Defaults to False.
        max_new_tokens (int, optional): Maximum number of new tokens to generate when
            generate=True. Defaults to 25.
        temperature (float, optional): Sampling temperature for generation. Higher values
            (>1.0) make output more random, lower values (<1.0) more deterministic.
            Set to 0 for greedy decoding. Defaults to 1.0.

    Returns:
        torch.Tensor | str: If generate=False, returns the model's logits as a tensor.
            If generate=True, returns the generated text as a string.

    Note:
        The hook is automatically removed after execution to prevent interference
        with subsequent forward passes. For generation mode, the hook only affects
        the first forward pass to avoid applying the intervention to newly generated tokens.
    """
    modified_activations = modified_activations.to(self.device)
    call_count = [0]
    def replacement_hook(module, input, output):
        if generate and call_count[0] > 0:
            return output
        call_count[0] += 1
        return modified_activations

    mlp_module = self.get_module_for_replacement_hook(layer_idx=layer)
    hook = mlp_module.register_forward_hook(replacement_hook)
    tokens = self.tokenizer(text, return_tensors='pt').to(self.device)
    try:
        if generate:
            generated_ids = self.model.generate(
                **tokens,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=temperature > 0
            )
            return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        else:
            out = self.model(**tokens)
            return out.logits
    finally:
        hook.remove()

deeplens.sae

SparseAutoencoder

Bases: Module

Sparse Autoencoder for learning interpretable features from neural network activations.

This implementation supports multiple sparsity methods (L1 regularization or top-k activation), optional weight tying between encoder and decoder, and unit-norm decoder constraints for improved feature interpretability. Designed for mechanistic interpretability research on transformer models.

The architecture consists of: - Optional input normalization layer - Encoder: Linear projection to expanded feature space with nonlinear activation - Decoder: Linear projection back to original space (optionally tied to encoder) - Sparsity constraint: Either L1 penalty or top-k activation selection

Source code in deeplens\sae.py

class SparseAutoencoder(nn.Module):
    """Sparse Autoencoder for learning interpretable features from neural network activations.

    This implementation supports multiple sparsity methods (L1 regularization or top-k activation),
    optional weight tying between encoder and decoder, and unit-norm decoder constraints for
    improved feature interpretability. Designed for mechanistic interpretability research
    on transformer models.

    The architecture consists of:
    - Optional input normalization layer
    - Encoder: Linear projection to expanded feature space with nonlinear activation
    - Decoder: Linear projection back to original space (optionally tied to encoder)
    - Sparsity constraint: Either L1 penalty or top-k activation selection
    """
    def __init__(
            self, 
            input_dims: int = 512, 
            n_features: int = 2048, 
            activation: str = "relu",
            input_norm: bool = True,
            k: int | None = None,
            beta_l1: float | None = None,
            tie_weights: bool = False,
            unit_norm_decoder: bool = True
        ) -> None:
        """Initialize the Sparse Autoencoder with specified architecture and sparsity settings.

        Args:
            input_dims (int, optional): Dimensionality of input activations (e.g., 3072 for
                GPT-2 layer 3 MLP output). Defaults to 512.
            n_features (int, optional): Number of learned features in the latent space.
                Typically set as expansion_factor × input_dims where expansion_factor is 2-8.
                Defaults to 2048.
            activation (str, optional): Nonlinearity applied after encoder. Must be 'relu'
                or 'silu'. ReLU is standard for interpretability; SiLU may improve reconstruction.
                Defaults to "relu".
            input_norm (bool, optional): Whether to apply LayerNorm to inputs before encoding.
                Helps stabilize training with varying activation magnitudes. Defaults to True.
            k (int | None, optional): If set, uses top-k sparsity (keeps only k largest activations
                per sample) instead of L1 regularization. Mutually exclusive with beta_l1.
                Defaults to None.
            beta_l1 (float | None, optional): L1 regularization coefficient for sparsity.
                Higher values encourage sparser activations. Ignored if k is set. Defaults to None.
            tie_weights (bool, optional): If True, decoder weights are the transpose of encoder
                weights (no separate decoder parameters). Reduces parameters but may hurt performance.
                Defaults to False.
            unit_norm_decoder (bool, optional): If True, constrains decoder weight columns to
                unit norm. Improves feature interpretability by removing scale ambiguity.
                Defaults to True.
        """
        super().__init__()
        self.norm = nn.LayerNorm(input_dims) if input_norm else nn.Identity()
        self.encoder = nn.Linear(input_dims, n_features, bias=True)
        self.decoder = None if tie_weights else nn.Linear(n_features, input_dims, bias=False)
        self.unit_norm_decoder = unit_norm_decoder
        self.input_norm = input_norm

        if activation == "relu":
            self.activation = nn.ReLU()
            kaiming_activation = "relu"
        elif activation == "silu":
            self.activation = nn.SiLU()
            kaiming_activation = "linear"
        else:
            raise ValueError("Activation must be 'relu' or 'silu'")

        nn.init.kaiming_normal_(self.encoder.weight, nonlinearity=kaiming_activation)
        if self.decoder is not None:
            nn.init.xavier_uniform_(self.decoder.weight)
            if self.unit_norm_decoder:
                self._renorm_decoder()

        self.k = k
        self.beta_l1 = beta_l1
        self.tie_weights = tie_weights

    @torch.no_grad()
    def _renorm_decoder(self, eps: float = 1e-8) -> None:
        """Normalize decoder weight columns to unit norm for improved interpretability.

        Constrains each feature's decoder weight vector to have L2 norm of 1, removing
        scale ambiguity from the learned features. This is a common technique in dictionary
        learning and sparse coding to ensure features represent directions rather than
        directions with varying magnitudes.

        Args:
            eps (float, optional): Small epsilon value to prevent division by zero for
                near-zero norm weights. Defaults to 1e-8.

        Returns:
            None: Modifies decoder weights in-place. No-op if decoder is None or
                unit_norm_decoder is False.
        """
        if self.decoder is not None and self.unit_norm_decoder:
            W = self.decoder.weight.data
            norms = W.norm(dim=1, keepdim=True).clamp_min(eps)
            self.decoder.weight.data = W / norms

    def encode(self, x) -> torch.Tensor:
        """Encode input activations into the sparse latent feature space.

        Applies optional normalization, linear transformation to expanded feature space,
        and nonlinear activation function.

        Args:
            x (torch.Tensor): Input activation tensor with shape (..., input_dims).
                Typically (batch_size, seq_length, input_dims) or (batch_size, input_dims).

        Returns:
            torch.Tensor: Encoded features with shape (..., n_features) after applying
                normalization (if enabled), linear encoding, and activation function.
                These are the pre-sparsity latent activations.
        """
        x = self.norm(x)
        return self.activation(self.encoder(x))

    def decode(self, z) -> torch.Tensor:
        """Decode sparse latent features back to the original activation space.

        Applies linear transformation from feature space back to input space. Uses either
        a separate decoder or transposed encoder weights depending on tie_weights setting.

        Args:
            z (torch.Tensor): Latent feature activations with shape (..., n_features).
                Typically the output of encode() or after applying sparsity constraints.

        Returns:
            torch.Tensor: Reconstructed activations with shape (..., input_dims).
                Should approximate the original input when z contains sufficient information.
        """
        if self.tie_weights:
            return F.linear(z, self.encoder.weight.t(), bias=None)
        else:
            return self.decoder(z)

    def post_step(self) -> None:
        """Renormalize decoder weights after each optimization step.

        Should be called after each parameter update (e.g., after optimizer.step()) to
        maintain the unit norm constraint on decoder weights. This ensures the constraint
        is enforced throughout training.

        Returns:
            None: Modifies decoder weights in-place if unit_norm_decoder is True.

        Note:
            This is a no-op if unit_norm_decoder is False or tie_weights is True.
        """
        self._renorm_decoder()

    def forward(self, x) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Perform a complete forward pass through the sparse autoencoder.

        Encodes input, applies sparsity constraint (top-k or none), and decodes to
        reconstruct the input. Returns both the reconstruction and intermediate activations.

        Args:
            x (torch.Tensor): Input activation tensor with shape (..., input_dims).

        Returns:
            tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
                - x_hat: Reconstructed activations with shape (..., input_dims)
                - z: Sparse latent activations after sparsity constraint, shape (..., n_features)
                - z_pre: Dense latent activations before sparsity constraint, shape (..., n_features)

        Note:
            If k is None, z and z_pre are identical. If k is set, z contains only the
            top-k activations (others are zeroed).
        """
        z_pre = self.encode(x)
        z = self.topk_mask(z_pre, self.k) if self.k is not None else z_pre
        x_hat = self.decode(z)
        return x_hat, z, z_pre

    def loss(self, x: torch.Tensor) -> tuple[torch.Tensor, dict]:
        """Compute the training loss with reconstruction and optional sparsity penalty.

        Calculates MSE reconstruction loss and, if using L1 sparsity (k is None), adds
        L1 penalty on latent activations. Also computes diagnostic metrics for logging.

        Args:
            x (torch.Tensor): Input activation tensor with shape (..., input_dims).

        Returns:
            tuple[torch.Tensor, dict]: A tuple containing:
                - total_loss: Scalar loss tensor for backpropagation. If k is None,
                    equals mse + beta_l1 * l1_sparsity. If k is set, equals mse only.
                - logs: Dictionary of detached metrics for logging:
                    - 'mse': Reconstruction loss (MSE between x_hat and x)
                    - 'l1': L1 norm of activations (only if k is None)
                    - 'non_zero_frac': Fraction of non-zero latent activations

        Note:
            The logs dictionary is intended for monitoring training progress and all
            values are detached from the computation graph.
        """
        x_hat, z, _ = self.forward(x)
        recon = F.mse_loss(x_hat, x)
        if self.k is None:
            sparsity = z.abs().mean()
            total = recon + self.beta_l1 * sparsity
            logs = {
                "mse": recon.detach(),
                "l1": sparsity.detach(),
                "non_zero_frac": (z != 0).float().mean().detach()
            }
        else:
            total = recon
            logs = {
                "mse": recon.detach(),
                "non_zero_frac": (z != 0).float().mean().detach()
            }
        return total, logs

    def topk_mask(self, z: torch.Tensor, k: int) -> torch.Tensor:
        """Apply top-k sparsity constraint by keeping only the k largest activations.

        Zeros out all but the k largest magnitude activations in the latent space,
        enforcing a fixed sparsity level. This is an alternative to L1 regularization
        that provides more predictable and controllable sparsity.

        Args:
            z (torch.Tensor): Dense latent activations with shape (..., n_features).
            k (int): Number of top activations to keep per sample. If None, ≤0, or
                ≥n_features, returns z unchanged.

        Returns:
            torch.Tensor: Sparse latent activations with same shape as z, but with all
                except the k largest (by absolute value) activations set to zero. The
                values of the top-k activations are preserved from the input.

        Note:
            Selection is based on absolute value, but original signed values are preserved
            for the top-k features.
        """
        if k is None or k <= 0 or k >= z.size(-1):
            return z
        vals, idx = torch.topk(z.abs(), k, dim=-1)
        out = torch.zeros_like(z)
        return out.scatter(-1, idx, z.gather(-1, idx))

__init__(input_dims=512, n_features=2048, activation='relu', input_norm=True, k=None, beta_l1=None, tie_weights=False, unit_norm_decoder=True)

Initialize the Sparse Autoencoder with specified architecture and sparsity settings.

Parameters:

Name	Type	Description	Default
input_dims	int	Dimensionality of input activations (e.g., 3072 for GPT-2 layer 3 MLP output). Defaults to 512.	512
n_features	int	Number of learned features in the latent space. Typically set as expansion_factor × input_dims where expansion_factor is 2-8. Defaults to 2048.	2048
activation	str	Nonlinearity applied after encoder. Must be 'relu' or 'silu'. ReLU is standard for interpretability; SiLU may improve reconstruction. Defaults to "relu".	'relu'
input_norm	bool	Whether to apply LayerNorm to inputs before encoding. Helps stabilize training with varying activation magnitudes. Defaults to True.	True
k	int | None	If set, uses top-k sparsity (keeps only k largest activations per sample) instead of L1 regularization. Mutually exclusive with beta_l1. Defaults to None.	None
beta_l1	float | None	L1 regularization coefficient for sparsity. Higher values encourage sparser activations. Ignored if k is set. Defaults to None.	None
tie_weights	bool	If True, decoder weights are the transpose of encoder weights (no separate decoder parameters). Reduces parameters but may hurt performance. Defaults to False.	False
unit_norm_decoder	bool	If True, constrains decoder weight columns to unit norm. Improves feature interpretability by removing scale ambiguity. Defaults to True.	True

Source code in deeplens\sae.py

def __init__(
        self, 
        input_dims: int = 512, 
        n_features: int = 2048, 
        activation: str = "relu",
        input_norm: bool = True,
        k: int | None = None,
        beta_l1: float | None = None,
        tie_weights: bool = False,
        unit_norm_decoder: bool = True
    ) -> None:
    """Initialize the Sparse Autoencoder with specified architecture and sparsity settings.

    Args:
        input_dims (int, optional): Dimensionality of input activations (e.g., 3072 for
            GPT-2 layer 3 MLP output). Defaults to 512.
        n_features (int, optional): Number of learned features in the latent space.
            Typically set as expansion_factor × input_dims where expansion_factor is 2-8.
            Defaults to 2048.
        activation (str, optional): Nonlinearity applied after encoder. Must be 'relu'
            or 'silu'. ReLU is standard for interpretability; SiLU may improve reconstruction.
            Defaults to "relu".
        input_norm (bool, optional): Whether to apply LayerNorm to inputs before encoding.
            Helps stabilize training with varying activation magnitudes. Defaults to True.
        k (int | None, optional): If set, uses top-k sparsity (keeps only k largest activations
            per sample) instead of L1 regularization. Mutually exclusive with beta_l1.
            Defaults to None.
        beta_l1 (float | None, optional): L1 regularization coefficient for sparsity.
            Higher values encourage sparser activations. Ignored if k is set. Defaults to None.
        tie_weights (bool, optional): If True, decoder weights are the transpose of encoder
            weights (no separate decoder parameters). Reduces parameters but may hurt performance.
            Defaults to False.
        unit_norm_decoder (bool, optional): If True, constrains decoder weight columns to
            unit norm. Improves feature interpretability by removing scale ambiguity.
            Defaults to True.
    """
    super().__init__()
    self.norm = nn.LayerNorm(input_dims) if input_norm else nn.Identity()
    self.encoder = nn.Linear(input_dims, n_features, bias=True)
    self.decoder = None if tie_weights else nn.Linear(n_features, input_dims, bias=False)
    self.unit_norm_decoder = unit_norm_decoder
    self.input_norm = input_norm

    if activation == "relu":
        self.activation = nn.ReLU()
        kaiming_activation = "relu"
    elif activation == "silu":
        self.activation = nn.SiLU()
        kaiming_activation = "linear"
    else:
        raise ValueError("Activation must be 'relu' or 'silu'")

    nn.init.kaiming_normal_(self.encoder.weight, nonlinearity=kaiming_activation)
    if self.decoder is not None:
        nn.init.xavier_uniform_(self.decoder.weight)
        if self.unit_norm_decoder:
            self._renorm_decoder()

    self.k = k
    self.beta_l1 = beta_l1
    self.tie_weights = tie_weights

decode(z)

Decode sparse latent features back to the original activation space.

Applies linear transformation from feature space back to input space. Uses either a separate decoder or transposed encoder weights depending on tie_weights setting.

Parameters:

Name	Type	Description	Default
z	Tensor	Latent feature activations with shape (..., n_features). Typically the output of encode() or after applying sparsity constraints.	required

Returns:

Type	Description
Tensor	torch.Tensor: Reconstructed activations with shape (..., input_dims). Should approximate the original input when z contains sufficient information.

Source code in deeplens\sae.py

def decode(self, z) -> torch.Tensor:
    """Decode sparse latent features back to the original activation space.

    Applies linear transformation from feature space back to input space. Uses either
    a separate decoder or transposed encoder weights depending on tie_weights setting.

    Args:
        z (torch.Tensor): Latent feature activations with shape (..., n_features).
            Typically the output of encode() or after applying sparsity constraints.

    Returns:
        torch.Tensor: Reconstructed activations with shape (..., input_dims).
            Should approximate the original input when z contains sufficient information.
    """
    if self.tie_weights:
        return F.linear(z, self.encoder.weight.t(), bias=None)
    else:
        return self.decoder(z)

encode(x)

Encode input activations into the sparse latent feature space.

Applies optional normalization, linear transformation to expanded feature space, and nonlinear activation function.

Parameters:

Name	Type	Description	Default
x	Tensor	Input activation tensor with shape (..., input_dims). Typically (batch_size, seq_length, input_dims) or (batch_size, input_dims).	required

Returns:

Type	Description
Tensor	torch.Tensor: Encoded features with shape (..., n_features) after applying normalization (if enabled), linear encoding, and activation function. These are the pre-sparsity latent activations.

Source code in deeplens\sae.py

def encode(self, x) -> torch.Tensor:
    """Encode input activations into the sparse latent feature space.

    Applies optional normalization, linear transformation to expanded feature space,
    and nonlinear activation function.

    Args:
        x (torch.Tensor): Input activation tensor with shape (..., input_dims).
            Typically (batch_size, seq_length, input_dims) or (batch_size, input_dims).

    Returns:
        torch.Tensor: Encoded features with shape (..., n_features) after applying
            normalization (if enabled), linear encoding, and activation function.
            These are the pre-sparsity latent activations.
    """
    x = self.norm(x)
    return self.activation(self.encoder(x))

forward(x)

Perform a complete forward pass through the sparse autoencoder.

Encodes input, applies sparsity constraint (top-k or none), and decodes to reconstruct the input. Returns both the reconstruction and intermediate activations.

Parameters:

Name	Type	Description	Default
x	Tensor	Input activation tensor with shape (..., input_dims).	required

Returns:

Type	Description
tuple[Tensor, Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing: - x_hat: Reconstructed activations with shape (..., input_dims) - z: Sparse latent activations after sparsity constraint, shape (..., n_features) - z_pre: Dense latent activations before sparsity constraint, shape (..., n_features)

Note

If k is None, z and z_pre are identical. If k is set, z contains only the top-k activations (others are zeroed).

Source code in deeplens\sae.py

def forward(self, x) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Perform a complete forward pass through the sparse autoencoder.

    Encodes input, applies sparsity constraint (top-k or none), and decodes to
    reconstruct the input. Returns both the reconstruction and intermediate activations.

    Args:
        x (torch.Tensor): Input activation tensor with shape (..., input_dims).

    Returns:
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
            - x_hat: Reconstructed activations with shape (..., input_dims)
            - z: Sparse latent activations after sparsity constraint, shape (..., n_features)
            - z_pre: Dense latent activations before sparsity constraint, shape (..., n_features)

    Note:
        If k is None, z and z_pre are identical. If k is set, z contains only the
        top-k activations (others are zeroed).
    """
    z_pre = self.encode(x)
    z = self.topk_mask(z_pre, self.k) if self.k is not None else z_pre
    x_hat = self.decode(z)
    return x_hat, z, z_pre

loss(x)

Compute the training loss with reconstruction and optional sparsity penalty.

Calculates MSE reconstruction loss and, if using L1 sparsity (k is None), adds L1 penalty on latent activations. Also computes diagnostic metrics for logging.

Parameters:

Name	Type	Description	Default
x	Tensor	Input activation tensor with shape (..., input_dims).	required

Returns:

Type	Description
tuple[Tensor, dict]	tuple[torch.Tensor, dict]: A tuple containing: - total_loss: Scalar loss tensor for backpropagation. If k is None, equals mse + beta_l1 * l1_sparsity. If k is set, equals mse only. - logs: Dictionary of detached metrics for logging: - 'mse': Reconstruction loss (MSE between x_hat and x) - 'l1': L1 norm of activations (only if k is None) - 'non_zero_frac': Fraction of non-zero latent activations

Note

The logs dictionary is intended for monitoring training progress and all values are detached from the computation graph.

Source code in deeplens\sae.py

def loss(self, x: torch.Tensor) -> tuple[torch.Tensor, dict]:
    """Compute the training loss with reconstruction and optional sparsity penalty.

    Calculates MSE reconstruction loss and, if using L1 sparsity (k is None), adds
    L1 penalty on latent activations. Also computes diagnostic metrics for logging.

    Args:
        x (torch.Tensor): Input activation tensor with shape (..., input_dims).

    Returns:
        tuple[torch.Tensor, dict]: A tuple containing:
            - total_loss: Scalar loss tensor for backpropagation. If k is None,
                equals mse + beta_l1 * l1_sparsity. If k is set, equals mse only.
            - logs: Dictionary of detached metrics for logging:
                - 'mse': Reconstruction loss (MSE between x_hat and x)
                - 'l1': L1 norm of activations (only if k is None)
                - 'non_zero_frac': Fraction of non-zero latent activations

    Note:
        The logs dictionary is intended for monitoring training progress and all
        values are detached from the computation graph.
    """
    x_hat, z, _ = self.forward(x)
    recon = F.mse_loss(x_hat, x)
    if self.k is None:
        sparsity = z.abs().mean()
        total = recon + self.beta_l1 * sparsity
        logs = {
            "mse": recon.detach(),
            "l1": sparsity.detach(),
            "non_zero_frac": (z != 0).float().mean().detach()
        }
    else:
        total = recon
        logs = {
            "mse": recon.detach(),
            "non_zero_frac": (z != 0).float().mean().detach()
        }
    return total, logs

post_step()

Renormalize decoder weights after each optimization step.

Should be called after each parameter update (e.g., after optimizer.step()) to maintain the unit norm constraint on decoder weights. This ensures the constraint is enforced throughout training.

Returns:

Name	Type	Description
None	None	Modifies decoder weights in-place if unit_norm_decoder is True.

Note

This is a no-op if unit_norm_decoder is False or tie_weights is True.

Source code in deeplens\sae.py

def post_step(self) -> None:
    """Renormalize decoder weights after each optimization step.

    Should be called after each parameter update (e.g., after optimizer.step()) to
    maintain the unit norm constraint on decoder weights. This ensures the constraint
    is enforced throughout training.

    Returns:
        None: Modifies decoder weights in-place if unit_norm_decoder is True.

    Note:
        This is a no-op if unit_norm_decoder is False or tie_weights is True.
    """
    self._renorm_decoder()

topk_mask(z, k)

Apply top-k sparsity constraint by keeping only the k largest activations.

Zeros out all but the k largest magnitude activations in the latent space, enforcing a fixed sparsity level. This is an alternative to L1 regularization that provides more predictable and controllable sparsity.

Parameters:

Name	Type	Description	Default
z	Tensor	Dense latent activations with shape (..., n_features).	required
k	int	Number of top activations to keep per sample. If None, ≤0, or ≥n_features, returns z unchanged.	required

Returns:

Type	Description
Tensor	torch.Tensor: Sparse latent activations with same shape as z, but with all except the k largest (by absolute value) activations set to zero. The values of the top-k activations are preserved from the input.

Note

Selection is based on absolute value, but original signed values are preserved for the top-k features.

Source code in deeplens\sae.py

def topk_mask(self, z: torch.Tensor, k: int) -> torch.Tensor:
    """Apply top-k sparsity constraint by keeping only the k largest activations.

    Zeros out all but the k largest magnitude activations in the latent space,
    enforcing a fixed sparsity level. This is an alternative to L1 regularization
    that provides more predictable and controllable sparsity.

    Args:
        z (torch.Tensor): Dense latent activations with shape (..., n_features).
        k (int): Number of top activations to keep per sample. If None, ≤0, or
            ≥n_features, returns z unchanged.

    Returns:
        torch.Tensor: Sparse latent activations with same shape as z, but with all
            except the k largest (by absolute value) activations set to zero. The
            values of the top-k activations are preserved from the input.

    Note:
        Selection is based on absolute value, but original signed values are preserved
        for the top-k features.
    """
    if k is None or k <= 0 or k >= z.size(-1):
        return z
    vals, idx = torch.topk(z.abs(), k, dim=-1)
    out = torch.zeros_like(z)
    return out.scatter(-1, idx, z.gather(-1, idx))

deeplens.train

SAETrainer

Training framework for Sparse Autoencoders with comprehensive logging and checkpointing.

Handles the complete training loop including optimization, learning rate scheduling, gradient clipping, mixed precision training (bfloat16), periodic evaluation, model checkpointing, and Weights & Biases logging. Designed specifically for training sparse autoencoders on neural network activation data.

Source code in deeplens\train.py

class SAETrainer():
    """Training framework for Sparse Autoencoders with comprehensive logging and checkpointing.

    Handles the complete training loop including optimization, learning rate scheduling,
    gradient clipping, mixed precision training (bfloat16), periodic evaluation, model
    checkpointing, and Weights & Biases logging. Designed specifically for training
    sparse autoencoders on neural network activation data.
    """
    def __init__(
            self, 
            train_dataloader: DataLoader = None, 
            eval_dataloader: DataLoader = None, 
            model: torch.nn.Module = None, 
            model_name: str = "sae",
            optim: torch.optim.Optimizer = torch.optim.Adam,
            epochs: int = 20, 
            bf16: bool = False,
            random_seed: int = 42,
            save_checkpoints: bool = True,
            device: str = "auto",
            grad_clip_norm: float = None,
            lrs_type: str = None,
            eval_steps: int = 5000,
            warmup_fraction: float = 0.1,
            save_best_only: bool = True,
            log_to_wandb: bool = True
        ) -> None:
        """Initialize the SAETrainer with model, data, and training configuration.

        Sets up the training environment including device placement, learning rate scheduling,
        and Weights & Biases logging if enabled.

        Args:
            train_dataloader (DataLoader, optional): DataLoader for training data containing
                batches of activation tensors. Defaults to None.
            eval_dataloader (DataLoader, optional): DataLoader for evaluation/validation data.
                Used for periodic evaluation during training. Defaults to None.
            model (torch.nn.Module, optional): Sparse autoencoder model to train. Should be
                an instance of SparseAutoencoder or compatible architecture. Defaults to None.
            model_name (str, optional): Name used for organizing saved checkpoints and W&B
                project naming. Creates directory structure at saved_models/{model_name}/.
                Defaults to "sae".
            optim (torch.optim.Optimizer, optional): Initialized optimizer instance with
                model parameters already attached. Defaults to torch.optim.Adam.
            epochs (int, optional): Number of complete passes through the training dataset.
                Defaults to 20.
            bf16 (bool, optional): Whether to use bfloat16 mixed precision training for
                faster computation and reduced memory usage. Requires CUDA support.
                Defaults to False.
            random_seed (int, optional): Random seed for reproducibility across numpy,
                PyTorch, and CUDNN. Defaults to 42.
            save_checkpoints (bool, optional): Whether to save model checkpoints during
                training. Checkpoints saved when evaluation loss improves. Defaults to True.
            device (str, optional): Device for training. Can be "auto" for automatic selection
                (prefers cuda > mps > cpu), "cuda", "mps", or "cpu". Defaults to "auto".
            grad_clip_norm (float, optional): Maximum gradient norm for gradient clipping.
                Helps prevent gradient explosion. If None, no clipping is applied.
                Defaults to None.
            lrs_type (str, optional): Learning rate scheduler type. Options: 'cosine',
                'plateau', 'linear'. If None, uses constant learning rate. Defaults to None.
            eval_steps (int, optional): Evaluate model every N training steps and save
                checkpoint if loss improves. Defaults to 5000.
            warmup_fraction (float, optional): Fraction of total training steps to use for
                learning rate warmup. Applies to 'cosine' and 'linear' schedulers.
                Defaults to 0.1 (10% warmup).
            save_best_only (bool, optional): If True, only saves the best checkpoint
                (overwrites previous best). If False, saves all improving checkpoints.
                Defaults to True.
            log_to_wandb (bool, optional): Whether to log metrics to Weights & Biases.
                Logs training/eval loss, non-zero fraction, learning rate, and config.
                Defaults to True.

        Note:
            The optimizer must be initialized with the model's parameters before passing
            to the trainer. Learning rate and other optimizer settings should be configured
            in the optimizer instance.
        """
        self.model = model
        self.model_name = model_name
        self.optim = optim
        self.epochs = epochs
        self.train_dataloader = train_dataloader
        self.eval_dataloader = eval_dataloader
        self.bf16 = bf16
        self.random_seed = random_seed
        self.save_checkpoints = save_checkpoints
        self.grad_clip_norm = grad_clip_norm
        self.eval_steps = eval_steps
        self.warmup_fraction = warmup_fraction
        self.save_best_only = save_best_only
        self.log_wandb = log_to_wandb

        self.device = get_device(device)
        print(f"Running on device: {self.device}")

        if lrs_type is not None:
            self.scheduler = self.set_lr_scheduler(lrs_type)
        else:
            self.scheduler = None

        if log_to_wandb:
            time = datetime.now().strftime("%Y%m%d_%H%M%S")
            wandb.init(
                project=f"sparse-autoencoder",
                name=f"run-{self.model_name}-{time}",
                config={
                    "epochs": epochs,
                    "in_dims": self.model.encoder.in_features,
                    "sparse_dims": self.model.encoder.out_features,
                    "activations": self.model.activation,
                    "input_norm": self.model.input_norm,
                    "top_k": self.model.k, 
                    "l1": self.model.beta_l1,
                    "tie_weights": self.model.tie_weights,
                    "unit_norm_decoder": self.model.unit_norm_decoder,
                    "lr": optim.param_groups[0]['lr'],
                    "lr_scheduler": lrs_type,
                    "warmup_fraction": warmup_fraction,
                    "seed": random_seed,
                    "grad_clip_norm": grad_clip_norm,
                    "bf16": bf16
                }
            )

    def train_one_epoch(
            self, 
            model: torch.nn.Module, 
            train_dataloader: torch.utils.data.DataLoader, 
            optim: torch.optim.Optimizer, 
            bf16: bool = True,
            global_step: int = 0,
            timestamp: str = None,
            best_loss: float = float('inf')
        ) -> tuple[int, float]:
        """Execute one complete training epoch with periodic evaluation and checkpointing.

        Iterates through all training batches, performs forward/backward passes with optional
        mixed precision, applies gradient clipping, updates learning rate, and periodically
        evaluates and saves checkpoints.

        Args:
            model (torch.nn.Module): The sparse autoencoder model to train.
            train_dataloader (torch.utils.data.DataLoader): DataLoader providing training batches.
            optim (torch.optim.Optimizer): Optimizer for updating model parameters.
            bf16 (bool, optional): Whether to use bfloat16 mixed precision. Defaults to True.
            global_step (int, optional): Current global training step count across all epochs.
                Used for logging and checkpoint naming. Defaults to 0.
            timestamp (str, optional): Timestamp string for checkpoint directory naming.
                Format: "YYYYMMDD_HHMMSS". Defaults to None.
            best_loss (float, optional): Best evaluation loss achieved so far. Used to
                determine when to save new checkpoints. Defaults to float('inf').

        Returns:
            tuple[int, float]: A tuple containing:
                - global_step: Updated global step count after this epoch
                - best_loss: Updated best evaluation loss (may be unchanged)

        Note:
            Prints training metrics every 100 steps. Evaluates every eval_steps and saves
            checkpoints when evaluation loss improves. Logs to W&B if enabled.
        """
        model.train()

        if bf16:
            scaler = torch.amp.GradScaler("cuda")
            device_type = "cuda" if torch.cuda.is_available() else "cpu"

        for idx, inputs in enumerate(train_dataloader):
            optim.zero_grad()
            if torch.cuda.is_available() and bf16:
                with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
                    inputs = inputs.to(self.device)
                    loss, logs = model.loss(inputs)
                scaler.scale(loss).backward()

                if self.grad_clip_norm is not None:
                    scaler.unscale_(optim)
                    torch.nn.utils.clip_grad_norm_(
                        model.parameters(), self.grad_clip_norm
                    )

                scaler.step(optim)
                scaler.update()

            else:
                inputs = inputs.to(self.device)
                loss, logs = model.loss(inputs)

                if self.grad_clip_norm is not None:
                    torch.nn.utils.clip_grad_norm_(
                        model.parameters(), self.grad_clip_norm
                    )

                loss.backward()
                optim.step()

            if self.scheduler is not None and not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau):
                self.scheduler.step()

            model.post_step()
            global_step += 1

            if self.log_wandb:
                wandb.log({
                    "train/loss": logs['mse'].item(),
                    "train/non_zero_frac": logs['non_zero_frac'].item(),
                    "train/lr": self.optim.param_groups[0]['lr'],
                    "global_step": global_step
                }, step=global_step)

            if (idx % 100) == 0:
                current_lr = self.optim.param_groups[0]['lr']
                print(
                    f"Step [{idx}/{len(train_dataloader)}] - "
                    f"train_loss: {round(logs['mse'].item(), 3)} - "
                    f"train_nz_frac: {round(logs['non_zero_frac'].item(), 3)} - "
                    f"lr: {current_lr:.2e}"
                )

            if global_step % self.eval_steps == 0:
                print(f"\n{'='*60}")
                print(f"Intermediate Evaluation at step {global_step}")
                print(f"{'='*60}")
                eval_loss = self.evaluate(
                    model=model,
                    eval_dataloader=self.eval_dataloader,
                    bf16=bf16
                )

                if self.log_wandb:
                    wandb.log({
                        "eval/loss": eval_loss,
                        "global_step": global_step
                    }, step=global_step)

                if self.save_checkpoints and eval_loss < best_loss:
                    if self.save_best_only:
                        save_path = f"saved_models/{self.model_name}/run_{timestamp}/best_model.pt"
                    else:
                        save_path = f"saved_models/{self.model_name}/run_{timestamp}/sae_step_{global_step}.pt"
                    torch.save(model.state_dict(), save_path)
                    print(f"New best model saved (loss: {eval_loss:.6f})")
                    best_loss = eval_loss

                model.train()

        return global_step, best_loss

    @torch.no_grad()
    def evaluate(
            self,
            model: torch.nn.Module, 
            eval_dataloader: torch.utils.data.DataLoader, 
            bf16: bool = True
        ) -> float:
        """Evaluate the model on the validation dataset without gradient computation.

        Computes average loss across all evaluation batches with optional mixed precision.
        Useful for monitoring generalization and preventing overfitting.

        Args:
            model (torch.nn.Module): The sparse autoencoder model to evaluate.
            eval_dataloader (torch.utils.data.DataLoader): DataLoader providing evaluation batches.
            bf16 (bool, optional): Whether to use bfloat16 mixed precision for evaluation.
                Should match training setting. Defaults to True.

        Returns:
            float: Average evaluation loss across all batches. Computed as total loss
                divided by number of batches.

        Note:
            Prints evaluation metrics every 100 steps. Sets model to eval mode and
            restores training mode afterward if called during training.
        """
        model.eval()
        n_batches = 0
        total_loss = 0.0

        if bf16:
            device_type = "cuda" if torch.cuda.is_available() else "cpu"

        for idx, inputs in enumerate(eval_dataloader):
            if torch.cuda.is_available() and bf16:
                with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
                    inputs = inputs.to(self.device)
                    loss, logs = model.loss(inputs)
            else:
                inputs = inputs.to(self.device)
                loss, logs = model.loss(inputs)

            total_loss += loss.item()
            n_batches += 1

            if (idx % 100) == 0:
                current_lr = self.optim.param_groups[0]['lr']
                print(
                    f"Step [{idx}/{len(eval_dataloader)}] - "
                    f"eval_loss: {round(logs['mse'].item(), 3)} - "
                    f"eval_nz_frac: {round(logs['non_zero_frac'].item(), 3)} - "
                    f"lr: {current_lr:.2e}"
                )

        avg_loss = total_loss / n_batches
        print(f"Avg loss: {avg_loss:.3f}")
        return avg_loss

    def train(self) -> None:
        """Execute the complete training loop for all epochs.

        Main training orchestrator that:
        1. Sets random seed for reproducibility
        2. Moves model to appropriate device
        3. Creates checkpoint directories if saving enabled
        4. Runs training epochs with periodic evaluation
        5. Saves best models based on evaluation loss
        6. Applies learning rate scheduling
        7. Logs metrics to W&B if enabled
        8. Finalizes W&B run on completion

        Note:
            Checkpoints saved to saved_models/{model_name}/run_{timestamp}/.
            Each epoch includes full pass through training data followed by evaluation.
            Best model determined by lowest evaluation loss.
        """
        self.set_seed(self.random_seed)
        self.model.to(self.device)

        if self.save_checkpoints:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            os.makedirs(f"saved_models/{self.model_name}/run_{timestamp}", exist_ok=True)

        best_loss = float('inf')
        global_step = 0

        for epoch in range(self.epochs):
            print(f"\nEpoch [{epoch+1}/{self.epochs}]")
            global_step, best_loss = self.train_one_epoch(
                model=self.model, 
                train_dataloader=self.train_dataloader, 
                optim=self.optim, 
                bf16=self.bf16,
                global_step=global_step,
                timestamp=timestamp,
                best_loss=best_loss
            )

            print(f"\n{'='*60}")
            print(f"End of epoch {epoch+1} evaluation")
            print(f"{'='*60}")
            loss = self.evaluate(
                model=self.model,
                eval_dataloader=self.eval_dataloader,
                bf16=self.bf16
            )

            if self.log_wandb:
                wandb.log({
                    "epoch": epoch + 1,
                    "eval/epoch_loss": loss,
                }, step=global_step)

            if self.save_checkpoints and loss < best_loss:
                if self.save_best_only:
                    save_path = f"saved_models/{self.model_name}/run_{timestamp}/best_model.pt"
                else:
                    save_path = f"saved_models/{self.model_name}/run_{timestamp}/sae_epoch_{epoch+1}.pt"
                torch.save(self.model.state_dict(), save_path)
                print(f"New best model saved (loss: {loss:.3f})")
                best_loss = loss

            if self.scheduler is not None and isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau):
                self.scheduler.step(loss)

        if self.log_wandb:
            wandb.finish()

        print("Finished training!")

    def set_lr_scheduler(self, lr_type: str = 'cosine') -> lr_scheduler:
        """Configure and initialize learning rate scheduler with warmup.

        Creates a learning rate scheduler based on the specified type. 'cosine' and 'linear'
        schedulers include warmup phase for training stability. 'plateau' reduces learning
        rate when validation loss plateaus.

        Args:
            lr_type (str, optional): Type of learning rate scheduler. Options:
                - 'cosine': Linear warmup followed by cosine annealing to 10% of initial LR
                - 'plateau': Reduces LR by factor when validation loss stops improving
                - 'linear': Linear warmup followed by linear decay to 10% of initial LR
                Defaults to 'cosine'.

        Returns:
            lr_scheduler: Configured PyTorch learning rate scheduler instance.

        Note:
            For 'cosine' and 'linear', warmup steps = warmup_fraction * total_steps.
            'plateau' scheduler requires manual .step(loss) calls, which are handled
            automatically at epoch end.
        """
        assert lr_type in ['cosine', 'plateau', 'linear'], "Use 'cosine', 'plateau', or 'linear'"

        total_steps = self.epochs * len(self.train_dataloader)
        warmup_steps = int(total_steps * self.warmup_fraction)

        if lr_type == 'cosine':
            warmup = lr_scheduler.LinearLR(
                self.optim, start_factor=0.01, end_factor=1.0,
                total_iters=warmup_steps
            )
            cosine = lr_scheduler.CosineAnnealingLR(
                self.optim, T_max=total_steps - warmup_steps,
                eta_min=self.optim.param_groups[0]['lr'] * 0.1
            )
            self.scheduler = lr_scheduler.SequentialLR(
                self.optim, schedulers=[warmup, cosine],
                milestones=[warmup_steps]
            )
        elif lr_type == 'plateau':
            self.scheduler = lr_scheduler.ReduceLROnPlateau(
                self.optim, patience=10, threshold=1e-4, min_lr=1e-6
            )
        else:
            warmup = lr_scheduler.LinearLR(
                self.optim, start_factor=0.01, end_factor=1.0,
                total_iters=warmup_steps
            )
            decay = lr_scheduler.LinearLR(
                self.optim, start_factor=1.0, end_factor=0.1,
                total_iters=total_steps - warmup_steps
            )
            self.scheduler = lr_scheduler.SequentialLR(
                self.optim, schedulers=[warmup, decay],
                milestones=[warmup_steps]
            )
        return self.scheduler

    def set_seed(self, seed: int = 1) -> None:
        """Set random seeds for reproducible training across all libraries.

        Configures random number generators for NumPy, PyTorch (CPU and CUDA), and
        makes CUDNN operations deterministic. Essential for experiment reproducibility.

        Args:
            seed (int, optional): Random seed value to use. Same seed should produce
                identical results across runs (assuming same hardware/software).
                Defaults to 1.

        Note:
            Setting deterministic=True may reduce performance. Disables CUDNN benchmarking
            for reproducibility.
        """
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        print(f"Using random seed: {seed}")

    @staticmethod
    def config_from_yaml(file: str) -> dict:
        """Load sparse autoencoder configuration from a YAML file.

        Static utility method for loading model or training configurations from YAML.
        Useful for maintaining configuration files separately from code.

        Args:
            file (str): Path to the YAML configuration file.

        Returns:
            dict: Dictionary containing parsed configuration parameters.
        """
        with open(file, "r") as f:
            config = yaml.safe_load(f)
        return config

__init__(train_dataloader=None, eval_dataloader=None, model=None, model_name='sae', optim=torch.optim.Adam, epochs=20, bf16=False, random_seed=42, save_checkpoints=True, device='auto', grad_clip_norm=None, lrs_type=None, eval_steps=5000, warmup_fraction=0.1, save_best_only=True, log_to_wandb=True)

Initialize the SAETrainer with model, data, and training configuration.

Sets up the training environment including device placement, learning rate scheduling, and Weights & Biases logging if enabled.

Parameters:

Name	Type	Description	Default
train_dataloader	DataLoader	DataLoader for training data containing batches of activation tensors. Defaults to None.	None
eval_dataloader	DataLoader	DataLoader for evaluation/validation data. Used for periodic evaluation during training. Defaults to None.	None
model	Module	Sparse autoencoder model to train. Should be an instance of SparseAutoencoder or compatible architecture. Defaults to None.	None
model_name	str	Name used for organizing saved checkpoints and W&B project naming. Creates directory structure at saved_models/{model_name}/. Defaults to "sae".	'sae'
optim	Optimizer	Initialized optimizer instance with model parameters already attached. Defaults to torch.optim.Adam.	Adam
epochs	int	Number of complete passes through the training dataset. Defaults to 20.	20
bf16	bool	Whether to use bfloat16 mixed precision training for faster computation and reduced memory usage. Requires CUDA support. Defaults to False.	False
random_seed	int	Random seed for reproducibility across numpy, PyTorch, and CUDNN. Defaults to 42.	42
save_checkpoints	bool	Whether to save model checkpoints during training. Checkpoints saved when evaluation loss improves. Defaults to True.	True
device	str	Device for training. Can be "auto" for automatic selection (prefers cuda > mps > cpu), "cuda", "mps", or "cpu". Defaults to "auto".	'auto'
grad_clip_norm	float	Maximum gradient norm for gradient clipping. Helps prevent gradient explosion. If None, no clipping is applied. Defaults to None.	None
lrs_type	str	Learning rate scheduler type. Options: 'cosine', 'plateau', 'linear'. If None, uses constant learning rate. Defaults to None.	None
eval_steps	int	Evaluate model every N training steps and save checkpoint if loss improves. Defaults to 5000.	5000
warmup_fraction	float	Fraction of total training steps to use for learning rate warmup. Applies to 'cosine' and 'linear' schedulers. Defaults to 0.1 (10% warmup).	0.1
save_best_only	bool	If True, only saves the best checkpoint (overwrites previous best). If False, saves all improving checkpoints. Defaults to True.	True
log_to_wandb	bool	Whether to log metrics to Weights & Biases. Logs training/eval loss, non-zero fraction, learning rate, and config. Defaults to True.	True

Note

The optimizer must be initialized with the model's parameters before passing to the trainer. Learning rate and other optimizer settings should be configured in the optimizer instance.

Source code in deeplens\train.py

def __init__(
        self, 
        train_dataloader: DataLoader = None, 
        eval_dataloader: DataLoader = None, 
        model: torch.nn.Module = None, 
        model_name: str = "sae",
        optim: torch.optim.Optimizer = torch.optim.Adam,
        epochs: int = 20, 
        bf16: bool = False,
        random_seed: int = 42,
        save_checkpoints: bool = True,
        device: str = "auto",
        grad_clip_norm: float = None,
        lrs_type: str = None,
        eval_steps: int = 5000,
        warmup_fraction: float = 0.1,
        save_best_only: bool = True,
        log_to_wandb: bool = True
    ) -> None:
    """Initialize the SAETrainer with model, data, and training configuration.

    Sets up the training environment including device placement, learning rate scheduling,
    and Weights & Biases logging if enabled.

    Args:
        train_dataloader (DataLoader, optional): DataLoader for training data containing
            batches of activation tensors. Defaults to None.
        eval_dataloader (DataLoader, optional): DataLoader for evaluation/validation data.
            Used for periodic evaluation during training. Defaults to None.
        model (torch.nn.Module, optional): Sparse autoencoder model to train. Should be
            an instance of SparseAutoencoder or compatible architecture. Defaults to None.
        model_name (str, optional): Name used for organizing saved checkpoints and W&B
            project naming. Creates directory structure at saved_models/{model_name}/.
            Defaults to "sae".
        optim (torch.optim.Optimizer, optional): Initialized optimizer instance with
            model parameters already attached. Defaults to torch.optim.Adam.
        epochs (int, optional): Number of complete passes through the training dataset.
            Defaults to 20.
        bf16 (bool, optional): Whether to use bfloat16 mixed precision training for
            faster computation and reduced memory usage. Requires CUDA support.
            Defaults to False.
        random_seed (int, optional): Random seed for reproducibility across numpy,
            PyTorch, and CUDNN. Defaults to 42.
        save_checkpoints (bool, optional): Whether to save model checkpoints during
            training. Checkpoints saved when evaluation loss improves. Defaults to True.
        device (str, optional): Device for training. Can be "auto" for automatic selection
            (prefers cuda > mps > cpu), "cuda", "mps", or "cpu". Defaults to "auto".
        grad_clip_norm (float, optional): Maximum gradient norm for gradient clipping.
            Helps prevent gradient explosion. If None, no clipping is applied.
            Defaults to None.
        lrs_type (str, optional): Learning rate scheduler type. Options: 'cosine',
            'plateau', 'linear'. If None, uses constant learning rate. Defaults to None.
        eval_steps (int, optional): Evaluate model every N training steps and save
            checkpoint if loss improves. Defaults to 5000.
        warmup_fraction (float, optional): Fraction of total training steps to use for
            learning rate warmup. Applies to 'cosine' and 'linear' schedulers.
            Defaults to 0.1 (10% warmup).
        save_best_only (bool, optional): If True, only saves the best checkpoint
            (overwrites previous best). If False, saves all improving checkpoints.
            Defaults to True.
        log_to_wandb (bool, optional): Whether to log metrics to Weights & Biases.
            Logs training/eval loss, non-zero fraction, learning rate, and config.
            Defaults to True.

    Note:
        The optimizer must be initialized with the model's parameters before passing
        to the trainer. Learning rate and other optimizer settings should be configured
        in the optimizer instance.
    """
    self.model = model
    self.model_name = model_name
    self.optim = optim
    self.epochs = epochs
    self.train_dataloader = train_dataloader
    self.eval_dataloader = eval_dataloader
    self.bf16 = bf16
    self.random_seed = random_seed
    self.save_checkpoints = save_checkpoints
    self.grad_clip_norm = grad_clip_norm
    self.eval_steps = eval_steps
    self.warmup_fraction = warmup_fraction
    self.save_best_only = save_best_only
    self.log_wandb = log_to_wandb

    self.device = get_device(device)
    print(f"Running on device: {self.device}")

    if lrs_type is not None:
        self.scheduler = self.set_lr_scheduler(lrs_type)
    else:
        self.scheduler = None

    if log_to_wandb:
        time = datetime.now().strftime("%Y%m%d_%H%M%S")
        wandb.init(
            project=f"sparse-autoencoder",
            name=f"run-{self.model_name}-{time}",
            config={
                "epochs": epochs,
                "in_dims": self.model.encoder.in_features,
                "sparse_dims": self.model.encoder.out_features,
                "activations": self.model.activation,
                "input_norm": self.model.input_norm,
                "top_k": self.model.k, 
                "l1": self.model.beta_l1,
                "tie_weights": self.model.tie_weights,
                "unit_norm_decoder": self.model.unit_norm_decoder,
                "lr": optim.param_groups[0]['lr'],
                "lr_scheduler": lrs_type,
                "warmup_fraction": warmup_fraction,
                "seed": random_seed,
                "grad_clip_norm": grad_clip_norm,
                "bf16": bf16
            }
        )

config_from_yaml(file) staticmethod

Load sparse autoencoder configuration from a YAML file.

Static utility method for loading model or training configurations from YAML. Useful for maintaining configuration files separately from code.

Parameters:

Name	Type	Description	Default
file	str	Path to the YAML configuration file.	required

Returns:

Name	Type	Description
dict	dict	Dictionary containing parsed configuration parameters.

Source code in deeplens\train.py

@staticmethod
def config_from_yaml(file: str) -> dict:
    """Load sparse autoencoder configuration from a YAML file.

    Static utility method for loading model or training configurations from YAML.
    Useful for maintaining configuration files separately from code.

    Args:
        file (str): Path to the YAML configuration file.

    Returns:
        dict: Dictionary containing parsed configuration parameters.
    """
    with open(file, "r") as f:
        config = yaml.safe_load(f)
    return config

evaluate(model, eval_dataloader, bf16=True)

Evaluate the model on the validation dataset without gradient computation.

Computes average loss across all evaluation batches with optional mixed precision. Useful for monitoring generalization and preventing overfitting.

Parameters:

Name	Type	Description	Default
model	Module	The sparse autoencoder model to evaluate.	required
eval_dataloader	DataLoader	DataLoader providing evaluation batches.	required
bf16	bool	Whether to use bfloat16 mixed precision for evaluation. Should match training setting. Defaults to True.	True

Returns:

Name	Type	Description
float	float	Average evaluation loss across all batches. Computed as total loss divided by number of batches.

Note

Prints evaluation metrics every 100 steps. Sets model to eval mode and restores training mode afterward if called during training.

Source code in deeplens\train.py

@torch.no_grad()
def evaluate(
        self,
        model: torch.nn.Module, 
        eval_dataloader: torch.utils.data.DataLoader, 
        bf16: bool = True
    ) -> float:
    """Evaluate the model on the validation dataset without gradient computation.

    Computes average loss across all evaluation batches with optional mixed precision.
    Useful for monitoring generalization and preventing overfitting.

    Args:
        model (torch.nn.Module): The sparse autoencoder model to evaluate.
        eval_dataloader (torch.utils.data.DataLoader): DataLoader providing evaluation batches.
        bf16 (bool, optional): Whether to use bfloat16 mixed precision for evaluation.
            Should match training setting. Defaults to True.

    Returns:
        float: Average evaluation loss across all batches. Computed as total loss
            divided by number of batches.

    Note:
        Prints evaluation metrics every 100 steps. Sets model to eval mode and
        restores training mode afterward if called during training.
    """
    model.eval()
    n_batches = 0
    total_loss = 0.0

    if bf16:
        device_type = "cuda" if torch.cuda.is_available() else "cpu"

    for idx, inputs in enumerate(eval_dataloader):
        if torch.cuda.is_available() and bf16:
            with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
                inputs = inputs.to(self.device)
                loss, logs = model.loss(inputs)
        else:
            inputs = inputs.to(self.device)
            loss, logs = model.loss(inputs)

        total_loss += loss.item()
        n_batches += 1

        if (idx % 100) == 0:
            current_lr = self.optim.param_groups[0]['lr']
            print(
                f"Step [{idx}/{len(eval_dataloader)}] - "
                f"eval_loss: {round(logs['mse'].item(), 3)} - "
                f"eval_nz_frac: {round(logs['non_zero_frac'].item(), 3)} - "
                f"lr: {current_lr:.2e}"
            )

    avg_loss = total_loss / n_batches
    print(f"Avg loss: {avg_loss:.3f}")
    return avg_loss

set_lr_scheduler(lr_type='cosine')

Configure and initialize learning rate scheduler with warmup.

Creates a learning rate scheduler based on the specified type. 'cosine' and 'linear' schedulers include warmup phase for training stability. 'plateau' reduces learning rate when validation loss plateaus.

Parameters:

Name	Type	Description	Default
lr_type	str	Type of learning rate scheduler. Options: - 'cosine': Linear warmup followed by cosine annealing to 10% of initial LR - 'plateau': Reduces LR by factor when validation loss stops improving - 'linear': Linear warmup followed by linear decay to 10% of initial LR Defaults to 'cosine'.	'cosine'

Returns:

Name	Type	Description
lr_scheduler	lr_scheduler	Configured PyTorch learning rate scheduler instance.

Note

For 'cosine' and 'linear', warmup steps = warmup_fraction * total_steps. 'plateau' scheduler requires manual .step(loss) calls, which are handled automatically at epoch end.

Source code in deeplens\train.py

def set_lr_scheduler(self, lr_type: str = 'cosine') -> lr_scheduler:
    """Configure and initialize learning rate scheduler with warmup.

    Creates a learning rate scheduler based on the specified type. 'cosine' and 'linear'
    schedulers include warmup phase for training stability. 'plateau' reduces learning
    rate when validation loss plateaus.

    Args:
        lr_type (str, optional): Type of learning rate scheduler. Options:
            - 'cosine': Linear warmup followed by cosine annealing to 10% of initial LR
            - 'plateau': Reduces LR by factor when validation loss stops improving
            - 'linear': Linear warmup followed by linear decay to 10% of initial LR
            Defaults to 'cosine'.

    Returns:
        lr_scheduler: Configured PyTorch learning rate scheduler instance.

    Note:
        For 'cosine' and 'linear', warmup steps = warmup_fraction * total_steps.
        'plateau' scheduler requires manual .step(loss) calls, which are handled
        automatically at epoch end.
    """
    assert lr_type in ['cosine', 'plateau', 'linear'], "Use 'cosine', 'plateau', or 'linear'"

    total_steps = self.epochs * len(self.train_dataloader)
    warmup_steps = int(total_steps * self.warmup_fraction)

    if lr_type == 'cosine':
        warmup = lr_scheduler.LinearLR(
            self.optim, start_factor=0.01, end_factor=1.0,
            total_iters=warmup_steps
        )
        cosine = lr_scheduler.CosineAnnealingLR(
            self.optim, T_max=total_steps - warmup_steps,
            eta_min=self.optim.param_groups[0]['lr'] * 0.1
        )
        self.scheduler = lr_scheduler.SequentialLR(
            self.optim, schedulers=[warmup, cosine],
            milestones=[warmup_steps]
        )
    elif lr_type == 'plateau':
        self.scheduler = lr_scheduler.ReduceLROnPlateau(
            self.optim, patience=10, threshold=1e-4, min_lr=1e-6
        )
    else:
        warmup = lr_scheduler.LinearLR(
            self.optim, start_factor=0.01, end_factor=1.0,
            total_iters=warmup_steps
        )
        decay = lr_scheduler.LinearLR(
            self.optim, start_factor=1.0, end_factor=0.1,
            total_iters=total_steps - warmup_steps
        )
        self.scheduler = lr_scheduler.SequentialLR(
            self.optim, schedulers=[warmup, decay],
            milestones=[warmup_steps]
        )
    return self.scheduler

set_seed(seed=1)

Set random seeds for reproducible training across all libraries.

Configures random number generators for NumPy, PyTorch (CPU and CUDA), and makes CUDNN operations deterministic. Essential for experiment reproducibility.

Parameters:

Name	Type	Description	Default
seed	int	Random seed value to use. Same seed should produce identical results across runs (assuming same hardware/software). Defaults to 1.	1

Note

Setting deterministic=True may reduce performance. Disables CUDNN benchmarking for reproducibility.

Source code in deeplens\train.py

def set_seed(self, seed: int = 1) -> None:
    """Set random seeds for reproducible training across all libraries.

    Configures random number generators for NumPy, PyTorch (CPU and CUDA), and
    makes CUDNN operations deterministic. Essential for experiment reproducibility.

    Args:
        seed (int, optional): Random seed value to use. Same seed should produce
            identical results across runs (assuming same hardware/software).
            Defaults to 1.

    Note:
        Setting deterministic=True may reduce performance. Disables CUDNN benchmarking
        for reproducibility.
    """
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True
    print(f"Using random seed: {seed}")

train()

Execute the complete training loop for all epochs.

Main training orchestrator that: 1. Sets random seed for reproducibility 2. Moves model to appropriate device 3. Creates checkpoint directories if saving enabled 4. Runs training epochs with periodic evaluation 5. Saves best models based on evaluation loss 6. Applies learning rate scheduling 7. Logs metrics to W&B if enabled 8. Finalizes W&B run on completion

Note

Checkpoints saved to saved_models/{model_name}/run_{timestamp}/. Each epoch includes full pass through training data followed by evaluation. Best model determined by lowest evaluation loss.

Source code in deeplens\train.py

def train(self) -> None:
    """Execute the complete training loop for all epochs.

    Main training orchestrator that:
    1. Sets random seed for reproducibility
    2. Moves model to appropriate device
    3. Creates checkpoint directories if saving enabled
    4. Runs training epochs with periodic evaluation
    5. Saves best models based on evaluation loss
    6. Applies learning rate scheduling
    7. Logs metrics to W&B if enabled
    8. Finalizes W&B run on completion

    Note:
        Checkpoints saved to saved_models/{model_name}/run_{timestamp}/.
        Each epoch includes full pass through training data followed by evaluation.
        Best model determined by lowest evaluation loss.
    """
    self.set_seed(self.random_seed)
    self.model.to(self.device)

    if self.save_checkpoints:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        os.makedirs(f"saved_models/{self.model_name}/run_{timestamp}", exist_ok=True)

    best_loss = float('inf')
    global_step = 0

    for epoch in range(self.epochs):
        print(f"\nEpoch [{epoch+1}/{self.epochs}]")
        global_step, best_loss = self.train_one_epoch(
            model=self.model, 
            train_dataloader=self.train_dataloader, 
            optim=self.optim, 
            bf16=self.bf16,
            global_step=global_step,
            timestamp=timestamp,
            best_loss=best_loss
        )

        print(f"\n{'='*60}")
        print(f"End of epoch {epoch+1} evaluation")
        print(f"{'='*60}")
        loss = self.evaluate(
            model=self.model,
            eval_dataloader=self.eval_dataloader,
            bf16=self.bf16
        )

        if self.log_wandb:
            wandb.log({
                "epoch": epoch + 1,
                "eval/epoch_loss": loss,
            }, step=global_step)

        if self.save_checkpoints and loss < best_loss:
            if self.save_best_only:
                save_path = f"saved_models/{self.model_name}/run_{timestamp}/best_model.pt"
            else:
                save_path = f"saved_models/{self.model_name}/run_{timestamp}/sae_epoch_{epoch+1}.pt"
            torch.save(self.model.state_dict(), save_path)
            print(f"New best model saved (loss: {loss:.3f})")
            best_loss = loss

        if self.scheduler is not None and isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau):
            self.scheduler.step(loss)

    if self.log_wandb:
        wandb.finish()

    print("Finished training!")

train_one_epoch(model, train_dataloader, optim, bf16=True, global_step=0, timestamp=None, best_loss=float('inf'))

Execute one complete training epoch with periodic evaluation and checkpointing.

Iterates through all training batches, performs forward/backward passes with optional mixed precision, applies gradient clipping, updates learning rate, and periodically evaluates and saves checkpoints.

Parameters:

Name	Type	Description	Default
model	Module	The sparse autoencoder model to train.	required
train_dataloader	DataLoader	DataLoader providing training batches.	required
optim	Optimizer	Optimizer for updating model parameters.	required
bf16	bool	Whether to use bfloat16 mixed precision. Defaults to True.	True
global_step	int	Current global training step count across all epochs. Used for logging and checkpoint naming. Defaults to 0.	0
timestamp	str	Timestamp string for checkpoint directory naming. Format: "YYYYMMDD_HHMMSS". Defaults to None.	None
best_loss	float	Best evaluation loss achieved so far. Used to determine when to save new checkpoints. Defaults to float('inf').	float('inf')

Returns:

Type	Description
tuple[int, float]	tuple[int, float]: A tuple containing: - global_step: Updated global step count after this epoch - best_loss: Updated best evaluation loss (may be unchanged)

Note

Prints training metrics every 100 steps. Evaluates every eval_steps and saves checkpoints when evaluation loss improves. Logs to W&B if enabled.

Source code in deeplens\train.py

def train_one_epoch(
        self, 
        model: torch.nn.Module, 
        train_dataloader: torch.utils.data.DataLoader, 
        optim: torch.optim.Optimizer, 
        bf16: bool = True,
        global_step: int = 0,
        timestamp: str = None,
        best_loss: float = float('inf')
    ) -> tuple[int, float]:
    """Execute one complete training epoch with periodic evaluation and checkpointing.

    Iterates through all training batches, performs forward/backward passes with optional
    mixed precision, applies gradient clipping, updates learning rate, and periodically
    evaluates and saves checkpoints.

    Args:
        model (torch.nn.Module): The sparse autoencoder model to train.
        train_dataloader (torch.utils.data.DataLoader): DataLoader providing training batches.
        optim (torch.optim.Optimizer): Optimizer for updating model parameters.
        bf16 (bool, optional): Whether to use bfloat16 mixed precision. Defaults to True.
        global_step (int, optional): Current global training step count across all epochs.
            Used for logging and checkpoint naming. Defaults to 0.
        timestamp (str, optional): Timestamp string for checkpoint directory naming.
            Format: "YYYYMMDD_HHMMSS". Defaults to None.
        best_loss (float, optional): Best evaluation loss achieved so far. Used to
            determine when to save new checkpoints. Defaults to float('inf').

    Returns:
        tuple[int, float]: A tuple containing:
            - global_step: Updated global step count after this epoch
            - best_loss: Updated best evaluation loss (may be unchanged)

    Note:
        Prints training metrics every 100 steps. Evaluates every eval_steps and saves
        checkpoints when evaluation loss improves. Logs to W&B if enabled.
    """
    model.train()

    if bf16:
        scaler = torch.amp.GradScaler("cuda")
        device_type = "cuda" if torch.cuda.is_available() else "cpu"

    for idx, inputs in enumerate(train_dataloader):
        optim.zero_grad()
        if torch.cuda.is_available() and bf16:
            with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
                inputs = inputs.to(self.device)
                loss, logs = model.loss(inputs)
            scaler.scale(loss).backward()

            if self.grad_clip_norm is not None:
                scaler.unscale_(optim)
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(), self.grad_clip_norm
                )

            scaler.step(optim)
            scaler.update()

        else:
            inputs = inputs.to(self.device)
            loss, logs = model.loss(inputs)

            if self.grad_clip_norm is not None:
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(), self.grad_clip_norm
                )

            loss.backward()
            optim.step()

        if self.scheduler is not None and not isinstance(self.scheduler, lr_scheduler.ReduceLROnPlateau):
            self.scheduler.step()

        model.post_step()
        global_step += 1

        if self.log_wandb:
            wandb.log({
                "train/loss": logs['mse'].item(),
                "train/non_zero_frac": logs['non_zero_frac'].item(),
                "train/lr": self.optim.param_groups[0]['lr'],
                "global_step": global_step
            }, step=global_step)

        if (idx % 100) == 0:
            current_lr = self.optim.param_groups[0]['lr']
            print(
                f"Step [{idx}/{len(train_dataloader)}] - "
                f"train_loss: {round(logs['mse'].item(), 3)} - "
                f"train_nz_frac: {round(logs['non_zero_frac'].item(), 3)} - "
                f"lr: {current_lr:.2e}"
            )

        if global_step % self.eval_steps == 0:
            print(f"\n{'='*60}")
            print(f"Intermediate Evaluation at step {global_step}")
            print(f"{'='*60}")
            eval_loss = self.evaluate(
                model=model,
                eval_dataloader=self.eval_dataloader,
                bf16=bf16
            )

            if self.log_wandb:
                wandb.log({
                    "eval/loss": eval_loss,
                    "global_step": global_step
                }, step=global_step)

            if self.save_checkpoints and eval_loss < best_loss:
                if self.save_best_only:
                    save_path = f"saved_models/{self.model_name}/run_{timestamp}/best_model.pt"
                else:
                    save_path = f"saved_models/{self.model_name}/run_{timestamp}/sae_step_{global_step}.pt"
                torch.save(model.state_dict(), save_path)
                print(f"New best model saved (loss: {eval_loss:.6f})")
                best_loss = eval_loss

            model.train()

    return global_step, best_loss

deeplens.pipeline

InterveneFeatures

Manipulate and intervene on sparse autoencoder latent features.

This class loads a trained sparse autoencoder and provides methods to analyze and modify its latent feature space, enabling causal analysis of feature effects on model behavior.

Source code in deeplens\intervene.py

class InterveneFeatures():
    """Manipulate and intervene on sparse autoencoder latent features.

    This class loads a trained sparse autoencoder and provides methods to analyze
    and modify its latent feature space, enabling causal analysis of feature effects
    on model behavior.
    """
    def __init__(
            self,
            sae_model: str = None,
            sae_config: str | dict = None,
            device: str = "auto"
        ):
        """Initialize the InterveneFeatures class for manipulating sparse autoencoder latent features.

        This class provides functionality to load a trained sparse autoencoder model and
        intervene on its latent feature space to analyze and modify activations.

        Args:
            sae_model (str, optional): Path to the trained sparse autoencoder model weights file.
                Should be a .pt or .pth file containing the model state dict. Defaults to None.
            sae_config (str | dict, optional): Configuration for the sparse autoencoder.
                Can be either a dictionary containing model hyperparameters or a path to a
                YAML configuration file. Defaults to None.
            device (str, optional): Device to run computations on. Can be "auto" for automatic
                selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
        """
        self.model_dir = sae_model

        self.device = get_device(device)
        print(f"Running on device: {self.device}")

        if isinstance(sae_config, dict):
            self.model_config = sae_config
        elif isinstance(sae_config, str) and sae_config.endswith(".yaml"):
            self.model_config = self.config_from_yaml(sae_config)
        else:
            raise ValueError("sae_config must be dict or path to .yaml file.")

        self.model = self.load_model()

    @torch.no_grad()
    def get_decoded(self, activations) -> torch.Tensor:
        """Encode input activations through the sparse autoencoder to get latent features.

        Passes the input activations through the sparse autoencoder's forward pass and
        returns the latent feature representation (z) from the encoded space.

        Args:
            activations (torch.Tensor | array-like): Input activations to encode. Can be a
                PyTorch tensor or any array-like structure that can be converted to a tensor.

        Returns:
            torch.Tensor: The latent feature representation (z) from the sparse autoencoder's
                encoded space.
        """
        if not isinstance(activations, torch.Tensor):
            activations = torch.Tensor(activations)
        activations = activations.to(self.device)
        _, z, _ = self.model(activations)
        return z

    @torch.no_grad()
    def get_alive_features(
            self, 
            activations: torch.Tensor, 
            token_position: int = -1, 
            k: int | None = None,
            return_values: bool = False
        ) -> torch.Tensor:
        """Get indices of non-zero (active) features in the latent space for a specific token.

        Encodes the input activations through the sparse autoencoder and identifies which
        latent features are active (non-zero) at the specified token position.

        Args:
            activations (torch.Tensor | array-like): Input activations to encode. Can be a
                PyTorch tensor or any array-like structure that can be converted to a tensor.
            token_position (int, optional): Position of the token in the sequence to analyze.
                Use -1 for the last token. Defaults to -1.
            k (int, optional): If provided, returns only the top-k most active features
                instead of all non-zero features. Defaults to None.
            return_values (bool, optional): If True, returns both indices and values.
                Defaults to False.

        Returns:
            torch.Tensor | tuple[torch.Tensor, torch.Tensor]: If return_values is False,
                returns a 1D tensor containing the indices of non-zero features. If True,
                returns a tuple of (indices, values).
        """
        z = self.get_decoded(activations)
        z_token = z[token_position]
        if k is not None:
            topk_result = torch.topk(z_token, k=k)
            feature_idxs = topk_result.indices
            feature_vals = topk_result.values
        else:
            feature_idxs = torch.nonzero(z_token != 0, as_tuple=False).squeeze(-1)
            feature_vals = z_token[feature_idxs]
        if return_values:
            return feature_idxs, feature_vals
        return feature_idxs

    @torch.no_grad()
    def intervene_feature(
            self, 
            activations, 
            feature: int, 
            alpha: float = 2.0,
            token_positions: int | list[int] | None = None
        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Intervene on a specific latent feature by scaling its activation.

        Encodes the input activations, multiplies the specified feature by alpha at the
        given token positions, and returns both the original and modified decoded outputs
        for comparison.

        Args:
            activations (torch.Tensor | array-like): Input activations to encode and modify.
                Can be a PyTorch tensor or any array-like structure.
            feature (int): Index of the latent feature to intervene on.
            alpha (float, optional): Scaling factor to apply to the feature. Values > 1
                amplify the feature, values < 1 suppress it. Defaults to 2.0.
            token_positions (int | list[int] | None, optional): Token position(s) at which
                to apply the intervention. If None, applies to all tokens. If int, applies
                to a single position. If list, applies to multiple positions. Defaults to None.

        Returns:
            tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
                - activations: The original input activations
                - original_decoded: Decoded output without intervention
                - modified_decoded: Decoded output with the feature intervention applied
        """
        if not isinstance(activations, torch.Tensor):
            activations = torch.Tensor(activations).unsqueeze(0)

        activations = activations.to(self.device)
        _, z, _ = self.model(activations)
        modified = z.clone()

        if token_positions is None:
            modified[:, feature] *= alpha
        elif isinstance(token_positions, int):
            modified[token_positions, feature] *= alpha
        else:
            for pos in token_positions:
                modified[pos, feature] *= alpha

        original_decoded = self.model.decode(z)
        modified_decoded = self.model.decode(modified)
        return activations, original_decoded, modified_decoded

    def load_model(self) -> torch.nn.Module:
        """Load the sparse autoencoder model from disk.

        Loads the model weights from the specified path and initializes a
        SparseAutoencoder instance with the provided configuration.

        Returns:
            torch.nn.Module: The loaded sparse autoencoder model moved to the
                appropriate device.
        """
        weights = torch.load(self.model_dir, map_location=self.device)
        model = SparseAutoencoder(**self.model_config)
        model.load_state_dict(state_dict=weights)
        return model.to(self.device)

    def config_from_yaml(self, file) -> dict:
        """Load sparse autoencoder configuration from a YAML file.

        Reads and parses a YAML configuration file containing the hyperparameters
        for the sparse autoencoder model.

        Args:
            file (str): Path to the YAML configuration file.

        Returns:
            dict: Dictionary containing the model configuration parameters.
        """
        try:
            with open(file, "r") as f:
                config = yaml.safe_load(f)
            return config
        except FileNotFoundError:
            raise FileNotFoundError(f"Config file not found: {file}")
        except yaml.YAMLError as e:
            raise ValueError(f"Invalid YAML in {file}: {e}")

__init__(sae_model=None, sae_config=None, device='auto')

Initialize the InterveneFeatures class for manipulating sparse autoencoder latent features.

This class provides functionality to load a trained sparse autoencoder model and intervene on its latent feature space to analyze and modify activations.

Parameters:

Name	Type	Description	Default
sae_model	str	Path to the trained sparse autoencoder model weights file. Should be a .pt or .pth file containing the model state dict. Defaults to None.	None
sae_config	str | dict	Configuration for the sparse autoencoder. Can be either a dictionary containing model hyperparameters or a path to a YAML configuration file. Defaults to None.	None
device	str	Device to run computations on. Can be "auto" for automatic selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".	'auto'

Source code in deeplens\intervene.py

def __init__(
        self,
        sae_model: str = None,
        sae_config: str | dict = None,
        device: str = "auto"
    ):
    """Initialize the InterveneFeatures class for manipulating sparse autoencoder latent features.

    This class provides functionality to load a trained sparse autoencoder model and
    intervene on its latent feature space to analyze and modify activations.

    Args:
        sae_model (str, optional): Path to the trained sparse autoencoder model weights file.
            Should be a .pt or .pth file containing the model state dict. Defaults to None.
        sae_config (str | dict, optional): Configuration for the sparse autoencoder.
            Can be either a dictionary containing model hyperparameters or a path to a
            YAML configuration file. Defaults to None.
        device (str, optional): Device to run computations on. Can be "auto" for automatic
            selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
    """
    self.model_dir = sae_model

    self.device = get_device(device)
    print(f"Running on device: {self.device}")

    if isinstance(sae_config, dict):
        self.model_config = sae_config
    elif isinstance(sae_config, str) and sae_config.endswith(".yaml"):
        self.model_config = self.config_from_yaml(sae_config)
    else:
        raise ValueError("sae_config must be dict or path to .yaml file.")

    self.model = self.load_model()

config_from_yaml(file)

Load sparse autoencoder configuration from a YAML file.

Reads and parses a YAML configuration file containing the hyperparameters for the sparse autoencoder model.

Parameters:

Name	Type	Description	Default
file	str	Path to the YAML configuration file.	required

Returns:

Name	Type	Description
dict	dict	Dictionary containing the model configuration parameters.

Source code in deeplens\intervene.py

def config_from_yaml(self, file) -> dict:
    """Load sparse autoencoder configuration from a YAML file.

    Reads and parses a YAML configuration file containing the hyperparameters
    for the sparse autoencoder model.

    Args:
        file (str): Path to the YAML configuration file.

    Returns:
        dict: Dictionary containing the model configuration parameters.
    """
    try:
        with open(file, "r") as f:
            config = yaml.safe_load(f)
        return config
    except FileNotFoundError:
        raise FileNotFoundError(f"Config file not found: {file}")
    except yaml.YAMLError as e:
        raise ValueError(f"Invalid YAML in {file}: {e}")

get_alive_features(activations, token_position=-1, k=None, return_values=False)

Get indices of non-zero (active) features in the latent space for a specific token.

Encodes the input activations through the sparse autoencoder and identifies which latent features are active (non-zero) at the specified token position.

Parameters:

Name	Type	Description	Default
activations	Tensor | array - like	Input activations to encode. Can be a PyTorch tensor or any array-like structure that can be converted to a tensor.	required
token_position	int	Position of the token in the sequence to analyze. Use -1 for the last token. Defaults to -1.	-1
k	int	If provided, returns only the top-k most active features instead of all non-zero features. Defaults to None.	None
return_values	bool	If True, returns both indices and values. Defaults to False.	False

Returns:

Type	Description
Tensor	torch.Tensor | tuple[torch.Tensor, torch.Tensor]: If return_values is False, returns a 1D tensor containing the indices of non-zero features. If True, returns a tuple of (indices, values).

Source code in deeplens\intervene.py

@torch.no_grad()
def get_alive_features(
        self, 
        activations: torch.Tensor, 
        token_position: int = -1, 
        k: int | None = None,
        return_values: bool = False
    ) -> torch.Tensor:
    """Get indices of non-zero (active) features in the latent space for a specific token.

    Encodes the input activations through the sparse autoencoder and identifies which
    latent features are active (non-zero) at the specified token position.

    Args:
        activations (torch.Tensor | array-like): Input activations to encode. Can be a
            PyTorch tensor or any array-like structure that can be converted to a tensor.
        token_position (int, optional): Position of the token in the sequence to analyze.
            Use -1 for the last token. Defaults to -1.
        k (int, optional): If provided, returns only the top-k most active features
            instead of all non-zero features. Defaults to None.
        return_values (bool, optional): If True, returns both indices and values.
            Defaults to False.

    Returns:
        torch.Tensor | tuple[torch.Tensor, torch.Tensor]: If return_values is False,
            returns a 1D tensor containing the indices of non-zero features. If True,
            returns a tuple of (indices, values).
    """
    z = self.get_decoded(activations)
    z_token = z[token_position]
    if k is not None:
        topk_result = torch.topk(z_token, k=k)
        feature_idxs = topk_result.indices
        feature_vals = topk_result.values
    else:
        feature_idxs = torch.nonzero(z_token != 0, as_tuple=False).squeeze(-1)
        feature_vals = z_token[feature_idxs]
    if return_values:
        return feature_idxs, feature_vals
    return feature_idxs

get_decoded(activations)

Encode input activations through the sparse autoencoder to get latent features.

Passes the input activations through the sparse autoencoder's forward pass and returns the latent feature representation (z) from the encoded space.

Parameters:

Name	Type	Description	Default
activations	Tensor | array - like	Input activations to encode. Can be a PyTorch tensor or any array-like structure that can be converted to a tensor.	required

Returns:

Type	Description
Tensor	torch.Tensor: The latent feature representation (z) from the sparse autoencoder's encoded space.

Source code in deeplens\intervene.py

@torch.no_grad()
def get_decoded(self, activations) -> torch.Tensor:
    """Encode input activations through the sparse autoencoder to get latent features.

    Passes the input activations through the sparse autoencoder's forward pass and
    returns the latent feature representation (z) from the encoded space.

    Args:
        activations (torch.Tensor | array-like): Input activations to encode. Can be a
            PyTorch tensor or any array-like structure that can be converted to a tensor.

    Returns:
        torch.Tensor: The latent feature representation (z) from the sparse autoencoder's
            encoded space.
    """
    if not isinstance(activations, torch.Tensor):
        activations = torch.Tensor(activations)
    activations = activations.to(self.device)
    _, z, _ = self.model(activations)
    return z

intervene_feature(activations, feature, alpha=2.0, token_positions=None)

Intervene on a specific latent feature by scaling its activation.

Encodes the input activations, multiplies the specified feature by alpha at the given token positions, and returns both the original and modified decoded outputs for comparison.

Parameters:

Name	Type	Description	Default
activations	Tensor | array - like	Input activations to encode and modify. Can be a PyTorch tensor or any array-like structure.	required
feature	int	Index of the latent feature to intervene on.	required
alpha	float	Scaling factor to apply to the feature. Values > 1 amplify the feature, values < 1 suppress it. Defaults to 2.0.	2.0
token_positions	int | list[int] | None	Token position(s) at which to apply the intervention. If None, applies to all tokens. If int, applies to a single position. If list, applies to multiple positions. Defaults to None.	None

Returns:

Type	Description
tuple[Tensor, Tensor, Tensor]	tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing: - activations: The original input activations - original_decoded: Decoded output without intervention - modified_decoded: Decoded output with the feature intervention applied

Source code in deeplens\intervene.py

@torch.no_grad()
def intervene_feature(
        self, 
        activations, 
        feature: int, 
        alpha: float = 2.0,
        token_positions: int | list[int] | None = None
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Intervene on a specific latent feature by scaling its activation.

    Encodes the input activations, multiplies the specified feature by alpha at the
    given token positions, and returns both the original and modified decoded outputs
    for comparison.

    Args:
        activations (torch.Tensor | array-like): Input activations to encode and modify.
            Can be a PyTorch tensor or any array-like structure.
        feature (int): Index of the latent feature to intervene on.
        alpha (float, optional): Scaling factor to apply to the feature. Values > 1
            amplify the feature, values < 1 suppress it. Defaults to 2.0.
        token_positions (int | list[int] | None, optional): Token position(s) at which
            to apply the intervention. If None, applies to all tokens. If int, applies
            to a single position. If list, applies to multiple positions. Defaults to None.

    Returns:
        tuple[torch.Tensor, torch.Tensor, torch.Tensor]: A tuple containing:
            - activations: The original input activations
            - original_decoded: Decoded output without intervention
            - modified_decoded: Decoded output with the feature intervention applied
    """
    if not isinstance(activations, torch.Tensor):
        activations = torch.Tensor(activations).unsqueeze(0)

    activations = activations.to(self.device)
    _, z, _ = self.model(activations)
    modified = z.clone()

    if token_positions is None:
        modified[:, feature] *= alpha
    elif isinstance(token_positions, int):
        modified[token_positions, feature] *= alpha
    else:
        for pos in token_positions:
            modified[pos, feature] *= alpha

    original_decoded = self.model.decode(z)
    modified_decoded = self.model.decode(modified)
    return activations, original_decoded, modified_decoded

load_model()

Load the sparse autoencoder model from disk.

Loads the model weights from the specified path and initializes a SparseAutoencoder instance with the provided configuration.

Returns:

Type	Description
Module	torch.nn.Module: The loaded sparse autoencoder model moved to the appropriate device.

Source code in deeplens\intervene.py

def load_model(self) -> torch.nn.Module:
    """Load the sparse autoencoder model from disk.

    Loads the model weights from the specified path and initializes a
    SparseAutoencoder instance with the provided configuration.

    Returns:
        torch.nn.Module: The loaded sparse autoencoder model moved to the
            appropriate device.
    """
    weights = torch.load(self.model_dir, map_location=self.device)
    model = SparseAutoencoder(**self.model_config)
    model.load_state_dict(state_dict=weights)
    return model.to(self.device)

ReinjectSingleSample

Reinject modified activations into a language model for causal inference.

This class enables injecting modified activations back into a transformer model's forward pass to observe the causal effects on model predictions and generated text. Useful for validating feature interventions and conducting mechanistic interpretability experiments.

Source code in deeplens\intervene.py

class ReinjectSingleSample():
    """Reinject modified activations into a language model for causal inference.

    This class enables injecting modified activations back into a transformer model's
    forward pass to observe the causal effects on model predictions and generated text.
    Useful for validating feature interventions and conducting mechanistic interpretability
    experiments.
    """
    def __init__(
            self, 
            hf_model: str, 
            device: str = "auto", 
            cache_dir: str = 'cache'
        ):
        """Initialize the ReinjectSingleSample class for causal inference with modified activations.

        Loads a HuggingFace causal language model and tokenizer to enable reinjection of
        modified activations into the model's forward pass for text generation and analysis.

        Args:
            hf_model (str): Name or path of the HuggingFace model to load.
                Should be a valid model identifier from the HuggingFace model hub
                (e.g., "gpt2", "meta-llama/Llama-2-7b").
            device (str, optional): Device to run computations on. Can be "auto" for automatic
                selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
            cache_dir (str, optional): Directory to cache downloaded models.
                Defaults to 'cache'.
        """
        self.device = get_device(device)
        print(f"Running on device: {self.device}")

        os.makedirs(cache_dir, exist_ok=True)
        self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
        self.model.eval()

    @torch.no_grad()
    def reinject_and_generate(
            self, 
            text, 
            modified_activations, 
            layer: int = 3, 
            generate: bool = False, 
            max_new_tokens: int = 25, 
            temperature: float = 1.0
        ) -> torch.Tensor | str:
        """Reinject modified activations into a model layer and optionally generate text.

        Replaces the activations at the specified layer with the provided modified activations
        during the forward pass. Can either return logits for the input text or generate
        new tokens autoregressively.

        Args:
            text (str): Input text to tokenize and process through the model.
            modified_activations (torch.Tensor): The modified activations to inject at the
                specified layer. Should have the appropriate shape for the layer's output.
            layer (int, optional): Index of the transformer layer where activations should
                be replaced. Defaults to 3.
            generate (bool, optional): If True, generates new tokens autoregressively.
                If False, only returns logits for the input. Defaults to False.
            max_new_tokens (int, optional): Maximum number of new tokens to generate when
                generate=True. Defaults to 25.
            temperature (float, optional): Sampling temperature for generation. Higher values
                (>1.0) make output more random, lower values (<1.0) more deterministic.
                Set to 0 for greedy decoding. Defaults to 1.0.

        Returns:
            torch.Tensor | str: If generate=False, returns the model's logits as a tensor.
                If generate=True, returns the generated text as a string.

        Note:
            The hook is automatically removed after execution to prevent interference
            with subsequent forward passes. For generation mode, the hook only affects
            the first forward pass to avoid applying the intervention to newly generated tokens.
        """
        modified_activations = modified_activations.to(self.device)
        call_count = [0]
        def replacement_hook(module, input, output):
            if generate and call_count[0] > 0:
                return output
            call_count[0] += 1
            return modified_activations

        mlp_module = self.get_module_for_replacement_hook(layer_idx=layer)
        hook = mlp_module.register_forward_hook(replacement_hook)
        tokens = self.tokenizer(text, return_tensors='pt').to(self.device)
        try:
            if generate:
                generated_ids = self.model.generate(
                    **tokens,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature,
                    do_sample=temperature > 0
                )
                return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
            else:
                out = self.model(**tokens)
                return out.logits
        finally:
            hook.remove()

    def get_module_for_replacement_hook(self, layer_idx) -> torch.nn.Module:
        """Get the MLP activation module for a specific layer.

        Retrieves the MLP activation function module at the specified layer,
        which can be used to register forward hooks for activation replacement.

        Args:
            layer_idx (int): Index of the transformer layer (0-indexed).

        Returns:
            torch.nn.Module: The MLP activation module at the specified layer.
        """
        if isinstance(self.model, (
            transformers.GPT2LMHeadModel, 
            transformers.FalconForCausalLM
        )):
            module = self.model.transformer.h[layer_idx].mlp.act
        elif isinstance(self.model, (
            transformers.LlamaForCausalLM, 
            transformers.MistralForCausalLM, 
            transformers.Gemma3ForCausalLM, 
            transformers.GemmaForCausalLM, 
            transformers.Qwen2ForCausalLM,
            transformers.Qwen3ForCausalLM
        )):
            module = self.model.model.layers[layer_idx].mlp.act_fn
        elif isinstance(self.model, (
            transformers.PhiForCausalLM, 
            transformers.Phi3ForCausalLM
        )):
            module = self.model.model.layers[layer_idx].mlp.activation_fn
        else:
            raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

        return module

__init__(hf_model, device='auto', cache_dir='cache')

Initialize the ReinjectSingleSample class for causal inference with modified activations.

Loads a HuggingFace causal language model and tokenizer to enable reinjection of modified activations into the model's forward pass for text generation and analysis.

Parameters:

Name	Type	Description	Default
hf_model	str	Name or path of the HuggingFace model to load. Should be a valid model identifier from the HuggingFace model hub (e.g., "gpt2", "meta-llama/Llama-2-7b").	required
device	str	Device to run computations on. Can be "auto" for automatic selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".	'auto'
cache_dir	str	Directory to cache downloaded models. Defaults to 'cache'.	'cache'

Source code in deeplens\intervene.py

def __init__(
        self, 
        hf_model: str, 
        device: str = "auto", 
        cache_dir: str = 'cache'
    ):
    """Initialize the ReinjectSingleSample class for causal inference with modified activations.

    Loads a HuggingFace causal language model and tokenizer to enable reinjection of
    modified activations into the model's forward pass for text generation and analysis.

    Args:
        hf_model (str): Name or path of the HuggingFace model to load.
            Should be a valid model identifier from the HuggingFace model hub
            (e.g., "gpt2", "meta-llama/Llama-2-7b").
        device (str, optional): Device to run computations on. Can be "auto" for automatic
            selection, "cuda" for GPU, or "cpu" for CPU. Defaults to "auto".
        cache_dir (str, optional): Directory to cache downloaded models.
            Defaults to 'cache'.
    """
    self.device = get_device(device)
    print(f"Running on device: {self.device}")

    os.makedirs(cache_dir, exist_ok=True)
    self.model = AutoModelForCausalLM.from_pretrained(hf_model, cache_dir=cache_dir).to(self.device)
    self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
    self.model.eval()

get_module_for_replacement_hook(layer_idx)

Get the MLP activation module for a specific layer.

Retrieves the MLP activation function module at the specified layer, which can be used to register forward hooks for activation replacement.

Parameters:

Name	Type	Description	Default
layer_idx	int	Index of the transformer layer (0-indexed).	required

Returns:

Type	Description
Module	torch.nn.Module: The MLP activation module at the specified layer.

Source code in deeplens\intervene.py

def get_module_for_replacement_hook(self, layer_idx) -> torch.nn.Module:
    """Get the MLP activation module for a specific layer.

    Retrieves the MLP activation function module at the specified layer,
    which can be used to register forward hooks for activation replacement.

    Args:
        layer_idx (int): Index of the transformer layer (0-indexed).

    Returns:
        torch.nn.Module: The MLP activation module at the specified layer.
    """
    if isinstance(self.model, (
        transformers.GPT2LMHeadModel, 
        transformers.FalconForCausalLM
    )):
        module = self.model.transformer.h[layer_idx].mlp.act
    elif isinstance(self.model, (
        transformers.LlamaForCausalLM, 
        transformers.MistralForCausalLM, 
        transformers.Gemma3ForCausalLM, 
        transformers.GemmaForCausalLM, 
        transformers.Qwen2ForCausalLM,
        transformers.Qwen3ForCausalLM
    )):
        module = self.model.model.layers[layer_idx].mlp.act_fn
    elif isinstance(self.model, (
        transformers.PhiForCausalLM, 
        transformers.Phi3ForCausalLM
    )):
        module = self.model.model.layers[layer_idx].mlp.activation_fn
    else:
        raise NotImplementedError(f"Model type {type(self.model).__name__} is not currently supported.")

    return module

reinject_and_generate(text, modified_activations, layer=3, generate=False, max_new_tokens=25, temperature=1.0)

Reinject modified activations into a model layer and optionally generate text.

Replaces the activations at the specified layer with the provided modified activations during the forward pass. Can either return logits for the input text or generate new tokens autoregressively.

Parameters:

Name	Type	Description	Default
text	str	Input text to tokenize and process through the model.	required
modified_activations	Tensor	The modified activations to inject at the specified layer. Should have the appropriate shape for the layer's output.	required
layer	int	Index of the transformer layer where activations should be replaced. Defaults to 3.	3
generate	bool	If True, generates new tokens autoregressively. If False, only returns logits for the input. Defaults to False.	False
max_new_tokens	int	Maximum number of new tokens to generate when generate=True. Defaults to 25.	25
temperature	float	Sampling temperature for generation. Higher values (>1.0) make output more random, lower values (<1.0) more deterministic. Set to 0 for greedy decoding. Defaults to 1.0.	1.0

Returns:

Type	Description
Tensor | str	torch.Tensor | str: If generate=False, returns the model's logits as a tensor. If generate=True, returns the generated text as a string.

Note

The hook is automatically removed after execution to prevent interference with subsequent forward passes. For generation mode, the hook only affects the first forward pass to avoid applying the intervention to newly generated tokens.

Source code in deeplens\intervene.py

@torch.no_grad()
def reinject_and_generate(
        self, 
        text, 
        modified_activations, 
        layer: int = 3, 
        generate: bool = False, 
        max_new_tokens: int = 25, 
        temperature: float = 1.0
    ) -> torch.Tensor | str:
    """Reinject modified activations into a model layer and optionally generate text.

    Replaces the activations at the specified layer with the provided modified activations
    during the forward pass. Can either return logits for the input text or generate
    new tokens autoregressively.

    Args:
        text (str): Input text to tokenize and process through the model.
        modified_activations (torch.Tensor): The modified activations to inject at the
            specified layer. Should have the appropriate shape for the layer's output.
        layer (int, optional): Index of the transformer layer where activations should
            be replaced. Defaults to 3.
        generate (bool, optional): If True, generates new tokens autoregressively.
            If False, only returns logits for the input. Defaults to False.
        max_new_tokens (int, optional): Maximum number of new tokens to generate when
            generate=True. Defaults to 25.
        temperature (float, optional): Sampling temperature for generation. Higher values
            (>1.0) make output more random, lower values (<1.0) more deterministic.
            Set to 0 for greedy decoding. Defaults to 1.0.

    Returns:
        torch.Tensor | str: If generate=False, returns the model's logits as a tensor.
            If generate=True, returns the generated text as a string.

    Note:
        The hook is automatically removed after execution to prevent interference
        with subsequent forward passes. For generation mode, the hook only affects
        the first forward pass to avoid applying the intervention to newly generated tokens.
    """
    modified_activations = modified_activations.to(self.device)
    call_count = [0]
    def replacement_hook(module, input, output):
        if generate and call_count[0] > 0:
            return output
        call_count[0] += 1
        return modified_activations

    mlp_module = self.get_module_for_replacement_hook(layer_idx=layer)
    hook = mlp_module.register_forward_hook(replacement_hook)
    tokens = self.tokenizer(text, return_tensors='pt').to(self.device)
    try:
        if generate:
            generated_ids = self.model.generate(
                **tokens,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                do_sample=temperature > 0
            )
            return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        else:
            out = self.model(**tokens)
            return out.logits
    finally:
        hook.remove()

pipeline(text, sae_model, sae_config, layer=3, hf_model='gpt2', feature=-1, alpha=5.0, tok_position=-1, generate=False, max_new_tokens=25, temperature=1.0, verbose=False)

End-to-end pipeline for extracting, intervening on, and analyzing SAE features.

This function provides a complete workflow for mechanistic interpretability analysis: extracts MLP activations from a language model, decodes them through a sparse autoencoder, intervenes on specific features, and reinjections the modified activations back into the model to observe behavioral changes.

Parameters:

Name	Type	Description	Default
text	str	Input text to process. Will be tokenized and processed by the language model.	required
sae_model	str	Path to the trained sparse autoencoder model weights file (.pt or .pth).	required
sae_config	str	Path to the SAE configuration YAML file or dictionary with hyperparameters.	required
layer	int	Transformer layer to extract activations from. Supports negative indexing (e.g., -1 for last layer). Defaults to 3.	3
hf_model	str	Name or path of the HuggingFace transformer model to use (e.g., "gpt2-medium"). Defaults to "gpt2".	'gpt2'
feature	int	Index of the SAE feature to intervene on from the alive features list. Supports negative indexing. Defaults to -1 (last alive feature).	-1
alpha	float	Intervention strength multiplier. Larger values create stronger feature manipulations. Can be negative to suppress features. Defaults to 5.0.	5.0
tok_position	int	Token position to analyze and intervene on. Supports negative indexing (e.g., -1 for last token). Defaults to -1.	-1
generate	bool	Whether to generate continuation text after reinjection. If False, returns only the reconstructed input. Defaults to False.	False
max_new_tokens	int	Maximum number of new tokens to generate if generate=True. Must be positive. Defaults to 25.	25
temperature	float	Sampling temperature for text generation. Higher values increase randomness. Must be positive. Defaults to 1.0.	1.0
verbose	bool	Whether to print diagnostic information during execution. Defaults to False.	False

Returns:

Type	Description
tuple[Tensor, Tensor, Tensor]	A 3-tuple containing: - original: Model output with original unmodified activations - decoded: Model output with SAE-reconstructed activations (no intervention) - modified: Model output with intervened feature activations

Source code in deeplens\pipeline.py

def pipeline(
        text: str,
        sae_model: str,
        sae_config: str,
        layer: int = 3,
        hf_model: str = "gpt2",
        feature: int = -1,
        alpha: float = 5.0,
        tok_position: int = -1,
        generate: bool = False,
        max_new_tokens: int = 25,
        temperature: float = 1.0,
        verbose: bool = False
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """End-to-end pipeline for extracting, intervening on, and analyzing SAE features.

    This function provides a complete workflow for mechanistic interpretability analysis:
    extracts MLP activations from a language model, decodes them through a sparse autoencoder,
    intervenes on specific features, and reinjections the modified activations back into the
    model to observe behavioral changes.

    Args:
        text (str): Input text to process. Will be tokenized and processed by the language model.
        sae_model (str): Path to the trained sparse autoencoder model weights file (.pt or .pth).
        sae_config (str): Path to the SAE configuration YAML file or dictionary with hyperparameters.
        layer (int, optional): Transformer layer to extract activations from. Supports negative
            indexing (e.g., -1 for last layer). Defaults to 3.
        hf_model (str, optional): Name or path of the HuggingFace transformer model to use
            (e.g., "gpt2-medium"). Defaults to "gpt2".
        feature (int, optional): Index of the SAE feature to intervene on from the alive features
            list. Supports negative indexing. Defaults to -1 (last alive feature).
        alpha (float, optional): Intervention strength multiplier. Larger values create stronger
            feature manipulations. Can be negative to suppress features. Defaults to 5.0.
        tok_position (int, optional): Token position to analyze and intervene on. Supports
            negative indexing (e.g., -1 for last token). Defaults to -1.
        generate (bool, optional): Whether to generate continuation text after reinjection.
            If False, returns only the reconstructed input. Defaults to False.
        max_new_tokens (int, optional): Maximum number of new tokens to generate if generate=True.
            Must be positive. Defaults to 25.
        temperature (float, optional): Sampling temperature for text generation. Higher values
            increase randomness. Must be positive. Defaults to 1.0.
        verbose (bool, optional): Whether to print diagnostic information during execution.
            Defaults to False.

    Returns:
        A 3-tuple containing:
            - original: Model output with original unmodified activations
            - decoded: Model output with SAE-reconstructed activations (no intervention)
            - modified: Model output with intervened feature activations
    """
    if not text or not isinstance(text, str):
        raise ValueError("Text must be a non-empty string")

    if not isinstance(layer, int):
        raise ValueError(f"Layer must be an integer, got {layer}")

    if not isinstance(feature, int):
        raise ValueError(f"Feature must be an integer, got {feature}")

    if not isinstance(alpha, (int, float)):
        raise ValueError(f"Alpha must be numeric, got {alpha}")

    if not isinstance(max_new_tokens, int) or max_new_tokens <= 0:
        raise ValueError(f"max_new_tokens must be positive integer, got {max_new_tokens}")


    try:
        extractor = ExtractSingleSample(hf_model=hf_model, layer=layer)
        intervene = InterveneFeatures(sae_model=sae_model, sae_config=sae_config)
        reinject = ReinjectSingleSample(hf_model=hf_model)

        acts = extractor.get_mlp_acts(text)
        if acts is None:
            raise RuntimeError("Failed to extract activations. Returned 'None'.")

        alive_features = intervene.get_alive_features(acts, token_position=tok_position)
        if alive_features is None or len(alive_features) == 0:
            raise RuntimeError(f"No alive features found at position {tok_position}")

        if feature < -len(alive_features) or feature >= len(alive_features):
            raise ValueError(
                f"Feature index {feature} out of bounds for {len(alive_features)} features"
            )

        if verbose:
            print(f"{len(alive_features)} alive features discovered at position {tok_position}.")
            print(f"Modifying feature {alive_features[feature].item()}")

        original, decoded, modified = intervene.intervene_feature(
            activations=acts,
            feature=alive_features[feature].item(),
            alpha=alpha, 
            token_positions=tok_position
        )

        if original is None or decoded is None or modified is None:
            raise RuntimeError("Feature intervention returned 'None'")

        feature_versions = [original, decoded, modified]
        outputs = []
        for i, feature_acts in enumerate(feature_versions):
            try:
                out = reinject.reinject_and_generate(
                    text=text,
                    modified_activations=feature_acts,
                    layer=layer,
                    generate=generate,
                    max_new_tokens=max_new_tokens,
                    temperature=temperature
                )
                outputs.append(out)
            except Exception as e:
                raise RuntimeError(f"Failed to reinject feature version {i}: {e}")      

        return tuple(outputs)

    except ValueError:
        raise
    except Exception as e:
        raise RuntimeError(f"Pipeline execution failed: {e}")