enable setting adam pars in config

jbloomAus · Apr 16, 2024 · 1e53ede · 1e53ede
1 parent c558849
commit 1e53ede
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/sae_lens/training/config.py b/sae_lens/training/config.py
@@ -54,6 +54,8 @@ class LanguageModelSAERunnerConfig:
     prepend_bos: bool = True
 
     # Training Parameters
+    adam_beta1: float | list[float] = 0
+    adam_beta2: float | list[float] = 0.999
     mse_loss_normalization: Optional[str] = None
     l1_coefficient: float | list[float] = 1e-3
     lp_norm: float | list[float] = 1

diff --git a/sae_lens/training/train_sae_on_language_model.py b/sae_lens/training/train_sae_on_language_model.py
@@ -248,7 +248,14 @@ def _build_train_context(
     )
     n_frac_active_tokens = 0
 
-    optimizer = Adam(sae.parameters(), lr=sae.cfg.lr)
+    optimizer = Adam(
+        sae.parameters(),
+        lr=sae.cfg.lr,
+        betas=(
+            sae.cfg.adam_beta1,  # type: ignore
+            sae.cfg.adam_beta2,  # type: ignore
+        ),
+    )
     assert sae.cfg.lr_end is not None  # this is set in config post-init
     scheduler = get_scheduler(
         sae.cfg.lr_scheduler_name,