attention

Documentation

Index

• func LinearAttention(g *ag.Graph, attIn QKV, mappingFunction MappingFunc, eps mat.Float) []ag.Node
• func ScaledDotProductAttention(g *ag.Graph, attIn QKV, scaleFactor mat.Float, useCausalMask bool) (context []ag.Node, prob []mat.Matrix)
• func ScaledDotProductAttentionConcurrent(g *ag.Graph, attIn QKV, scaleFactor mat.Float) (context []ag.Node, prob []mat.Matrix)
• type MappingFunc
• type QKV
  • func ToQKV(xs []ag.Node) QKV

Functions

func LinearAttention

func LinearAttention(g *ag.Graph, attIn QKV, mappingFunction MappingFunc, eps mat.Float) []ag.Node

LinearAttention performs the self-attention as a linear dot-product of kernel feature maps. It operates with O(N) complexity, where N is the sequence length. Reference: "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention" by Katharopoulos et al. (2020)

func ScaledDotProductAttention

func ScaledDotProductAttention(g *ag.Graph, attIn QKV, scaleFactor mat.Float, useCausalMask bool) (context []ag.Node, prob []mat.Matrix)

ScaledDotProductAttention is a self-attention mechanism relating different positions of a single sequence to compute a representation of the same sequence. This method requires that the query, the key and the value vectors have already been obtained from the input sequence. The scaled factor is the square root of the dimension of the key vectors.

func ScaledDotProductAttentionConcurrent

func ScaledDotProductAttentionConcurrent(g *ag.Graph, attIn QKV, scaleFactor mat.Float) (context []ag.Node, prob []mat.Matrix)

ScaledDotProductAttentionConcurrent does the same thing as ScaledDotProductAttention but processes input concurrently.

Types

type MappingFunc

type MappingFunc func(g *ag.Graph, x ag.Node) ag.Node

MappingFunc is a mapping function used by LinearAttention.

type QKV

type QKV struct {
	Queries []ag.Node
	Keys    []ag.Node
	Values  []ag.Node
}

QKV groups queries, keys and values useful for self-attention functions, as described in "Attention Is All You Need" (Vaswani et al., 2017 - http://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf).

func ToQKV

func ToQKV(xs []ag.Node) QKV

ToQKV create a new QKV struct with queries = keys = values = xs.