selfattention

Documentation

Index

• type Config
• type Model
  • func New(config Config) *Model
  • func (m *Model) Forward(qkv attention.QKV) attention.Output
  • func (m *Model) ForwardWithPastKeysValues(qkv attention.QKV, past attention.KeysValuesPair) attention.Output

Types

type Config

type Config struct {
	InputSize     int
	QuerySize     int
	KeySize       int
	ValueSize     int
	ScaleFactor   mat.Float
	UseCausalMask bool
}

Config provides configuration settings for a Self-Attention Model.

type Model

type Model struct {
	nn.BaseModel
	Config
	Query *linear.Model
	Key   *linear.Model
	Value *linear.Model
}

Model contains the serializable parameters.

func New

func New(config Config) *Model

New returns a new model with parameters initialized to zeros.

func (*Model) Forward

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

Forward performs the forward step for each input node and returns the result. It generates the queries, keys and values from the same input xs.

func (*Model) ForwardWithPastKeysValues

func (m *Model) ForwardWithPastKeysValues(qkv attention.QKV, past attention.KeysValuesPair) attention.Output

ForwardWithPastKeysValues performs the forward step for each input node and returns the result. It generates the queries, keys and values from the same input xs.