rla

Documentation

Overview

RLA (Recurrent Linear Attention) "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) NewProc(ctx nn.Context) nn.Processor
• type Processor
  • func (p *Processor) Forward(xs ...ag.Node) []ag.Node
  • func (p *Processor) LastState() *State
  • func (p *Processor) SetInitialState(state *State)
• type State

Types

type Config

type Config struct {
	InputSize int
}

type Model

type Model struct {
	Config
	Wk *nn.Param `type:"weights"`
	Bk *nn.Param `type:"biases"`
	Wv *nn.Param `type:"weights"`
	Bv *nn.Param `type:"biases"`
	Wq *nn.Param `type:"weights"`
	Bq *nn.Param `type:"biases"`
}

func New

func New(config Config) *Model

func (*Model) NewProc

func (m *Model) NewProc(ctx nn.Context) nn.Processor

type Processor

type Processor struct {
	nn.BaseProcessor

	States []*State
	// contains filtered or unexported fields
}

func (*Processor) Forward

func (p *Processor) Forward(xs ...ag.Node) []ag.Node

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

func (*Processor) LastState

func (p *Processor) LastState() *State

func (*Processor) SetInitialState

func (p *Processor) SetInitialState(state *State)

type State

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