PaperlayerView exact revision on arXiv

Source-linked AI summary

Bayesian Attention Belief Networks

Shujian Zhang, Xinjie Fan, Bo Chen, Mingyuan Zhou

arXiv:2106.05251v1cs.LGcs.CLstat.ML

TL;DR

Deterministic attention can miss complex dependencies and uncertainty, motivating Bayesian attention belief networks (BABN). Across language-understanding tasks, neural machine translation, and visual question answering, BABN outperforms deterministic and state-of-the-art stochastic attention in reported performance dimensions.

  • Problem

    Deterministic attention is common but can fail to capture complex dependencies or provide uncertainty estimation.

    Source: §1 · p. 1, §3.1 · p. 2

  • Method

    BABN models unnormalized attention weights with hierarchical gamma distributions in the decoder and Weibull distributions in a deterministic-upward, stochastic-downward encoder.

    Source: §1 · p. 1, §3 · p. 2

  • Results

    BABN outperforms vanilla deterministic and state-of-the-art stochastic attention in accuracy and uncertainty estimation, with strong domain generalization and adversarial robustness.

    Source: §1 · p. 2, §6 · p. 9

  • Takeaways & Limitations

    Existing deterministic attention models, including pretrained ones, can be converted to BABN while only slightly increasing memory and computational cost.

    Source: §1 · p. 2, §6 · p. 9

    The stated Gibbs-sampling route is conditional, and the decoder passage specifies a Poisson distribution.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

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.

1. Introduction

Attention became foundational to state-of-the-art models, but most attention networks use deterministic weights, limiting their ability to model complex dependencies or estimate uncertainty. The paper introduces Bayesian attention belief networks (BABN), a stochastic attention architecture that can be optimized with a differentiable variational objective and performs strongly across several tasks.

1. Introduction

  • Transformer-based attention has become foundational to state-of-the-art models across language, vision, and multimodal tasks.The cited examples include pretrained models trained on large unlabeled datasets.Source: §1 · p. 1
  • Deterministic attention is easy to optimize but often cannot model complex dependencies or provide uncertainty estimates, motivating stochastic belief networks.Stochastic belief networks have shown advantages in modeling structured data and providing uncertainty estimation.Source: §1 · p. 1
  • BABN models unnormalized attention weights as random variables, using a hierarchical gamma decoder and a Weibull encoder with deterministic-upward and stochastic-downward paths.Existing attention keys and queries parameterize the distributions, enabling conversion of deterministic models, including pretrained ones.Source: §1 · p. 1
  • The encoder's reparameterizable Weibull distributions support differentiable variational optimization, while an analytic gamma-to-Weibull KL divergence reduces gradient-estimation variance.The objective is an evidence lower bound.Source: §1 · pp. 1–2
  • Across language understanding, translation, and visual question answering, BABN outperforms deterministic and state-of-the-art stochastic attention in accuracy and uncertainty estimation, with strong domain generalization and adversarial robustness.The method is reported to increase memory and computational cost only slightly.Source: §1 · p. 2

2. Background on Attention Networks

Attention uses keys and queries to calculate weights, then aggregates values according to those weights. Attention layers can be stacked to form deep networks, with each layer's output passed onward as input.

2. Background on Attention Networks

  • Keys and queries calculate attention weights, and the weights aggregate values to produce the attention output.The framework is described for n key-value pairs and m queries.Source: §2 · p. 2
  • The attention function exponentiates scores into positive unnormalized weights, then normalizes them across keys with softmax.This makes the weights positive and sum to one across keys.Source: §2 · p. 2
  • Attention layers can be stacked by feeding each layer's output into the next layer as its input.In self-attention, the layer input is linearly projected into keys, queries, and values.Source: §2 · p. 2
  • Residual structure, feed-forward networks, and layer normalization are also indispensable details of the network, though they do not affect the general framework described.Source: §2 · p. 2

3. BABN: Bayesian Attention Belief Networks

BABN models unnormalized attention weights as latent variables, pairing a hierarchical gamma decoder with a Weibull inference network that uses deterministic upward and stochastic downward paths. The network uses variational inference and adapts existing attention architectures with lightweight added layers.

3.1. Deep Gamma Decoder Attention Networks

  • BABN uses hierarchical gamma distributions to model positive, unnormalized attention weights as latent variables in its decoder.The gamma decoder’s shape parameters are learned from preceding attention weights and the input.Source: §3.1 · p. 2, §3 · p. 2
  • The decoder uses keys from existing attention layers to parameterize gamma distributions, sharing computation across layers instead of using separate neural networks.A two-layer MLP transforms each layer’s key to obtain its gamma shape parameter; the hierarchical sharing is described as improving efficiency and helping prevent overfitting.Source: §3.1 · p. 3

3.2. Deep Weibull Encoder Attention Networks

  • The Weibull encoder approximates the posterior over unnormalized attention weights through a deterministic upward path and a stochastic downward path.The upward path conveys likelihood information and the downward path conveys prior information; the paper reports that removing the upward path often leads to unstable performance.Source: §3.2 · pp. 3–4
  • BABN retains the deterministic attention architecture up to keys, queries, and values, then adds lightweight layers for gamma priors and Weibull variational distributions.The paper says pretrained deterministic checkpoints can initialize BABN for fine-tuning, and describes an extension to multi-head attention that models each head’s unnormalized weights separately.Source: Figure 2 · p. 4, §3.2 · pp. 4–5
  • The encoder’s distribution depends on x rather than y at test time, allowing it to be used on new data points.In sequence generation, the encoder may use parts of y already observed during training.Source: §3.2 · p. 4

3.3. Learning Bayesian Attention Belief Networks

  • Training maximizes an evidence lower bound with a likelihood term and a prior-matching KL term, using reparameterized Weibull samples and analytic conditional KL divergences.The paper reports that reparameterization and semi-analytic KL reduce Monte Carlo gradient-estimation variance while keeping the estimate unbiased; it samples the uniform noise from Uniform (0.1, 0.9) to prevent numerical instability.Source: §3.3 · p. 5

4. Related Work

Prior stochastic-attention methods face optimization or flexibility constraints, while deep stochastic networks offer a way to model complex dependencies. This work applies deep stochastic-network structure to attention using existing efficient attention architectures.

Stochastic attentions

  • Unlike Fan et al. (2020), BABN uses sample-dependent Weibull scale and shape parameters with a deterministic-upward path, giving its inference network greater flexibility.Fan et al. use a hyperparameter for the shape parameter and omit a deterministic upward path.Source: §4 · p. 5

Deep stochastic networks

  • Stacking stochastic layers can capture data randomness, estimate uncertainty, and model more complicated dependencies than shallow probabilistic models.The related-work examples include gamma belief networks for topic modeling and deep Weibull networks for scalable inference.Source: §4 · p. 5
  • The paper applies a deep stochastic-network structure to attention models while leveraging existing efficient attention architectures to build scalable networks.Source: §4 · p. 6

5. Experimental Results

Experiments apply BABN across language understanding, domain generalization, adversarial robustness, neural machine translation, and visual question answering, including by converting pretrained models during finetuning. BABN improves on deterministic attention and stochastic-attention baselines across these tasks, with gains in accuracy, calibration, and robustness.

5.1.1. IN-DOMAIN PERFORMANCE EVALUATION

  • BABN outperforms deterministic attention and BAM on GLUE and SQuAD in-domain accuracy, even when applied only during finetuning.The reported results average five independent runs.Source: §5.1.1 · p. 6, Table 1 · p. 7

5.1.2. GENERALIZATION ACROSS DOMAINS

  • BABN improves both in-domain and out-of-domain accuracy for BERT-base and RoBERTa-base, with often larger gains out of domain and improved ECE.The evaluation covers natural language inference, paraphrase detection, and commonsense reasoning.Source: §5.1.2 · p. 7, Table 2 · p. 7

5.1.3. ROBUSTNESS TOWARDS ADVERSARIAL ATTACKS

  • Across Textfooler, Textbugger, and BAE attacks, BABN beats deterministic attention on most datasets and achieves better average accuracy.Textfooler is the strongest attacker; the authors attribute BABN’s greater robustness to its stochastic connections.Source: §5.1.3 · pp. 7–8, Table 3 · p. 8

5.2. Attention in Neural Machine Translation

  • On IWSLT, BABN achieves the best BLEU score among deterministic attention, variational attention, and BAM while retaining comparable model size.Its runtime is on par with BAM, slightly slower than deterministic attention, and faster than variational-attention methods.Source: §5.2 · p. 8, Table 4 · p. 8

5.3. Attention in Visual Question Answering

  • On original and noisy VQA-v2 data, BABN improves accuracy and PAvPU over deterministic attention and BAM, with larger gains on noisy data.Higher PAvPU reflects better uncertainty estimation; the authors interpret the noisy-data gains as evidence of greater robustness.Source: §5.3 · p. 8, Table 5 · p. 8
  • BABN learns query-key-specific attention uncertainty, unlike BAM’s fixed std/mean, which the authors link to stronger uncertainty estimation.The visualization uses normalized posterior mean for average importance and posterior std/mean for uncertainty.Source: §5.3 · p. 9, Figure 3 · p. 9
  • Setting either upward-downward structure parameter ρ or σ to zero reduces performance, especially for ρ.The ablation supports the necessity and effectiveness of the upward-downward structure.Source: §5.3 · p. 9

6. Conclusion

BABN models attention weights as hierarchically dependent random variables, using a generative network and an inference network built on existing attention architectures. The conclusion reports gains across several language-understanding measures and applications to neural machine translation and visual question answering.

6. Conclusion

  • BABN shows strong performance in accuracy, uncertainty estimation, domain generalization, and adversarial robustness across language-understanding tasks.The conclusion also reports consistent outperformance of corresponding baselines on neural machine translation and visual question answering.Source: §6 · p. 9
  • BABN uses a multistochastic-layer generative model and a deterministic-upward-stochastic-downward inference network built by leveraging existing attention architecture.The conclusion describes attention weights as hierarchically dependent random variables.Source: §6 · p. 9
  • Existing deterministic attention models, including pretrained ones, can be converted to BABN with only slightly increased memory and computational cost.The conclusion describes the architecture design as generic and efficient.Source: §6 · p. 9
  • Adding BABN only during finetuning already yields clear performance improvement.The conclusion identifies finetuning-stage addition as sufficient for the reported improvement.Source: §6 · p. 9

A.1.1. MODEL SPECIFICATIONS FOR IN-DOMAIN EVALUATION

The in-domain experiments use ALBERT-base as the pretrained model for extracting context embeddings, with 12 attention layers and hidden dimension 768.

A.1.1. MODEL SPECIFICATIONS FOR IN-DOMAIN EVALUATION

  • ALBERT is used as the pretrained model on large corpora to extract context embeddings.The passage describes ALBERT as a memory-efficient BERT version with parameter sharing and embedding factorization.Source: §A.1.1 · p. 13
  • The experiments use ALBERT-base with 12 attention layers and hidden dimension 768.The factorized embedding dimension is 128.Source: §A.1.1 · p. 13
  • ALBERT is described as a memory-efficient version of BERT using parameter sharing and embedding factorization.The embedding dimension for factorized embedding is 128.Source: §A.1.1 · p. 13

A.1.2. EXPERIMENTAL SETTINGS FOR IN-DOMAIN EVALUATION

The in-domain evaluation covers eight GLUE tasks and both SQuAD versions, using the HuggingFace Transformers codebase. The appendix points to tables for task settings and ALBERT-base efficiency.

A.1.2. EXPERIMENTAL SETTINGS FOR IN-DOMAIN EVALUATION

  • The experiments cover eight GLUE tasks and SQuAD v1.1 and v2.0.The GLUE tasks include MRPC, CoLA, RTE, MNLI, QNLI, QQP, SST, and STS.Source: §A.1.2 · p. 13
  • The experimental codebase is HuggingFace Transformers.The passage points to a footnote for the codebase.Source: §A.1.2 · p. 13
  • Table 6 summarizes experimental settings for each in-domain pretrained-language-model task.Its caption defines LR, BSZ, DR, TS, WS, and MSL as learning rate, batch size, dropout rate, training steps, warming steps, and maximum sentence length.Source: Table 6 · p. 13
  • Table 7 is identified as reporting efficiency on ALBERT-base models.Source: Table 7 · p. 13

A.1.3. MODEL SPECIFICATIONS FOR DOMAIN GENERALIZATIONS

For domain-generalization model specifications, the authors use BERT and RoBERTa pretrained baselines alongside two non-pretrained baselines, with separate finetuning settings for BERT and RoBERTa.

A.1.3. MODEL SPECIFICATIONS FOR DOMAIN GENERALIZATIONS

  • The pretrained baselines are bert-base-uncased and roberta-base, and the non-pretrained baselines are DA and ESIM.The DA and ESIM results come from Desai & Durrett (2020), while the pretrained models are provided by HuggingFace Transformers.Source: §A.1.3 · p. 13
  • BERT is finetuned for up to 3 epochs with batch size 16, learning rate 2e−5, gradient clip 1.0, and no weight decay.The passage says these settings largely follow Desai & Durrett (2020).Source: §A.1.3 · p. 13
  • RoBERTa is finetuned for up to 3 epochs with batch size 32, learning rate 1e−5, gradient clip 1.0, and weight decay 0.1.AdamW is used as the optimizer in the experiments.Source: §A.1.3 · p. 13

A.1.4. EXPERIMENTAL SETTINGS FOR DOMAIN GENERALIZATIONS

The appendix specifies evaluation settings for domain generalization, adversarial robustness, neural machine translation, and visual question answering. Its VQA ablation finds that removing either upward-downward weight parameter reduces performance, especially ρ.

A.1.4. EXPERIMENTAL SETTINGS FOR DOMAIN GENERALIZATIONS

  • Domain-generalization evaluations use SNLI, QQP, and SWAG for in-domain tasks, with MNLI, TwitterPPDB, and HellaSWAG as out-of-domain tests.The development set is split in half to create a held-out, non-blind test set.Source: §A.1.4 · p. 13

A.1.5. ADVERSARIAL ROBUSTNESS

A.2.1. MODEL SPECIFICATIONS

  • For NMT, a bidirectional LSTM encodes source sentences, while decoder-stage attention selects source positions to inform next-word prediction.Source: §A.2.1 · pp. 13–14

A.2.2. EXPERIMENTAL SETTINGS

  • NMT experiments use IWSLT, 14k BPE tokens, sequences up to 125 tokens, two-layer LSTMs, and beam search with size 10.The encoder has 512 units and the decoder 768 units.Source: §A.2.2 · p. 14

A.3.1. MODEL SPECIFICATIONS

  • The VQA model uses MCAN’s encoder-decoder structure with four cascaded co-attention layers combining question self-attention and question-image guided attention.Source: §A.3.1 · p. 14

A.3.2. EXPERIMENTAL SETTINGS

  • VQA-v2 evaluation uses the validation set because test labels are unavailable, and Gaussian noise with mean 0 and variance 5 is added to image features for noisy data.The dataset includes Yes/No, Number, and Other question types.Source: §A.3.2 · p. 14

A.3.3. ABLATION STUDY

  • Setting either ρ or σ to zero reduces VQA performance, especially for ρ; results are insensitive to ρ values from 1 to 4.The authors fix σ at 1.00e−6 across the experiments considered.Source: §A.3.3 · p. 14, Table 8 · p. 14
Loading 2106.05251v1…