ReLU²: Thresholded Activation for MoE

Hard · Solved in PyTorch · Machine learning coding practice

Problem

Recent research (ReLU2 Wins) suggests that the squared ReLU function f(x) = \text{ReLU}(x)^2, is highly effective for Sparse LLMs / MoE models. It improves the trade-off between model performance and activation sparsity.

To exploit this for hardware acceleration, production systems often apply a Sparsity Threshold (\epsilon). Any activation value below this threshold is "zeroed out" to allow sparse kernels to skip those neurons entirely.

Topics: sparse-computation, llm-optimization, autograd-custom, activation-sparsity

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve ReLU²: Thresholded Activation for MoE in the browser and run it against the tests.

More problems