danbraunai-apollo commented 10 months ago

Update tlens and misc config settings

Description

Bump transformer-lens to ~=1.10.0, and subsequent torch and torchvision bumps resulting from the tlens update. Notably, this update sets IGNORE to -inf in Attention instead of -1e5 (which was causing us issues in the old version).
Reduce truncation_threshold to 1e-15 from 1e-5 in all the configs. Tests in #215 showed that this allowed us to have a much smaller atol.
Used float64 instead of float32 in pythia and modular_arithmetic configs. For modadd, this gives some slightly different edge weights (in particular, float32 truncates more vectors than float64), but the resulting graphs look almost identical (see images below). I think for modadd it probably doesn't matter much, and for pythia it might. Regardless, I'd prefer if we just continue to use float64 everywhere.
Train a new modular arithmetic model with the above updates. Store it, along with its resulting rib graph, in sample_checkpoints and sample_graphs, respectively. I did not bother running the (slow) pythia graphs with these updates.

Related Issue

Closes #217; Closes #216; Closes #205 (@nix-apollo please confirm, I think this issue may have been caused by the IGNORE that was too large?).

Motivation and Context

Having the old INGORE=1e-5 caused issues in both transformerlens and also in RIB tests (requiring much larger atols in tests). The other changes are also likely to improve any precision issues we have.

How Has This Been Tested?

Updated tests to point to the new modular arithmetic graphs. Same tests pass.

Does this PR introduce a breaking change?

Technically no. But everybody please reinstall the package that has an updated requirements.txt, otherwise things will break.

float32 vs float64 modular arithmetic graph build

modular_arithmetic_rib_graph modular_arithmetic_fp32_rib_graph

nix-apollo commented 10 months ago

Imported updated requirements and ran tests successfully.

nix-apollo commented 10 months ago

I confirm this should close #205

nix-apollo commented 10 months ago

I assume you are planning to put the change to the default value of eps in a different PR?

danbraunai-apollo commented 10 months ago

I assume you are planning to put the change to the default value of eps in a different PR?

Yep will do. Just made an issue for it https://github.com/ApolloResearch/rib/issues/224.

ApolloResearch / rib