sae-lens 6.15.0__py3-none-any.whl → 6.24.1__py3-none-any.whl

sae_lens/saes/temporal_sae.py

@@ -0,0 +1,365 @@
+"""TemporalSAE: A Sparse Autoencoder with temporal attention mechanism.
+
+TemporalSAE decomposes activations into:
+1. Predicted codes (from attention over context)
+2. Novel codes (sparse features of the residual)
+
+See: https://arxiv.org/abs/2410.04185
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from typing_extensions import override
+
+from sae_lens import logger
+from sae_lens.saes.sae import SAE, SAEConfig
+
+
+def get_attention(query: torch.Tensor, key: torch.Tensor) -> torch.Tensor:
+    """Compute causal attention weights."""
+    L, S = query.size(-2), key.size(-2)
+    scale_factor = 1 / math.sqrt(query.size(-1))
+    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+    temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0)
+    attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+    attn_bias.to(query.dtype)
+
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    return torch.softmax(attn_weight, dim=-1)
+
+
+class ManualAttention(nn.Module):
+    """Manual attention implementation for TemporalSAE."""
+
+    def __init__(
+        self,
+        dimin: int,
+        n_heads: int = 4,
+        bottleneck_factor: int = 64,
+        bias_k: bool = True,
+        bias_q: bool = True,
+        bias_v: bool = True,
+        bias_o: bool = True,
+    ):
+        super().__init__()
+        assert dimin % (bottleneck_factor * n_heads) == 0
+
+        self.n_heads = n_heads
+        self.n_embds = dimin // bottleneck_factor
+        self.dimin = dimin
+
+        # Key, query, value projections
+        self.k_ctx = nn.Linear(dimin, self.n_embds, bias=bias_k)
+        self.q_target = nn.Linear(dimin, self.n_embds, bias=bias_q)
+        self.v_ctx = nn.Linear(dimin, dimin, bias=bias_v)
+        self.c_proj = nn.Linear(dimin, dimin, bias=bias_o)
+
+        # Normalize to match scale with representations
+        with torch.no_grad():
+            scaling = 1 / math.sqrt(self.n_embds // self.n_heads)
+            self.k_ctx.weight.copy_(
+                scaling
+                * self.k_ctx.weight
+                / (1e-6 + torch.linalg.norm(self.k_ctx.weight, dim=1, keepdim=True))
+            )
+            self.q_target.weight.copy_(
+                scaling
+                * self.q_target.weight
+                / (1e-6 + torch.linalg.norm(self.q_target.weight, dim=1, keepdim=True))
+            )
+
+            scaling = 1 / math.sqrt(self.dimin // self.n_heads)
+            self.v_ctx.weight.copy_(
+                scaling
+                * self.v_ctx.weight
+                / (1e-6 + torch.linalg.norm(self.v_ctx.weight, dim=1, keepdim=True))
+            )
+
+            scaling = 1 / math.sqrt(self.dimin)
+            self.c_proj.weight.copy_(
+                scaling
+                * self.c_proj.weight
+                / (1e-6 + torch.linalg.norm(self.c_proj.weight, dim=1, keepdim=True))
+            )
+
+    def forward(
+        self, x_ctx: torch.Tensor, x_target: torch.Tensor, get_attn_map: bool = False
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Compute projective attention output."""
+        k = self.k_ctx(x_ctx)
+        v = self.v_ctx(x_ctx)
+        q = self.q_target(x_target)
+
+        # Split into heads
+        B, T, _ = x_ctx.size()
+        k = k.view(B, T, self.n_heads, self.n_embds // self.n_heads).transpose(1, 2)
+        q = q.view(B, T, self.n_heads, self.n_embds // self.n_heads).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.dimin // self.n_heads).transpose(1, 2)
+
+        # Attention map (optional)
+        attn_map = None
+        if get_attn_map:
+            attn_map = get_attention(query=q, key=k)
+
+        # Scaled dot-product attention
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=None, dropout_p=0, is_causal=True
+        )
+
+        # Reshape and project
+        d_target = self.c_proj(
+            attn_output.transpose(1, 2).contiguous().view(B, T, self.dimin)
+        )
+
+        return d_target, attn_map
+
+
+@dataclass
+class TemporalSAEConfig(SAEConfig):
+    """Configuration for TemporalSAE inference.
+
+    Args:
+        d_in: Input dimension (dimensionality of the activations being encoded)
+        d_sae: SAE latent dimension (number of features)
+        n_heads: Number of attention heads in temporal attention
+        n_attn_layers: Number of attention layers
+        bottleneck_factor: Bottleneck factor for attention dimension
+        sae_diff_type: Type of SAE for novel codes ('relu' or 'topk')
+        kval_topk: K value for top-k sparsity (if sae_diff_type='topk')
+        tied_weights: Whether to tie encoder and decoder weights
+        activation_normalization_factor: Scalar factor for rescaling activations (used with normalize_activations='constant_scalar_rescale')
+    """
+
+    n_heads: int = 8
+    n_attn_layers: int = 1
+    bottleneck_factor: int = 64
+    sae_diff_type: Literal["relu", "topk"] = "topk"
+    kval_topk: int | None = None
+    tied_weights: bool = True
+    activation_normalization_factor: float = 1.0
+
+    def __post_init__(self):
+        # Call parent's __post_init__ first, but allow constant_scalar_rescale
+        if self.normalize_activations not in [
+            "none",
+            "expected_average_only_in",
+            "constant_norm_rescale",
+            "constant_scalar_rescale",  # Temporal SAEs support this
+            "layer_norm",
+        ]:
+            raise ValueError(
+                f"normalize_activations must be none, expected_average_only_in, layer_norm, constant_norm_rescale, or constant_scalar_rescale. Got {self.normalize_activations}"
+            )
+
+    @override
+    @classmethod
+    def architecture(cls) -> str:
+        return "temporal"
+
+
+class TemporalSAE(SAE[TemporalSAEConfig]):
+    """TemporalSAE: Sparse Autoencoder with temporal attention.
+
+    This SAE decomposes each activation x_t into:
+    - x_pred: Information aggregated from context {x_0, ..., x_{t-1}}
+    - x_novel: Novel information at position t (encoded sparsely)
+
+    The forward pass:
+    1. Uses attention layers to predict x_t from context
+    2. Encodes the residual (novel part) with a sparse SAE
+    3. Combines both for reconstruction
+    """
+
+    # Custom parameters (in addition to W_enc, W_dec, b_dec from base)
+    attn_layers: nn.ModuleList  # Attention layers
+    eps: float
+    lam: float
+
+    def __init__(self, cfg: TemporalSAEConfig, use_error_term: bool = False):
+        # Call parent init first
+        super().__init__(cfg, use_error_term)
+
+        # Initialize attention layers after parent init and move to correct device
+        self.attn_layers = nn.ModuleList(
+            [
+                ManualAttention(
+                    dimin=cfg.d_sae,
+                    n_heads=cfg.n_heads,
+                    bottleneck_factor=cfg.bottleneck_factor,
+                    bias_k=True,
+                    bias_q=True,
+                    bias_v=True,
+                    bias_o=True,
+                ).to(device=self.device, dtype=self.dtype)
+                for _ in range(cfg.n_attn_layers)
+            ]
+        )
+
+        self.eps = 1e-6
+        self.lam = 1 / (4 * self.cfg.d_in)
+
+    @override
+    def _setup_activation_normalization(self):
+        """Set up activation normalization functions for TemporalSAE.
+
+        Overrides the base implementation to handle constant_scalar_rescale
+        using the temporal-specific activation_normalization_factor.
+        """
+        if self.cfg.normalize_activations == "constant_scalar_rescale":
+            # Handle constant scalar rescaling for temporal SAEs
+            def run_time_activation_norm_fn_in(x: torch.Tensor) -> torch.Tensor:
+                return x * self.cfg.activation_normalization_factor
+
+            def run_time_activation_norm_fn_out(x: torch.Tensor) -> torch.Tensor:
+                return x / self.cfg.activation_normalization_factor
+
+            self.run_time_activation_norm_fn_in = run_time_activation_norm_fn_in
+            self.run_time_activation_norm_fn_out = run_time_activation_norm_fn_out
+        else:
+            # Delegate to parent for all other normalization types
+            super()._setup_activation_normalization()
+
+    @override
+    def initialize_weights(self) -> None:
+        """Initialize TemporalSAE weights."""
+        # Initialize D (decoder) and b (bias)
+        self.W_dec = nn.Parameter(
+            torch.randn(
+                (self.cfg.d_sae, self.cfg.d_in), dtype=self.dtype, device=self.device
+            )
+        )
+        self.b_dec = nn.Parameter(
+            torch.zeros((self.cfg.d_in), dtype=self.dtype, device=self.device)
+        )
+
+        # Initialize E (encoder) if not tied
+        if not self.cfg.tied_weights:
+            self.W_enc = nn.Parameter(
+                torch.randn(
+                    (self.cfg.d_in, self.cfg.d_sae),
+                    dtype=self.dtype,
+                    device=self.device,
+                )
+            )
+
+    def encode_with_predictions(
+        self, x: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Encode input to novel codes only.
+
+        Returns only the sparse novel codes (not predicted codes).
+        This is the main feature representation for TemporalSAE.
+        """
+        # Process input through SAELens preprocessing
+        x = self.process_sae_in(x)
+
+        B, L, _ = x.shape
+
+        if self.cfg.tied_weights:  # noqa: SIM108
+            W_enc = self.W_dec.T
+        else:
+            W_enc = self.W_enc
+
+        # Compute predicted codes using attention
+        x_residual = x
+        z_pred = torch.zeros((B, L, self.cfg.d_sae), device=x.device, dtype=x.dtype)
+
+        for attn_layer in self.attn_layers:
+            # Encode input to latent space
+            z_input = F.relu(torch.matmul(x_residual * self.lam, W_enc))
+
+            # Shift context (causal masking)
+            z_ctx = torch.cat(
+                (torch.zeros_like(z_input[:, :1, :]), z_input[:, :-1, :].clone()), dim=1
+            )
+
+            # Apply attention to get predicted codes
+            z_pred_, _ = attn_layer(z_ctx, z_input, get_attn_map=False)
+            z_pred_ = F.relu(z_pred_)
+
+            # Project predicted codes back to input space
+            Dz_pred_ = torch.matmul(z_pred_, self.W_dec)
+            Dz_norm_ = Dz_pred_.norm(dim=-1, keepdim=True) + self.eps
+
+            # Compute projection scale
+            proj_scale = (Dz_pred_ * x_residual).sum(
+                dim=-1, keepdim=True
+            ) / Dz_norm_.pow(2)
+
+            # Accumulate predicted codes
+            z_pred = z_pred + (z_pred_ * proj_scale)
+
+            # Remove prediction from residual
+            x_residual = x_residual - proj_scale * Dz_pred_
+
+        # Encode residual (novel part) with sparse SAE
+        z_novel = F.relu(torch.matmul(x_residual * self.lam, W_enc))
+        if self.cfg.sae_diff_type == "topk":
+            kval = self.cfg.kval_topk
+            if kval is not None:
+                _, topk_indices = torch.topk(z_novel, kval, dim=-1)
+                mask = torch.zeros_like(z_novel)
+                mask.scatter_(-1, topk_indices, 1)
+                z_novel = z_novel * mask
+
+        # Return only novel codes (these are the interpretable features)
+        return z_novel, z_pred
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        return self.encode_with_predictions(x)[0]
+
+    def decode(self, feature_acts: torch.Tensor) -> torch.Tensor:
+        """Decode novel codes to reconstruction.
+
+        Note: This only decodes the novel codes. For full reconstruction,
+        use forward() which includes predicted codes.
+        """
+        # Decode novel codes
+        sae_out = torch.matmul(feature_acts, self.W_dec)
+        sae_out = sae_out + self.b_dec
+
+        # Apply hook
+        sae_out = self.hook_sae_recons(sae_out)
+
+        # Apply output activation normalization (reverses input normalization)
+        sae_out = self.run_time_activation_norm_fn_out(sae_out)
+
+        # Add bias (already removed in process_sae_in)
+        logger.warning(
+            "NOTE this only decodes x_novel. The x_pred is missing, so we're not reconstructing the full x."
+        )
+        return sae_out
+
+    @override
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Full forward pass through TemporalSAE.
+
+        Returns complete reconstruction (predicted + novel).
+        """
+        # Encode
+        z_novel, z_pred = self.encode_with_predictions(x)
+
+        # Decode the sum of predicted and novel codes.
+        x_recons = torch.matmul(z_novel + z_pred, self.W_dec) + self.b_dec
+
+        # Apply output activation normalization (reverses input normalization)
+        x_recons = self.run_time_activation_norm_fn_out(x_recons)
+
+        return self.hook_sae_output(x_recons)
+
+    @override
+    def fold_W_dec_norm(self) -> None:
+        raise NotImplementedError("Folding W_dec_norm is not supported for TemporalSAE")
+
+    @override
+    @torch.no_grad()
+    def fold_activation_norm_scaling_factor(self, scaling_factor: float) -> None:
+        raise NotImplementedError(
+            "Folding activation norm scaling factor is not supported for TemporalSAE"
+        )

sae_lens/saes/topk_sae.py

@@ -4,7 +4,6 @@ from dataclasses import dataclass
 from typing import Any, Callable
 
 import torch
-from jaxtyping import Float
 from torch import nn
 from transformer_lens.hook_points import HookPoint
 from typing_extensions import override
@@ -235,9 +234,7 @@ class TopKSAE(SAE[TopKSAEConfig]):
         super().initialize_weights()
         _init_weights_topk(self)
 
-    def encode(
-        self, x: Float[torch.Tensor, "... d_in"]
-    ) -> Float[torch.Tensor, "... d_sae"]:
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
         """
         Converts input x into feature activations.
         Uses topk activation under the hood.
@@ -251,8 +248,8 @@ class TopKSAE(SAE[TopKSAEConfig]):
 
     def decode(
         self,
-        feature_acts: Float[torch.Tensor, "... d_sae"],
-    ) -> Float[torch.Tensor, "... d_in"]:
+        feature_acts: torch.Tensor,
+    ) -> torch.Tensor:
         """
         Reconstructs the input from topk feature activations.
         Applies optional finetuning scaling, hooking to recons, out normalization,
@@ -354,8 +351,8 @@ class TopKTrainingSAE(TrainingSAE[TopKTrainingSAEConfig]):
         _init_weights_topk(self)
 
     def encode_with_hidden_pre(
-        self, x: Float[torch.Tensor, "... d_in"]
-    ) -> tuple[Float[torch.Tensor, "... d_sae"], Float[torch.Tensor, "... d_sae"]]:
+        self, x: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Similar to the base training method: calculate pre-activations, then apply TopK.
         """
@@ -372,8 +369,8 @@ class TopKTrainingSAE(TrainingSAE[TopKTrainingSAEConfig]):
     @override
     def decode(
         self,
-        feature_acts: Float[torch.Tensor, "... d_sae"],
-    ) -> Float[torch.Tensor, "... d_in"]:
+        feature_acts: torch.Tensor,
+    ) -> torch.Tensor:
         """
         Decodes feature activations back into input space,
         applying optional finetuning scale, hooking, out normalization, etc.
@@ -534,7 +531,7 @@ def _fold_norm_topk(
     b_enc: torch.Tensor,
     W_dec: torch.Tensor,
 ) -> None:
-    W_dec_norm = W_dec.norm(dim=-1)
+    W_dec_norm = W_dec.norm(dim=-1).clamp(min=1e-8)
     b_enc.data = b_enc.data * W_dec_norm
     W_dec_norms = W_dec_norm.unsqueeze(1)
     W_dec.data = W_dec.data / W_dec_norms

sae_lens/saes/transcoder.py

@@ -368,3 +368,44 @@ class JumpReLUTranscoder(Transcoder):
     def from_dict(cls, config_dict: dict[str, Any]) -> "JumpReLUTranscoder":
         cfg = JumpReLUTranscoderConfig.from_dict(config_dict)
         return cls(cfg)
+
+
+@dataclass
+class JumpReLUSkipTranscoderConfig(JumpReLUTranscoderConfig):
+    """Configuration for JumpReLU transcoder."""
+
+    @classmethod
+    def architecture(cls) -> str:
+        """Return the architecture name for this config."""
+        return "jumprelu_skip_transcoder"
+
+    @classmethod
+    def from_dict(cls, config_dict: dict[str, Any]) -> "JumpReLUSkipTranscoderConfig":
+        """Create a JumpReLUSkipTranscoderConfig from a dictionary."""
+        # Filter to only include valid dataclass fields
+        filtered_config_dict = filter_valid_dataclass_fields(config_dict, cls)
+
+        # Create the config instance
+        res = cls(**filtered_config_dict)
+
+        # Handle metadata if present
+        if "metadata" in config_dict:
+            res.metadata = SAEMetadata(**config_dict["metadata"])
+
+        return res
+
+
+class JumpReLUSkipTranscoder(JumpReLUTranscoder, SkipTranscoder):
+    """
+    A transcoder with a learnable skip connection and JumpReLU activation function.
+    """
+
+    cfg: JumpReLUSkipTranscoderConfig  # type: ignore[assignment]
+
+    def __init__(self, cfg: JumpReLUSkipTranscoderConfig):
+        super().__init__(cfg)
+
+    @classmethod
+    def from_dict(cls, config_dict: dict[str, Any]) -> "JumpReLUSkipTranscoder":
+        cfg = JumpReLUSkipTranscoderConfig.from_dict(config_dict)
+        return cls(cfg)

sae_lens/training/activation_scaler.py

@@ -1,5 +1,6 @@
 import json
 from dataclasses import dataclass
+from pathlib import Path
 from statistics import mean
 
 import torch
@@ -51,3 +52,9 @@ class ActivationScaler:
 
         with open(file_path, "w") as f:
             json.dump({"scaling_factor": self.scaling_factor}, f)
+
+    def load(self, file_path: str | Path):
+        """load the state dict from a file in json format"""
+        with open(file_path) as f:
+            data = json.load(f)
+            self.scaling_factor = data["scaling_factor"]

sae_lens/training/activations_store.py

@@ -4,6 +4,7 @@ import json
 import os
 import warnings
 from collections.abc import Generator, Iterator, Sequence
+from pathlib import Path
 from typing import Any, Literal, cast
 
 import datasets
@@ -11,10 +12,9 @@ import torch
 from datasets import Dataset, DatasetDict, IterableDataset, load_dataset
 from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import HfHubHTTPError
-from jaxtyping import Float, Int
 from requests import HTTPError
-from safetensors.torch import save_file
-from tqdm import tqdm
+from safetensors.torch import load_file, save_file
+from tqdm.auto import tqdm
 from transformer_lens.hook_points import HookedRootModule
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
@@ -24,7 +24,7 @@ from sae_lens.config import (
     HfDataset,
     LanguageModelSAERunnerConfig,
 )
-from sae_lens.constants import DTYPE_MAP
+from sae_lens.constants import ACTIVATIONS_STORE_STATE_FILENAME, DTYPE_MAP
 from sae_lens.pretokenize_runner import get_special_token_from_cfg
 from sae_lens.saes.sae import SAE, T_SAE_CONFIG, T_TRAINING_SAE_CONFIG
 from sae_lens.tokenization_and_batching import concat_and_batch_sequences
@@ -166,9 +166,11 @@ class ActivationsStore:
         disable_concat_sequences: bool = False,
         sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = "bos",
     ) -> ActivationsStore:
+        if context_size is None:
+            context_size = sae.cfg.metadata.context_size
         if sae.cfg.metadata.hook_name is None:
             raise ValueError("hook_name is required")
-        if sae.cfg.metadata.context_size is None:
+        if context_size is None:
             raise ValueError("context_size is required")
         if sae.cfg.metadata.prepend_bos is None:
             raise ValueError("prepend_bos is required")
@@ -178,9 +180,7 @@ class ActivationsStore:
             d_in=sae.cfg.d_in,
             hook_name=sae.cfg.metadata.hook_name,
             hook_head_index=sae.cfg.metadata.hook_head_index,
-            context_size=sae.cfg.metadata.context_size
-            if context_size is None
-            else context_size,
+            context_size=context_size,
             prepend_bos=sae.cfg.metadata.prepend_bos,
             streaming=streaming,
             store_batch_size_prompts=store_batch_size_prompts,
@@ -230,7 +230,7 @@ class ActivationsStore:
             load_dataset(
                 dataset,
                 split="train",
-                streaming=streaming,
+                streaming=streaming,  # type: ignore
                 trust_remote_code=dataset_trust_remote_code,  # type: ignore
             )
             if isinstance(dataset, str)
@@ -318,7 +318,7 @@ class ActivationsStore:
                 )
         else:
             warnings.warn(
-                "Dataset is not tokenized. Pre-tokenizing will improve performance and allows for more control over special tokens. See https://jbloomaus.github.io/SAELens/training_saes/#pretokenizing-datasets for more info."
+                "Dataset is not tokenized. Pre-tokenizing will improve performance and allows for more control over special tokens. See https://decoderesearch.github.io/SAELens/training_saes/#pretokenizing-datasets for more info."
             )
 
         self.iterable_sequences = self._iterate_tokenized_sequences()
@@ -541,8 +541,8 @@ class ActivationsStore:
         d_in: int,
         raise_on_epoch_end: bool,
     ) -> tuple[
-        Float[torch.Tensor, "(total_size context_size) num_layers d_in"],
-        Int[torch.Tensor, "(total_size context_size)"] | None,
+        torch.Tensor,
+        torch.Tensor | None,
     ]:
         """
         Loads `total_size` activations from `cached_activation_dataset`
@@ -729,6 +729,48 @@ class ActivationsStore:
         """save the state dict to a file in safetensors format"""
         save_file(self.state_dict(), file_path)
 
+    def save_to_checkpoint(self, checkpoint_path: str | Path):
+        """Save the state dict to a checkpoint path"""
+        self.save(str(Path(checkpoint_path) / ACTIVATIONS_STORE_STATE_FILENAME))
+
+    def load_from_checkpoint(self, checkpoint_path: str | Path):
+        """Load the state dict from a checkpoint path"""
+        self.load(str(Path(checkpoint_path) / ACTIVATIONS_STORE_STATE_FILENAME))
+
+    def load(self, file_path: str):
+        """Load the state dict from a file in safetensors format"""
+
+        state_dict = load_file(file_path)
+
+        if "n_dataset_processed" in state_dict:
+            target_n_dataset_processed = state_dict["n_dataset_processed"].item()
+
+            # Only fast-forward if needed
+
+            if target_n_dataset_processed > self.n_dataset_processed:
+                logger.info(
+                    "Fast-forwarding through dataset samples to match checkpoint position"
+                )
+                samples_to_skip = target_n_dataset_processed - self.n_dataset_processed
+
+                pbar = tqdm(
+                    total=samples_to_skip,
+                    desc="Fast-forwarding through dataset",
+                    leave=False,
+                )
+                while target_n_dataset_processed > self.n_dataset_processed:
+                    start = self.n_dataset_processed
+                    try:
+                        # Just consume and ignore the values to fast-forward
+                        next(self.iterable_sequences)
+                    except StopIteration:
+                        logger.warning(
+                            "Dataset exhausted during fast-forward. Resetting dataset."
+                        )
+                        self.iterable_sequences = self._iterate_tokenized_sequences()
+                    pbar.update(self.n_dataset_processed - start)
+                pbar.close()
+
 
 def validate_pretokenized_dataset_tokenizer(
     dataset_path: str, model_tokenizer: PreTrainedTokenizerBase

sae_lens/training/optim.py

@@ -2,6 +2,8 @@
 Took the LR scheduler from my previous work: https://github.com/jbloomAus/DecisionTransformerInterpretability/blob/ee55df35cdb92e81d689c72fb9dd5a7252893363/src/decision_transformer/utils.py#L425
 """
 
+from typing import Any
+
 import torch.optim as optim
 import torch.optim.lr_scheduler as lr_scheduler
 
@@ -150,3 +152,12 @@ class CoefficientScheduler:
     def value(self) -> float:
         """Returns the current scalar value."""
         return self.current_value
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "current_step": self.current_step,
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]):
+        for k in state_dict:
+            setattr(self, k, state_dict[k])