Source-linked AI summary
Bayesian Attention Modules
Xinjie Fan, Shujian Zhang, Bo Chen, Mingyuan Zhou
TL;DR
Most attention models use deterministic weights, while stochastic attention can present optimization or implementation challenges. This paper proposes Bayesian attention built from normalized reparameterizable distributions and a key-dependent prior, and reports consistent improvements over baselines across varied tasks.
Problem
Attention weights are generally deterministic, while stochastic alternatives can involve optimization difficulties or require substantial model changes.
Source: §1 · p. 1Method
The method learns attention distributions by normalizing reparameterizable random variables and regularizes them with a key-dependent contextual prior in a Bayesian framework.
Source: §1 · p. 2, §3 · p. 3, §3.3 · p. 5Results
Takeaways & Limitations
The authors present the approach as a scalable option for large-scale applications that requires only a few modifications to deterministic attention.
Source: §1 · p. 2The authors flag real-world risks from biased, long-tailed data, train–test discrepancies, and undue trust or imprecise interpretation by inexperienced practitioners.
Source: Broader Impact · p. 10
Abstract
from arXiv · showhide
Attention modules, as simple and effective tools, have not only enabled deep neural networks to achieve state-of-the-art results in many domains, but also enhanced their interpretability. Most current models use deterministic attention modules due to their simplicity and ease of optimization. Stochastic counterparts, on the other hand, are less popular despite their potential benefits. The main reason is that stochastic attention often introduces optimization issues or requires significant model changes. In this paper, we propose a scalable stochastic version of attention that is easy to implement and optimize. We construct simplex-constrained attention distributions by normalizing reparameterizable distributions, making the training process differentiable. We learn their parameters in a Bayesian framework where a data-dependent prior is introduced for regularization. We apply the proposed stochastic attention modules to various attention-based models, with applications to graph node classification, visual question answering, image captioning, machine translation, and language understanding. Our experiments show the proposed method brings consistent improvements over the corresponding baselines.
1 Introduction
Attention is widely used in neural networks and can support interpretability, yet attention weights are usually deterministic. This section motivates stochastic attention and presents a Bayesian, differentiable approach evaluated across multiple tasks.
1 Introduction
- Attention modules are important across several neural-network applications and can aid model visualization and interpretability.The introduction cites natural language processing, computer vision, graph analysis, and multi-modal learning as application areas.Source: §1 · p. 1
- Most models treat attention weights as deterministic, leaving randomness primarily to the output layer, which may be insufficient for complex dependencies.Source: §1 · p. 1
- Stochastic attention can model complex dependencies and support Bayesian regularization and uncertainty analysis.The paper connects stochastic attention weights to latent alignment models and Bayesian inference.Source: §1 · pp. 1–2
- Prior stochastic attention approaches face optimization or design constraints, including high-variance REINFORCE gradients and distributions that violate simplex constraints or lack reparameterization.Source: §1 · p. 2
- The proposed Bayesian attention uses normalized, reparameterizable Lognormal or Weibull variables and a key-dependent prior with semi-analytic KL regularization.The construction makes training differentiable and is designed to reduce gradient-estimation variance.Source: §1 · p. 2
- Across graph classification, visual question answering, image captioning, machine translation, and language understanding, the method consistently outperforms baseline attention and improves uncertainty estimation.Source: §1 · p. 2
2 Preliminaries on attention modules
The preliminaries describe soft attention as a way to use queries and keys to weight values into output features. They also introduce attention scores, multiple heads, and stacked attention layers.
2 Preliminaries on attention modules
- Queries may come from the same source as keys and values, from a decoder paired with encoder outputs, or from another modality.Source: §2 · p. 2
- Attention uses queries and keys to produce weights that aggregate values into output features.Source: §2 · p. 2
- Each attention weight indicates the importance of a key to a query, with alignment scores computed using choices such as scaled dot-product or additive attention.Source: §2 · p. 2
- Multi-head attention projects queries, keys, and values into multiple representation subspaces, and layers can be stacked so one layer's output feeds the next.Source: §2 · p. 3
3 Bayesian attention modules: a general recipe for stochastic attention
The section frames stochastic attention as Bayesian inference over data-dependent attention distributions and gives a recipe for differentiable simplex-constrained weights with a key-dependent prior.
3.1 Learning attention distributions in a Bayesian way
- Attention weights W become input-dependent local latent variables sampled from qφ parameterized by queries and keys, enabling amortized variational inference without requiring unobserved future outputs.The queries and keys come from standard attention; qφ depends on x alone or on x and the part of y already observed or generated.Source: §3.1 · p. 3
- Using keys and values keeps qφ close to standard attention, supporting pretrained-model variational finetuning, test-time availability, low added cost, and cross-layer dependence.Unlike approaches that require information unavailable at test time, qφ can be used during testing because keys and values are available in both training and testing.Source: §3.1 · p. 3
3.2 Reparameterizable attention distributions
- Normalizing nonnegative Weibull or Lognormal random variables constructs simplex-constrained attention weights that are reparameterizable for gradient-based optimization.The paper selects these distributions because they make the optimization objectives simple to optimize.Source: §3.2 · p. 4
- Replacing samples with their expectations recovers vanilla soft attention as Weibull k approaches infinity or Lognormal σ approaches zero.In these limits, the variance of the unnormalized random weights goes to zero.Source: §3.2 · p. 5
3.3 Contextual prior: key-dependent Bayesian regularization
- Conditional decomposition makes KL regularization semi-analytic, using analytic conditional KL terms and Rao-Blackwellization to reduce Monte Carlo estimation variance.The KL is computed over unnormalized weights S because the normalized attention distribution may lack an analytic density.Source: §3.3 · p. 5, §3.2 · p. 5
- The key-dependent prior encodes global importance shared across queries, while posterior attention captures each query's local importance; Figure 1 illustrates this for image captioning.In captioning, the prior expresses image-area importance before generation, and attention updates at each step using the current generation state.Source: §3.3 · p. 5, Figure 1 · p. 6
3.4 Putting it all together
- One-sample Monte Carlo gradients are unbiased and low variance; posterior means provide point predictions, while multiple posterior samples calibrate estimation uncertainty.The gradient estimate uses one sample of noise ϵ; multiple posterior samples each produce a prediction-probability sample.Source: §3.4 · p. 6
- For a single layer and head, contextual-prior memory is O(dk dmid), while added sampling and KL computation is O(mn) and prior computation O(n dk dmid), compared with deterministic attention at O(mn dkdv).The paper describes the extra memory and computation as inconsiderable relative to deterministic attention, with dmid as small as 10 and much smaller than dv.Source: §3.4 · p. 6
4 Experiments
Experiments evaluate attention modules across graph, multimodal, and language tasks. The reported comparisons favor BAM across the task-specific settings, with VQA results also showing robustness and uncertainty-estimation gains.
4.1 Attention in graph neural networks
- BAM consistently improves graph-node classification over GAT; Weibull performs better than Lognormal, and a contextual prior outperforms fixed or no prior.The experiments report mean test-node accuracy over five runs.Source: Table 1 · p. 7, §4.1 · p. 7
4.2 Attention in visual question answering models
- On noisy VQA data, BAM outperforms soft attention in accuracy, and its uncertainty estimates are better on both original and noisy data.On the original dataset, BAM performs similarly to soft attention in accuracy.Source: Table 2 · p. 7, §4.2 · p. 7
4.3 Attention in image captioning models
- BAM consistently outperforms deterministic attention on image captioning, and BAM-WC outperforms BAM-LC.The authors report that hard attention is difficult to train due to high-variance gradients; directly modeling weights with a Gaussian can produce NAN results, while a fully sampled KL term makes training unstable.Source: Table 3 · p. 8, §4.3 · p. 8
4.4 Attention in neural machine translation
- BAM-WC outperforms deterministic soft attention in BLEU and all compared variational-attention variants, with less training cost.The authors attribute the cost difference to BAM not requiring a separate variational network.Source: §4.4 · pp. 8–9, Table 4 · p. 9
4.5 Attention in pretrained language models
- BAM-WC yields consistent improvements over deterministic soft attention on both GLUE and SQuAD when fine-tuning from the same ALBERT checkpoints.The experiments use BAM only at fine-tuning; applying it during pretraining is left for future work.Source: Table 5 · p. 9, §4.5 · p. 9
5 Conclusion
BAM uses normalized reparameterizable attention distributions with a key-dependent contextual prior, and the paper reports consistent gains across tasks, including language understanding.
5 Conclusion
- BAM consistently outperforms corresponding baselines across graph classification, VQA, image captioning, and machine translation, while improving uncertainty estimation at slightly higher computational and memory cost.The paper also reports better performance than deterministic soft attention when finetuning pretrained language models.Source: §5 · p. 9
- Attention weights come from normalized reparameterizable distributions parameterized by keys and queries, with a key-dependent contextual prior making the regularization KL term semi-analytic.Source: §5 · p. 9
- BAM requires few modifications to standard attention and can finetune pretrained deterministic attention for language understanding benchmarks.Source: §5 · p. 9
- Ablations support the effectiveness of each architecture component, and the authors present BAM as an efficient alternative to deterministic attention.Source: §5 · p. 9
Broader Impact
The paper argues that its framework addresses optimization difficulty and complicated design in probabilistic attention, while noting real-world risks and the potential role of uncertainty estimation. The supplied passages also sketch the training algorithm and GAT setup.
Broader Impact
- The framework is presented as addressing optimization difficulty and complicated design in probabilistic attention, with potential applicability to attention-based models broadly.Source: Broader Impact · p. 10
- Biased, long-tailed real-world data and train–test discrepancies pose risks; uncertainty estimation may help identify when human aid is needed.Source: Broader Impact · p. 10
A Algorithm
- Algorithm 1 samples λ = sigmoid(t * ρ), computes gradients, and updates θ, η, and φ until convergence.Source: A Algorithm · p. 16
B.1.1 Model descriptions
- The GAT description defines self-attention over node features and uses N_i for node i’s neighbors; || denotes concatenation.Source: §B.1.1 · p. 16
B.1.2 Detailed experimental settings
The appendix reports graph-experiment settings and describes uncertainty evaluation with PAvPU, including its statistical confidence test and MCAN model configuration.
B.1.2 Detailed experimental settings
- The graph experiments use eight attention heads in the first GAT layer and a single head in the second, with dataset-specific changes for Pubmed and extensive BAM hyperparameter settings.The passage also specifies dropout, L2 regularization, initialization, optimizer, learning rates, and early stopping.Source: §B.1.2 · p. 16
- Table 6 summarizes basic statistics for datasets used in graph node classification.Source: Table 6 · p. 17
B.2.1 Uncertainty evaluation via PAvPU
- Prediction confidence is assessed by testing whether the top two classes’ posterior-mean probabilities differ significantly using a two-sample t-test.A p-value threshold determines whether a prediction is considered certain.Source: §B.2.1 · p. 17
- PAvPU scores the share of samples that are accurate and certain or inaccurate and uncertain, using counts for those and the two mismatched cases.Source: §B.2.1 · p. 17
B.2.2 Model descriptions
- The VQA experiments use MCAN’s six co-attention layers, combining self-attention for within-modality interactions with guided attention across image and question features.Source: §B.2.2 · p. 17
B.2.3 Detailed experimental settings
The VQA appendix specifies dataset splits, evaluation conditions, and training settings for the MCAN experiments, including hyperparameter tuning for BAM.
B.2.3 Detailed experimental settings
- VQA-v2 evaluation uses the validation set because test labels are unavailable for uncertainty evaluation, and the noisy condition adds Gaussian noise with variance 5 to image features.The dataset is split into training, validation, and testing sets; the passage gives their image and question-answer pair counts.Source: §B.2.3 · p. 17
- MCAN training uses Adam, batch size 64, up to 13 epochs, and a learning rate schedule that decays after epoch 10.The passage also specifies feature dimensions, eight attention heads, and a 3129-answer vocabulary.Source: §B.2.3 · pp. 17–18
- BAM hyperparameters are tuned using a random 20% training holdout, then the model is trained on the full training set and evaluated on the validation set.Source: §B.2.3 · p. 18
B.2.4 More results
Figure 2 presents VQA predictions and uncertainty estimates for three image-question pairs, while the image-captioning model computes attention over image features at each decoding step [10].
B.2.4 More results
- Figure 2 shows three image-question pairs with human annotations, method predictions, and uncertainty estimates.Answers are evaluated with p-values, and bold highlights the preferred answer.Source: Figure 2 · p. 18
- The figure ranks answers from correct certain to incorrect certain, with correct uncertain and incorrect uncertain between them.The accompanying table compares attention modules on visual question answering.Source: Figure 2 · p. 18, Table 7 · p. 18
B.3.1 Model descriptions
- In Att2in, RNN states are aligned with image features to compute attention weights, and the aggregated image feature feeds the next hidden state [10].The model uses an RNN decoder and performs this aggregation at each decoding step.Source: §B.3.1 · pp. 18–19
B.3.2 Detailed experimental settings
Image-captioning experiments use Att2in on MS COCO with bounding-box visual features, while the translation setup follows Deng et al. [29].
B.3.2 Detailed experimental settings
- The image-captioning experiments use MS COCO’s 123,287 images, with 113,287 for training and 5,000 each for validation and testing.Images use bounding-box features from a pretrained Faster-RCNN, replacing the ResNet-encoded features in Rennie et al. [10].Source: §B.3.2 · p. 19
- Training uses batch size 10, Adam at learning rate 5E−4, dropout 0.5, and 30 epochs, with MLE loss only.Testing uses greedy search; BAM uses contextual prior with dmid = 10 and ρ = 1.Source: §B.3.2 · p. 19
B.4.1 Model descriptions
- The translation model follows Deng et al. [29], using a bidirectional LSTM encoder and an MLP to align decoder queries with source-token representations.The aggregated feature is used to produce the distribution over the next target word.Source: §B.4.1 · p. 19
B.4.2 Detailed experimental settings
The translation experiments use the IWSLT setup from Deng et al. [29], and the language-understanding experiments use ALBERT-base.
B.4.2 Detailed experimental settings
- IWSLT is preprocessed with Byte Pair Encoding to a 14,000-token vocabulary, and training sequences are limited to length 125.The setup uses a two-layer bidirectional LSTM, batch size 6, dropout 0.3, and Adam at learning rate 3E−4; testing uses beam size 10.Source: §B.4.2 · p. 19
- The NMT comparison table reports step time and parameter counts for variational attention [26] and BAM.BAM-WC uses k = 5, β = 1E−6, ρ = 1, and dmid = 5.Source: Table 8 · p. 19, §B.4.2 · p. 19
B.5.1 Model descriptions
- The language-understanding experiments use ALBERT-base, with 12 attention layers, hidden dimension 768, and 12M parameters.Its factorized embedding dimension is 128; the passage contrasts its parameter count with BERT-base’s 108M.Source: §B.5.1 · p. 19
B.5.2 Detailed experimental settings
The experiments evaluate pretrained language models on GLUE and SQuAD and list the training-setting abbreviations used in the accompanying table.
B.5.2 Detailed experimental settings
- Evaluation covers eight GLUE tasks and both SQuAD v1.1 and SQuAD v2.0.The listed GLUE tasks include CoLA, SST, MRPC, STS-B, QQP, MNLI, QNLI, and RTE.Source: §B.5.2 · pp. 19–20
- The pretrained model is not KL-annealed; experiments use β = 1E−2 and ddim = 5, settings reported as performance-insensitive.The setup follows Lan et al. [4].Source: §B.5.2 · p. 20
- Table 9 defines LR, BSZ, DR, TS, WS, and MSL as learning rate, batch size, dropout rate, training steps, warmup steps, and maximum sentence length.The table presents experiment settings for the pretrained language model.Source: Table 9 · p. 20