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

TL;DR

This paper introduces a nonlinear attention mechanism for permutation-invariant neural networks in reinforcement learning, significantly speeding up training convergence while maintaining performance.

Contribution

It proposes a novel nonlinear attention method that enhances representational capacity and accelerates convergence in RL agents with permutation-invariant architectures.

Findings

Faster convergence in RL training

Maintains baseline performance

Improved training efficiency

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.