Support custom loss functions

[norabelrose]

We all want to change the CCS loss function in various ways, so we need a flexible way of defining and specifying loss functions.

We need some sort of API, maybe a class that can be inherited from, for defining entirely new loss functions programmatically and passing them into the CCS class. We also need a small library of predefined loss functions that can be accessed by name from the command line.

The custom losses need to be able to specify what inputs they take— for example, a conjunction/disjunction consistency loss would need hidden states from N independent propositions. Prompt invariance losses will take M different variants of the same proposition. We'll need some sort of data collation logic to piece together the prompts required by the given loss and then extract the hidden states from the model.

This is a big task which should probably be split into multiple PRs.

[FabienRoger]

Additionally, you would want to train on multiple datasets at once? (either with different losses, or at least with different numbers of variants)

[lauritowal]

Additionally, you would want to train on multiple datasets at once? (either with different losses, or at least with different numbers of variants)

You mean train one probe on multiple datasets?

[FabienRoger]

Exactly. It can help

• having a purely consistency based probe (because when the number of classes vary, an easy way to be consistent is to be truthlike)
• exploring whether you can find a direction which is truthlike according to many datasets/losses simultaneously.

[FabienRoger]

But if it's too complicated, maybe start by having sth that works in simpler cases?