Support different bases and edges formulas

[stefan-apollo]

Re-implementing hacky changes from feature/new_norm_november_A and feature/new_norm_november_B See doc for details about the corresponding math, and the overleaf should soon also have info.

Replaces #188.

Changes:

• Implement new basis and attribution method
  • Allow switching bases with ig_formula (config) ["(1-alpha)^2", "(1-0)*alpha"]
  • New basis is default now
  • Allow switching attribution method with edge_formula (config) ["functional", "squared"]
  • Remove edge_form function
  • Replace first half of it (converting RIB to neurons, calling module, converting back) with module_hat (utils)
  • Use module_hat_partial like edge_norm_partial
  • But the actual norm computation now just happens in integrated_gradient_trapezoidal_jacobian
  • This comes with a bit of refactoring of the associated code
  • E.g. no longer needed to pass outputs_const around in arguments

Tests:

• Made build tests for mod add and MNIST run for new and old settings
• Float precision test TestPythiaFloatingPointErrors::test_interaction_rotations but I don't think this is a big deal
• Added build_graph tests
  • Check that new basis edges are >= 0
  • Skip check of sum of incoming edges for new basis or new edges as they no longer hold. See Slack discussion for details.
• Test runtime is rather large now, should reduce number of tests before merging

  42.82s call     tests/test_build_graph.py::test_mnist_build_graph_new_basis_old_attribution
  42.65s call     tests/test_build_graph.py::test_mnist_build_graph_old_basis_old_attribution
  42.58s call     tests/test_build_graph.py::test_mnist_build_graph_new_basis_and_new_attribution
  26.63s setup    tests/test_float_precision.py::TestPythiaFloatingPointErrors::test_gram_matrices
  24.84s setup    tests/test_float_precision.py::TestPythiaFloatingPointErrors::test_ablation_result_float_precision
  16.86s call     tests/test_build_graph.py::test_pythia_14m_build_graph
  15.67s call     tests/test_build_graph.py::test_modular_arithmetic_build_graph_new_basis_and_new_attribution
  11.79s call     tests/test_build_graph.py::test_modular_arithmetic_build_graph_old_basis_and_old_attribution
  11.59s call     tests/test_build_graph.py::test_modular_arithmetic_build_graph_new_basis_and_old_attribution
  9.65s call     tests/test_ablations.py::test_run_mnist_ablations[orthogonal]

  Minor:

• Fix launch.json for MNIST
• Make integral_boundary_relative_epsilon a config argument

[stefan-apollo]

So far implemented alternative integrated_gradient_trapezoidal_norm formula. Only changes attn slightly, as expected by Lucius.

Next up: Changes in edges

[stefan-apollo]

I see to have introduced a sign error in test_integrated_gradient_trapezoidal_jacobian_jacrev, can't figure out why yet

[stefan-apollo]

Found the bug, with Jake's help: In the test we were computing the old version of the formula (without the const - ...). This happened in both the jacrev function _integrated_gradient_jacobian_with_jacrev and our function integrated_gradient_trapezoidal_jacobian. But with the new changes integrated_gradient_trapezoidal_jacobianalways implements the new (const - ... alpha) one while the test had the old one hardcoded. Thus the test and our function disagreed. Updating _integrated_gradient_jacobian_with_jacrev to match fixes this.

[stefan-apollo]

test_modular_arithmetic_build_graph is supposed to fail now with the new bases, according to Lucius. This property should no longer hold

[stefan-apollo]

The floating point test for the interaction_rotations fails for the new basis. However the test was already mostly failing for the old basis. So I suggest we scrap that particular test.

[nix-apollo]

Test runtime is rather large now, will reduce number of tests before merging 42.82s call tests/test_build_graph.py::test_mnist_build_graph_new_basis_old_attribution 42.65s call tests/test_build_graph.py::test_mnist_build_graph_old_basis_old_attribution 42.58s call tests/test_build_graph.py::test_mnist_build_graph_new_basis_and_new_attribution

With #223 at least mnist build tests can be made faster by reducing data amount

[stefan-apollo]

And yeah these tests take a long time, maybe we can disable some of them on CI / reduce data amount later

With https://github.com/ApolloResearch/rib/pull/223 at least mnist build tests can be made faster by reducing data amount

80.00s call     tests/test_build_graph.py::test_mnist_build_graph_old_basis_old_attribution[(1-0)*alpha-squared]
76.74s call     tests/test_build_graph.py::test_mnist_build_graph_old_basis_old_attribution[(1-alpha)^2-squared]
43.66s call     tests/test_build_graph.py::test_mnist_build_graph_old_basis_old_attribution[(1-alpha)^2-functional]
43.15s call     tests/test_build_graph.py::test_mnist_build_graph_old_basis_old_attribution[(1-0)*alpha-functional]
27.06s setup    tests/test_float_precision.py::TestPythiaFloatingPointErrors::test_gram_matrices
20.19s setup    tests/test_float_precision.py::TestPythiaFloatingPointErrors::test_ablation_result_float_precision
17.76s call     tests/test_build_graph.py::test_pythia_14m_build_graph
12.01s call     tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-alpha)^2-squared]
11.77s call     tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-alpha)^2-functional]
11.64s call     tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-0)*alpha-functional]
11.35s call     tests/test_ablations.py::test_run_mnist_ablations[orthogonal]
11.01s call     tests/test_build_graph.py::test_modular_arithmetic_build_graph[(1-0)*alpha-squared]
8.81s call     tests/test_parallel.py::TestDistributed::test_edges_are_same
7.84s call     tests/test_ablations.py::test_run_mnist_ablations[rib]
6.84s call     tests/test_train_mnist.py::test_main_accuracy

[stefan-apollo]

I do have a problem with the variable_position_dimension which has 12 occurances in the repo, and is only there in order to initialise this thing:

Oh I absolutely agree with this, I somehow missed your comment -- will change that

[stefan-apollo]

Test finished, merging!