rla

Documentation

Overview

Package rla provides an implementation of RLA (Recurrent Linear Attention). See: "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention" by Katharopoulos et al., 2020. TODO: support arbitrary mapping functions

Index

• type Config
• type Model
  • func New(config Config) *Model
  • func (m *Model) Forward(xs ...ag.Node) []ag.Node
  • func (m *Model) LastState() *State
  • func (m *Model) SetInitialState(state *State)
• type State

Types

type Config

type Config struct {
	InputSize int
}

Config provides configuration settings for a RLA Model.

type Model

type Model struct {
	nn.BaseModel
	Config
	Wk     nn.Param `spago:"type:weights"`
	Bk     nn.Param `spago:"type:biases"`
	Wv     nn.Param `spago:"type:weights"`
	Bv     nn.Param `spago:"type:biases"`
	Wq     nn.Param `spago:"type:weights"`
	Bq     nn.Param `spago:"type:biases"`
	States []*State `spago:"scope:processor"`
}

Model contains the serializable parameters for an RLA neural network.

func New

func New(config Config) *Model

New returns a new RLA Model, initialized according to the given configuration.

func (*Model) Forward

func (m *Model) Forward(xs ...ag.Node) []ag.Node

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

func (*Model) LastState

func (m *Model) LastState() *State

LastState returns the last state of the recurrent network. It returns nil if there are no states.

func (*Model) SetInitialState

func (m *Model) SetInitialState(state *State)

SetInitialState sets the initial state of the recurrent network. It panics if one or more states are already present.

type State

type State struct {
	S ag.Node
	Z ag.Node
	Y ag.Node
}

State represent a state of the RLA recurrent network.