/ml-probabilistic-attention

Primary LanguagePythonOtherNOASSERTION

Project Name

This software project accompanies the research paper Probabilistic Attention for Interactive Segmentation and its predecessor Probabilistic Transformers.

It contains implementation of a Pytorch module for the probabilistic attention update proposed in the above paper(s).

Documentation

Runs an update of the probabilistic version of attention based on a Mixture of Gaussians model.
It accepts the following parameters during a forward pass:

  • q: A tensor of queries with dims N, G, C, H
  • zeta: A tensor of keys (query/key Gaussian means) with dims N, G, C, H
  • alpha: A scalar (see special case above) or tensor of query/key Gaussian precisions with dims N, G, C, H
  • mu: A tensor of value Gaussian means with dims N, G, Cv, H
  • beta: A scalar (see special case above) or tensor of value Gaussian precisions with dims N, G, C, H
  • pi: A tensor of mixture component priors with dims N, G, H, H
  • v_init: A tensor of initial vals for the values with dims N, G, Cv, H (optional)
  • v_fixed: A tensor of fixed vals for the values with dims N, G, (Cv+1), H (optional). The extra (last) channel is an indicator for the fixed val locations
  • zeta_prior_precision: A tensor of precisions for the Gaussian prior over zeta with dims N, G, C, H (optional)
  • mu_prior_precision: A tensor of precisions for the Gaussian prior over mu with dims N, G, Cv, H (optional)
  • q_pos_emb: A tensor of query positional embeddings with dims C, H, H
  • zeta_pos_emb: A tensor of key positional embeddings with dims C, H, H
  • v_pos_emb: A tensor of value positional embeddings with dims Cv, H, H
  • nonzero_wts_mask: A boolean indexing tensor for setting weight matrix values to zero (where mask value is false) with dims H, H

And returns the following output tensor:

  • Updated values with dims N, G, Cv, H if no position embedding (v_pos_emb=None) else N, G, 2*Cv, H

Notably, this layer is equivalent to a standard dot product attention (without position embeddings) when:

  • uniform_query_precision = True
  • uniform_value_precision = True
  • magnitude_priors = True
  • alpha = 1/sqrt(C) (Could be a scalar to save some memory)
  • beta = 0 (Could be a scalar to save some memory)
  • v_init = None
  • v_fixed = None

Getting Started

The module is in the file probabilisticattention.py. It can be imported as any other Pytorch layer.