torch_blue.vi.VIMultiheadAttention

class torch_blue.vi.VIMultiheadAttention(embed_dim: int, num_heads: int, bias: bool = True, add_bias_kv: bool = False, add_zero_attn: bool = False, kdim: int | None = None, vdim: int | None = None, batch_first: bool = True, variational_distribution: Distribution | List[Distribution] = MeanFieldNormal(), prior: Distribution | List[Distribution] = MeanFieldNormal(), kaiming_initialization: bool = True, prior_initialization: bool = False, rescale_prior: bool = True, return_log_probs: bool = True, device: torch.device | None = None, dtype: torch.dtype | None = None)

Bases: torch_blue.vi.base.VIModule

Allows the model to jointly attend to information from different representation subspaces.

Equivalent of nn.MultiheadAttention with variational inference. See its documentation for usage.

This does not support the dropout argument. In addition to all other arguments, this class accepts VIkwargs.

This module’s random variables are:

  • (“in_proj_weight”, “out_proj_weight”) if kdim and vdim are None or equal to embed_dim.

  • (“q_proj_weight”, “k_proj_weight”, “v_proj_weight”, “out_proj_weight”) else.

Additional random variables are appended in the following order based on the given arguments:

  • (“in_proj_bias”, “out_proj_bias”) if bias is True.

  • (“bias_k”, “bias_v”) if add_bias_kv is True.

Parameters:

  • torch_args – The same arguments and keyword arguments as the pytorch version MultiheadAttention (documentation here), except dropout, which is not relevant for BNNs.

  • VIkwargs – Several standard keyword arguments. See VIkwargs for details.

forward(query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: torch.Tensor | None = None, key_padding_mask: torch.Tensor | None = None, average_attn_weights: bool = True, is_causal: bool = False) → Tuple[torch.Tensor, torch.Tensor | None]

Compute attention outputs using query, key, and value embeddings.

Supports optional parameters for padding, masks and attention weights. See documentation of nn.MultiheadAttention for details. For technical reasons need_weights is hardcoded as True.

This implementation also currently does not support the torch fastpath.