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

2025/05/31 by Junaid Muzaffar, Khubaib Ahmed, Muzaffar, Junaid +6
Computer Science · Psychology · #Artificial Intelligence (cs.AI) #Emotion and Mood Recognition #FOS: Computer and information sciences #Machine Learning (cs.LG) #Neural Networks and Reservoir Computing #Reinforcement Learning in Robotics

paper · pdf · doi:10.48550/arxiv.2506.00691

openalex publication_date 2025/05/31 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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.

Citations

• Perceiver: General Perception with Iterative Attention
• Taming Transformers for High-Resolution Image Synthesis
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
• Language Models are Few-Shot Learners
• Neuroevolution of self-interpretable agents
• Network Randomization: A Simple Technique for Generalization in Deep Reinforcement Learning
• Learning to Predict Without Looking Ahead: World Models Without Forward Prediction
• Learning Video Representations using Contrastive Bidirectional Transformer
• Weight Agnostic Neural Networks
• Towards Interpretable Reinforcement Learning Using Attention Augmented Agents
• Video Action Transformer Network
• Learning Latent Dynamics for Planning from Pixels
• Differentiable MPC for End-to-end Planning and Control
• Recurrent World Models Facilitate Policy Evolution
• Multi-focus Attention Network for Efficient Deep Reinforcement Learning
• Attention Is All You Need
• Using Fast Weights to Attend to the Recent Past
• HyperNetworks
• Layer Normalization
• OpenAI Gym
• Deep Attention Recurrent Q-Network
• Effective Approaches to Attention-based Neural Machine Translation
• Neural Turing Machines
• Long Short-Term Memory
• BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding

Related

• The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning
• Neural architecture impact on identifying temporally extended Reinforcement Learning tasks
• Learning as Reinforcement: Applying Principles of Neuroscience for More General Reinforcement Learning Agents
• Advantages of biologically-inspired adaptive neural activation in RNNs during learning
• Attention-Gated Reinforcement Learning of Internal Representations for Classification