Reproduce “Spark Transformer: Reactivating Sparsity in FFN and Attention” on OpenResearch

About this paper

Chong You, Kan Wu, Zhipeng Jia, Lin Chen, Srinadh Bhojanapalli, Jiaxian Guo, Utku Evci, Jan Wassenberg and 11 more

Spark Transformer reduces Transformer computation by predicting which feed-forward network (FFN) neurons and context tokens to use, then computing only those entries. Its predictor reuses existing model dimensions, and its Gemma-2 2B version keeps the parameter count and training procedure. The authors report competitive benchmark quality and 2.5× fewer FLOPs at 8k context, with speedups limited by hardware and sparse-operation implementations.

Key insights

• High levels of activation sparsity can be reintroduced into modern Transformer FFNs and attention mechanisms while preserving model quality by leveraging a parameter-efficient predictor and a statistical sparsity operator.
• The Statistical Top-k operator provides an efficient, differentiable, and hardware-accelerator-friendly alternative to sorting-based top-k, effectively preventing training slowdowns.
• The Spark Transformer conceptually extends Mixture of Experts to a 'Mixture of Minimum Experts,' where individual neurons act as fine-grained experts with an internal routing mechanism, offering unstructured sparsity benefits.

Summary by alphaXiv.