gmlp

Documentation

Overview

Package gmlp implements a model composed by basic MLP layers with gating mechanism. Reference: `Pay Attention to MLPs` by Liu et al, 2021 (https://arxiv.org/pdf/2105.08050.pdf)

Index

• type Block
  • func NewBlock(config BlockConfig) *Block
• type BlockConfig
• type Config
• type Model
  • func New(config Config) *Model
  • func (m *Model) Forward(xs ...ag.Node) []ag.Node
• type PreNorm
  • func NewPreNorm(dim int, block *Block) *PreNorm
  • func (m *PreNorm) Forward(xs ...ag.Node) []ag.Node
• type Residual
  • func NewResidual(preNorm *PreNorm) *Residual
  • func (m *Residual) Forward(xs ...ag.Node) []ag.Node

Types

type Block

type Block struct {
	nn.BaseModel
	*stack.Model
}

Block is the core model of the gMLP.

func NewBlock

func NewBlock(config BlockConfig) *Block

NewBlock returns a new Block.

type BlockConfig

type BlockConfig struct {
	Dim        int
	DimFF      int
	SeqLen     int
	Activation ag.OpName
}

BlockConfig provides configuration parameters for a single Block of the gMLP Model.

type Config

type Config struct {
	Dim        int
	Depth      int
	SeqLen     int
	FFMult     int
	Activation ag.OpName
}

Config provides configuration parameters for a the gMLP Model.

type Model

type Model struct {
	Config Config
	*stack.Model
}

Model contains the serializable parameters.

func New

func New(config Config) *Model

New returns a new Model.

func (*Model) Forward

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

Forward performs the forward step. It adds pads if necessary.

type PreNorm

type PreNorm struct {
	nn.BaseModel
	Block *Block
	Norm  *layernorm.Model
}

PreNorm is a helper model to perform pre-normalization.

func NewPreNorm

func NewPreNorm(dim int, block *Block) *PreNorm

NewPreNorm returns a new PreNorm.

func (*PreNorm) Forward

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

Forward performs the forward step.

type Residual

type Residual struct {
	nn.BaseModel
	PreNorm *PreNorm
}

Residual is a helper model to perform residual connections.

func NewResidual

func NewResidual(preNorm *PreNorm) *Residual

NewResidual returns a new Residual.

func (*Residual) Forward

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

Forward performs the forward step.