mlsanityData to 14 Oct 2026

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

arXiv 2506.00691 ↗First posted 31 May 2025cs.LG cs.AIcited by 0

On the map: Deep reinforcement learning methods → Deep RL algorithms and representations

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.

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.

Builds on

27 of 37 references.

References · 27

  1. 2017Attention Is All You Needcited by 60,361
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
  3. 2020Language Models are Few-Shot Learnerscited by 32,419
  4. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  5. 2016Layer Normalizationcited by 5,669
  6. 2020Taming Transformers for High-Resolution Image Synthesiscited by 3,253
  7. 2016OpenAI Gymcited by 1,736
  8. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,264
  9. 2021Perceiver: General Perception with Iterative Attentioncited by 1,104
  10. 2016HyperNetworkscited by 1,027
  11. 2015Effective Approaches to Attention-based Neural Machine Translationcited by 1,026
  12. 2018Recurrent World Models Facilitate Policy Evolutioncited by 785
15 more
  1. 2014Neural Turing Machinescited by 644
  2. 2018Video Action Transformer Networkcited by 316
  3. 2018Differentiable MPC for End-to-end Planning and Controlcited by 189
  4. 2016Using Fast Weights to Attend to the Recent Pastcited by 170
  5. 2019Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learningcited by 104
  6. 2018Differentiable plasticity: training plastic neural networks with backpropagationcited by 84
  7. 2019Towards Interpretable Reinforcement Learning Using Attention Augmented Agentscited by 82
  8. 2019Weight Agnostic Neural Networkscited by 71
  9. 2020Neuroevolution of Self-Interpretable Agentscited by 49
  10. 2020Meta-Learning through Hebbian Plasticity in Random Networkscited by 47
  11. 2020Backpropamine: training self-modifying neural networks with differentiable neuromodulated plasticitycited by 41
  12. 2019Learning Video Representations using Contrastive Bidirectional Transformercited by 39
  13. 2015Deep Attention Recurrent Q-Networkcited by 27
  14. 2019Learning to Predict Without Looking Ahead: World Models Without Forward Predictioncited by 15
  15. 2017Multi-focus Attention Network for Efficient Deep Reinforcement Learningcited by 10

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025A Method of Selective Attention for Reservoir Based Agents
  2. 2025Attention Trajectories as a Diagnostic Axis for Deep Reinforcement Learning
  3. 2025Learning to Focus: Prioritizing Informative Histories with Structured Attention Mechanisms in Partially Observable Reinforcement Learning
  4. 2024Attention-Driven Multi-Agent Reinforcement Learning: Enhancing Decisions with Expertise-Informed Tasks
  5. 2025SPECTra: Scalable Multi-Agent Reinforcement Learning with Permutation-Free Networks
  6. 2025Interpretable Learning Dynamics in Unsupervised Reinforcement Learning
  7. 2024Revealing the Learning Process in Reinforcement Learning Agents Through Attention-Oriented Metrics
  8. 2025Meta-reinforcement learning with minimum attention