nix-apollo commented 9 months ago

Centred rib

Description

Incorporates a centring matrix Y (aka Gamma) into the calculation of the rib rotation
Isolates the bias direction when center = true, returning it in the 0th position.
Reorganizes the order of the residual stream so the last position is always the bias position.

Related Issue

Closes #248

Motivation and Context

This is a first version of centered rib. It seems likely that we might want to handle lambda and/or edge calculation differently in the future to make the calculation more principled. For instance, the baseline for IG is still the 0 point, instead of the mean activation.

How Has This Been Tested?

I have a test that checks invariants for the output of the centered rib build. Including that there is a single constant direction pointing in the direction we expect, and that activations in all other rib directions are centered.

This code was also used for various analysis in the OP report, where it seemed to do reasonable things.

Does this PR introduce a breaking change?

Residual stream reorder may break some analysis code. No interface changes.

nix-apollo commented 8 months ago

Todos:

[x] Change so that the bias rib direction is always sorted first (fix computation code, update tests, remove logic)
[x] Look over code. Fix anything ugly
[x] Somewhat better tests. In particular test more layers in modular addition
[x] Rename bias pos to const pos (?)

nix-apollo commented 7 months ago

We only ever need a bias position for both:

making a shift matrix
identifying the constant rib dir (to move it to the 0th in ordering) both of these only need a single constant pos. A good simplification would be to have get_dataset_means return a single int for the bias pos instead of a tensor. Even better this could be a separate helper function that only needs the means as input.

Possibly I'll postpone this and not improve the current situation in this PR.