multiheadattention

Documentation

Index

• type KeysValuesPairs
• type Model
  • func New(size, numOfHeads int, useCausalMask bool) *Model
  • func (m *Model) Forward(qkv attention.QKV) Output
  • func (m *Model) ForwardWithPastKeysValues(qkv attention.QKV, pastProjKeysValues KeysValuesPairs) Output
• type Output

Types

type KeysValuesPairs

type KeysValuesPairs = []attention.KeysValuesPair

KeysValuesPairs contains the attention.KeysValuesPair for each attention head.

type Model

type Model struct {
	nn.BaseModel
	Attention   []*selfattention.Model
	OutputMerge *linear.Model
	NumOfHeads  int // number of heads
	Dm          int // input and output vectors dimension
	Dk          int // hidden vectors dimension (Dm / NumOfHeads)
}

Model contains the serializable parameters.

func New

func New(size, numOfHeads int, useCausalMask bool) *Model

New returns a new model with parameters initialized to zeros.

func (*Model) Forward

func (m *Model) Forward(qkv attention.QKV) Output

Forward performs the forward step for each input node and returns the result.

func (*Model) ForwardWithPastKeysValues

func (m *Model) ForwardWithPastKeysValues(qkv attention.QKV, pastProjKeysValues KeysValuesPairs) Output

ForwardWithPastKeysValues performs the forward step for each input node and returns the result.

type Output

type Output struct {
	// Result of the multi-head attention.
	AttOutput []ag.Node
	// AttWeights attention scores.
	AttWeights [][]mat.Matrix
	// ProjKeysValues contains the attention.KeysValuesPair for each attention head.
	ProjKeysValues KeysValuesPairs
}

Output aggregates the multiple output of the multi-head attentions, incl. attention scores and last projected keys and values for each head.