feat: Add linear probe trainer

[tom-pollak]

Description

@chanind @jbloomAus

Created linear probe trainer which can take a hf dataset (created like in #321) and train a linear probe (provided we have labels for the activations)

This can make a pretty wandb run with all your stats by default. similar to training the SAE.

(Training run is clearly very bad)

So far just tried to keep the code simple, don't know if there's any interest in a probe trainer like this.

---

Type of change

Please delete options that are not relevant.

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

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 unit 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 unit tests (acceptance tests not currently in use)

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

[chanind]

I would personally view this as out-of-scope of this library, since there's nothing SAE-specific about training a linear probe. Curious to hear others' thoughts though!

[hijohnnylin]

I would personally view this as out-of-scope of this library, since there's nothing SAE-specific about training a linear probe. Curious to hear others' thoughts though!

FWIW, I'm a +1 to including this in SAELens, and/or some repo similarly easily accessible. possibly the scope of SAELens should be widened. I would prefer not to have a separate library for TranscoderLens, CrosscoderLens, ProbeLens etc. Though I defer to you, @curt-tigges, etc.

[curt-tigges]

We've had some discussion internally about creating and maintaining a probe-specific library, which I think would be a better fit than SAELens for this. I think one of the main benefits of having a library is having common code etc. but probing is rather independent and there isn't really anything shared with the rest of the code. There's not really a specific reason to have it in SAELens.

Proposal: I'm inclined to just go ahead and create a sketch for an initial probing library and ask @tom-pollak to resubmit this code to that. @jbloomAus any thoughts?