adds fields for sequence logic in ActivationsStore #492

anthonyduong9 · 2025-06-10T06:58:03Z

Description

Adds a field, where if the value is True, excludes the BOS token between concatenated sequences, and a field, where if the value is True, disables concatenating sequences and ignores sequences shorter than the context size. Gemma 2 wasn't trained on BOS tokens and we want to match this during training. We can do this using PretokenizeRunner, but not without it.

Fixes #472

Type of change

Please delete options that are not relevant.

New feature (non-breaking change which adds functionality)

Checklist:

I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing tests pass locally with my changes
I have not rewritten tests relating to key interfaces which would affect backward compatibility

You have tested formatting, typing and tests

I have run make check-ci to check format and linting. (you can run make format to format code if needed.)

Performance Check.

If you have implemented a training change, please indicate precisely how performance changes with respect to the following metrics:

L0
CE Loss
MSE Loss
Feature Dashboard Interpretability

Please links to wandb dashboards with a control and test group.

codecov · 2025-06-10T07:03:42Z

Codecov Report

Attention: Patch coverage is 72.00000% with 7 lines in your changes missing coverage. Please review.

Project coverage is 85.55%. Comparing base (ce310f5) to head (f0b2a2a).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
sae_lens/config.py	50.00%	5 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #492      +/-   ##
==========================================
- Coverage   85.65%   85.55%   -0.10%     
==========================================
  Files          28       28              
  Lines        3568     3593      +25     
  Branches      443      448       +5     
==========================================
+ Hits         3056     3074      +18     
- Misses        331      336       +5     
- Partials      181      183       +2

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

chanind · 2025-06-16T15:04:49Z

sae_lens/training/activations_store.py

+                if self.prepend_bos and not self.exclude_bos_between_sequences:
+                    sequence_separator_token_id = bos_token_id
+
+                yield from concat_and_batch_sequences(


It would be better to adapt the concat_and_batch_sequences function to handle disable_concat_sequences, so then the pretokenize runner could also gain this capability too.

Good idea, I've added disable_concat_sequences for concat_and_batch_sequences() and PretokenizeRunnerConfig.

chanind · 2025-06-16T15:08:34Z

sae_lens/config.py

    store_batch_size_prompts: int = 32
    seqpos_slice: tuple[int | None, ...] = (None,)
+    disable_concat_sequences: bool = False
+    exclude_bos_between_sequences: bool = False


Does it make sense to have the same options both here and in pretokenize_runner? I could see an argument that the pretokenize_runner arguments are too confusing / complicated if it's best to stick with just bos here, but can also see it being nice if we have the same options and thus same capabilities regardless of pretokenizing or not

Are you asking whether we should add exclude_bos_between_sequences to PretokenizeRunnerConfig, or change exclude_bos_between_sequences to be like PretokenizeRunnerConfig.sequence_separator_token?

I think we should do the latter. If that sounds good, I can make the changes.

Yeah I was thinking the latter, it is sort of strange these have different options, but can see the argument that maybe the pretokenize runner version is too complex if we want to keep them separate. It looks like the 6.0 release will have a lot of breaking changes so if we're going to unify these now is probably a good time to do it.

chanind · 2025-06-17T11:17:23Z

tests/training/test_tokenization_and_batching.py

+        [3, 4, 5, 6, 7],
+        [10, 11, 12, 13, 14],
+    ]
+    assert batches.tolist() == expected


chanind

LGTM!

anthonyduong9 · 2025-06-29T00:51:26Z

@chanind I realized I didn't change the logic for ActivationsStore.from_sae() and ActivationsStore.from_cache_activations().

I'm thinking I should add disable_concat_sequences and sequence_separator_token to SAEMetadata, and then change ActivationsStore.from_sae() to read those fields and pass the values to cls(), so that the logic for concatenating and ignoring sequences, and the sequence separator token, when loading a given SAE, are the same as they were while the SAE was trained.

And then for ActivationsStore.from_cache_activations(), I'm thinking we don't need to change anything.

Does that sound good to you?

chanind · 2025-06-30T15:44:51Z

Sounds reasonable!

anthonyduong9 · 2025-07-18T07:56:22Z

Because of jbloomAus/SAEDashboard#46 (comment), I've actually made changes so ActivationsStore.from_sae() has params disable_concat_sequences and sequence_separator_token. For SAEs not trained with SAELens, callers need to be able to pass values without reading from SAEMetadata. Because of this, and that it's tricky for the function to distinguish between None passed by the caller and no value passed at all, it's tricky to make the SAEMetadata values the default. For SAEs trained with SAELens, the caller will have to pass the values that were passed for training (which are on SAEMetadata), if they're to be the same for loading.

chanind · 2025-07-18T12:14:53Z

You should be able to check if a field is explicitly set on SAEMetadata with the in operator, e.g.: https://github.com/jbloomAus/SAELens/blob/main/tests/saes/test_sae.py#L72-L79. But regardless, letting the user pass in options seems fine. IIRC ActivationsStore.from_sae() is only ever used if the user is running evals on a random SAE from CLI anyway.

anthonyduong9 · 2025-07-18T23:16:57Z

In the case of sequence_separator_token, what we have now (a default value of "bos" for backwards compatibility, no token when None is passed, and otherwise, whatever value the caller passes) is easiest to reason about, IMO, so I'll just merge this as-is.

…rch#492) * adds fields for sequence logic in ActivationsStore * adds docstrings to LanguageModelSAERunnerConfig * groups params for ActivationsStore together * adds disable_concat_sequences to concat_and_batch_sequences() * adds disable_concat_sequences to PretokenizedDatasetMetadata * replaces exclude_bos_between_sequences with sequence_separator_token * updates test name * fixes tests * adds params to ActivationsStore.from_sae() * adds fields to SAEMetdata

anthonyduong9 marked this pull request as ready for review June 10, 2025 21:51

anthonyduong9 requested a review from chanind June 10, 2025 21:51

chanind reviewed Jun 16, 2025

View reviewed changes

anthonyduong9 requested a review from chanind June 17, 2025 06:40

chanind reviewed Jun 17, 2025

View reviewed changes

anthonyduong9 requested a review from chanind June 17, 2025 22:27

chanind approved these changes Jun 23, 2025

View reviewed changes

anthonyduong9 mentioned this pull request Jun 25, 2025

Option to zero out BOS tokens jbloomAus/SAEDashboard#46

Open

chanind force-pushed the alpha branch from f147df9 to 31332c7 Compare July 14, 2025 19:36

Base automatically changed from alpha to main July 14, 2025 22:25

anthonyduong9 added 10 commits July 17, 2025 14:49

adds fields for sequence logic in ActivationsStore

ef5081a

adds docstrings to LanguageModelSAERunnerConfig

cb8c262

groups params for ActivationsStore together

c67d499

adds disable_concat_sequences to concat_and_batch_sequences()

c74f881

adds disable_concat_sequences to PretokenizedDatasetMetadata

87f496d

replaces exclude_bos_between_sequences with sequence_separator_token

54ce4f2

updates test name

a914972

fixes tests

2b05652

adds params to ActivationsStore.from_sae()

4649418

adds fields to SAEMetdata

f0b2a2a

anthonyduong9 force-pushed the add-fields-for-sequence-logic-in-ActivationsStore branch from e1270c1 to f0b2a2a Compare July 18, 2025 06:30

anthonyduong9 merged commit 25fc5c3 into main Jul 18, 2025
3 of 5 checks passed

anthonyduong9 deleted the add-fields-for-sequence-logic-in-ActivationsStore branch July 18, 2025 23:18

adds fields for sequence logic in ActivationsStore #492

adds fields for sequence logic in ActivationsStore #492

Uh oh!

Conversation

anthonyduong9 commented Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Checklist:

You have tested formatting, typing and tests

Performance Check.

Uh oh!

codecov bot commented Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

chanind Jun 16, 2025

Choose a reason for hiding this comment

Uh oh!

anthonyduong9 Jun 17, 2025

Choose a reason for hiding this comment

Uh oh!

chanind Jun 16, 2025

Choose a reason for hiding this comment

Uh oh!

anthonyduong9 Jun 17, 2025

Choose a reason for hiding this comment

Uh oh!

chanind Jun 17, 2025

Choose a reason for hiding this comment

Uh oh!

chanind Jun 17, 2025

Choose a reason for hiding this comment

Uh oh!

chanind left a comment

Choose a reason for hiding this comment

Uh oh!

anthonyduong9 commented Jun 29, 2025

Uh oh!

chanind commented Jun 30, 2025

Uh oh!

anthonyduong9 commented Jul 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

chanind commented Jul 18, 2025

Uh oh!

anthonyduong9 commented Jul 18, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

anthonyduong9 commented Jun 10, 2025 •

edited

Loading

codecov bot commented Jun 10, 2025 •

edited

Loading

anthonyduong9 commented Jul 18, 2025 •

edited

Loading