# Optimizing Sensory Neurons: Nonlinear Attention Mechanisms for Accelerated Convergence in Permutation-Invariant Neural Networks for Reinforcement Learning

arXiv: https://arxiv.org/abs/2506.00691 · HTML: https://mlsanity.com/paper/2506.00691/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 27 of its 37 references that are in the map.
- First posted on arXiv 31 May 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 27 of 27)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Taming Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2012.09841/) (2020)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Perceiver: General Perception with Iterative Attention](https://mlsanity.com/paper/2103.03206/) (2021)
- [HyperNetworks](https://mlsanity.com/paper/1609.09106/) (2016)
- [Effective Approaches to Attention-based Neural Machine Translation](https://mlsanity.com/paper/1508.04025/) (2015)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Neural Turing Machines](https://mlsanity.com/paper/1410.5401/) (2014)
- [Video Action Transformer Network](https://mlsanity.com/paper/1812.02707/) (2018)
- [Differentiable MPC for End-to-end Planning and Control](https://mlsanity.com/paper/1810.13400/) (2018)
- [Using Fast Weights to Attend to the Recent Past](https://mlsanity.com/paper/1610.06258/) (2016)
- [Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learning](https://mlsanity.com/paper/1910.05396/) (2019)
- [Differentiable plasticity: training plastic neural networks with backpropagation](https://mlsanity.com/paper/1804.02464/) (2018)
- [Towards Interpretable Reinforcement Learning Using Attention Augmented Agents](https://mlsanity.com/paper/1906.02500/) (2019)
- [Weight Agnostic Neural Networks](https://mlsanity.com/paper/1906.04358/) (2019)
- [Neuroevolution of Self-Interpretable Agents](https://mlsanity.com/paper/2003.08165/) (2020)
- [Meta-Learning through Hebbian Plasticity in Random Networks](https://mlsanity.com/paper/2007.02686/) (2020)
- [Backpropamine: training self-modifying neural networks with differentiable neuromodulated plasticity](https://mlsanity.com/paper/2002.10585/) (2020)
- [Learning Video Representations using Contrastive Bidirectional Transformer](https://mlsanity.com/paper/1906.05743/) (2019)
- [Deep Attention Recurrent Q-Network](https://mlsanity.com/paper/1512.01693/) (2015)
- [Learning to Predict Without Looking Ahead: World Models Without Forward Prediction](https://mlsanity.com/paper/1910.13038/) (2019)
- [Multi-focus Attention Network for Efficient Deep Reinforcement Learning](https://mlsanity.com/paper/1712.04603/) (2017)

## Abstract (arXiv)

Training reinforcement learning (RL) agents often requires significant computational resources and prolonged training durations. To address this challenge, we build upon prior work that introduced a neural architecture with permutation-invariant sensory processing. We propose a modified attention mechanism that applies a non-linear transformation to the key vectors (K), producing enriched representations (K') through a custom mapping function. This Nonlinear Attention (NLA) mechanism enhances the representational capacity of the attention layer, enabling the agent to learn more expressive feature interactions. As a result, our model achieves significantly faster convergence and improved training efficiency, while maintaining performance on par with the baseline. These results highlight the potential of nonlinear attention mechanisms to accelerate reinforcement learning without sacrificing effectiveness.