Anish Mudide(Anish Mudide)

• Anish Mudide 23 Jul 2024 20:54 UTC

  3 points

  0

  in reply to: Aaron_Scher’s comment on: Effi­cient Dic­tionary Learn­ing with Switch Sparse Autoencoders

  Thanks for your comment! I believe your concern was echoed by Lee and Arthur in their comments and is completely valid. This work is primarily a proof-of-concept that we can successfully scale SAEs by directly applying MoE, but I suspect that we will need to make tweaks to the architecture.

• Anish Mudide 23 Jul 2024 20:47 UTC

  3 points

  0

  in reply to: Sheikh Abdur Raheem Ali’s comment on: Effi­cient Dic­tionary Learn­ing with Switch Sparse Autoencoders

  Yes, you can train a Switch SAE in any scenario where you can train a standard SAE. @hugofry has a nice blog post on training SAEs for ViT.

• Anish Mudide 23 Jul 2024 20:44 UTC

  4 points

  0

  in reply to: phenomanon’s comment on: Effi­cient Dic­tionary Learn­ing with Switch Sparse Autoencoders

  Thanks for the question -- $f$ is calculated over an entire batch of inputs, not a single $x$. Figure 1 shows how the Switch SAE processes a single residual stream activation $x$.

• Anish Mudide 23 Jul 2024 20:32 UTC

  5 points

  0

  in reply to: Arthur Conmy’s comment on: Effi­cient Dic­tionary Learn­ing with Switch Sparse Autoencoders

  Hi Lee and Arthur, thanks for the feedback! I agree that routing to a single expert will force redundant features and will experiment with Arthur’s suggestion. I haven’t taken a close look at the router/​expert geometry yet but plan to do so soon.

• Anish Mudide 23 Jul 2024 20:17 UTC

  3 points

  1

  in reply to: Noa Nabeshima’s comment on: Effi­cient Dic­tionary Learn­ing with Switch Sparse Autoencoders

  Thanks for the comment—I trained TopK SAEs with various widths (all fitting within a single GPU) and observed wider SAEs take substantially longer to train, which leads me to believe that the encoder forward pass is a major bottleneck for wall-clock time. The Switch SAE also improves memory efficiency because we do not need to store all $M$ latents.

  I’m currently working on implementing expert-parallelism, which I hope will lead to substantial improvements to wall-clock time.

Effi­cient Dic­tionary Learn­ing with Switch Sparse Autoencoders