In this post you seem to imply that the slow training is due to a lack of parallelization, but don’t MP-SAEs also require more total flops?
At each iteration you need to recompute the encoder dot products using a matmul with the encoder matrix (a look at your code confirms this), so I would think that the total flops would scale almost linearly as you increase the number of iterations.
Very true! Each iteration of matching pursuit uses as much compute as the entire encode() of a standard SAE, so it’s not only a parallelism problem (although it doesn’t help either). I’ll update the wording in the post.
In this post you seem to imply that the slow training is due to a lack of parallelization, but don’t MP-SAEs also require more total flops?
At each iteration you need to recompute the encoder dot products using a matmul with the encoder matrix (a look at your code confirms this), so I would think that the total flops would scale almost linearly as you increase the number of iterations.
Very true! Each iteration of matching pursuit uses as much compute as the entire
encode()of a standard SAE, so it’s not only a parallelism problem (although it doesn’t help either). I’ll update the wording in the post.