Adds option to use stochastic noise sources for lms instead of iterating over all output position indices when calculating edges.
Phis are distribution as 1s or -1s with equal probability.
Phis are the same across all alpha integration values.
Implementation is simply a calc_edge_squared formula in linalg, with the option to pass "stochastic" as an argument to edge_formula in lm_rib_build.
Motivation and Context
The current edge attribution method needs to iterate over output position dimensions. With stochastic sources, we only need to iterate over the stochastic source dimension, which should be << output_pos_dim.
How Has This Been Tested?
tests.test_build_graph.test_stochastic_source_single_pos_modadd checks that we only need 1 stochastic source to exactly replicate the "true" edges over the MLP layer in modular arithmetic. This is because there is only a single position dimension, and our stochastic sources are 1 or -1 (which is then squared).
tests.test_build_graph.test_stochastic_source_modadd_convergence checks that, as we increase the number of stochastic sources, the edges converge to those obtained using the squared edge formula. Runs on modular addition mlp layer when keeping the 3 position dimension (i.e. last_pos_module_type="unembed")
Add Stochastic Noise Sources (SNS) for Edges
Description
calc_edge_squaredformula in linalg, with the option to pass "stochastic" as an argument toedge_formulain lm_rib_build.Motivation and Context
The current edge attribution method needs to iterate over output position dimensions. With stochastic sources, we only need to iterate over the stochastic source dimension, which should be << output_pos_dim.
How Has This Been Tested?
tests.test_build_graph.test_stochastic_source_single_pos_modaddchecks that we only need 1 stochastic source to exactly replicate the "true" edges over the MLP layer in modular arithmetic. This is because there is only a single position dimension, and our stochastic sources are 1 or -1 (which is then squared).tests.test_build_graph.test_stochastic_source_modadd_convergencechecks that, as we increase the number of stochastic sources, the edges converge to those obtained using the squared edge formula. Runs on modular addition mlp layer when keeping the 3 position dimension (i.e. last_pos_module_type="unembed")There are more extensive comparison/testing code in https://github.com/ApolloResearch/interp/blob/main/interp/dan/rib/stochastic_edge_comparison.py (and results thread here). This code shows that the stochastic sources do seem to converge to the true edges for a variety of settings (in pythia).
Right now, we're not doing very thorough testing and scaling laws for this.
Does this PR introduce a breaking change?
No