Reproduce “Normalized Attention Without Probability Cage” on OpenResearch

Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.

Universal app for Apple silicon and Intel. Requires macOS 11 or later.

About this paper

Oliver Richter, Roger Wattenhofer

Researchers at ETH Zurich introduce Normalized Attention Pooling (NAP), an alternative to the softmax function in Transformer architectures that removes the probability simplex constraint. Their work demonstrates that NAP achieves superior performance, increased robustness to hyperparameter choices, and better generalization across various synthetic tasks compared to traditional softmax-based Transformers.

Key insights

  • The 'probability cage' imposed by softmax normalization, which constrains attention outputs to the convex hull of value vectors, limits the representational capacity of Transformer models.
  • Transformers with softmax attention exhibit a sequence length-dependent bias at initialization and are highly sensitive to hyperparameter choices like learning rate.
  • Replacing softmax with simpler vector normalization can lead to more robust and generalizable attention mechanisms.

Summary by alphaXiv.

Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.