Bayesian Attention Belief Networks
This work addresses the need for better stochastic attention mechanisms in AI models, offering a method that enhances performance and reliability across various attention-based tasks, though it is incremental as it builds on existing attention frameworks.
The paper tackled the problem of stochastic attention in neural networks, which is less explored due to optimization difficulties, by introducing Bayesian attention belief networks that model attention weights with gamma and Weibull distributions, resulting in improved accuracy, uncertainty estimation, generalization, and robustness on language tasks, with demonstrations on translation and visual question answering.
Attention-based neural networks have achieved state-of-the-art results on a wide range of tasks. Most such models use deterministic attention while stochastic attention is less explored due to the optimization difficulties or complicated model design. This paper introduces Bayesian attention belief networks, which construct a decoder network by modeling unnormalized attention weights with a hierarchy of gamma distributions, and an encoder network by stacking Weibull distributions with a deterministic-upward-stochastic-downward structure to approximate the posterior. The resulting auto-encoding networks can be optimized in a differentiable way with a variational lower bound. It is simple to convert any models with deterministic attention, including pretrained ones, to the proposed Bayesian attention belief networks. On a variety of language understanding tasks, we show that our method outperforms deterministic attention and state-of-the-art stochastic attention in accuracy, uncertainty estimation, generalization across domains, and robustness to adversarial attacks. We further demonstrate the general applicability of our method on neural machine translation and visual question answering, showing great potential of incorporating our method into various attention-related tasks.