metaformer

MetaFormer: Meta-Learning with Transformers

Overview

MetaFormer is a transformer-based model that incorporates both spatial mixing and channel mixing blocks. The architecture is designed to learn from few examples and generalize to new tasks.

For more info, refer to the original paper MetaFormer: Meta-Learning with Transformers.

Parameters are in helia_edge.models.metaformer_params.

Functions

Name	Description
`build`	MetaFormer model from `MetaFormerParams`
`patch_embedding`	Patch embedding layer
`pool_token_mixer`	Token mixer using average pooling
`conv_token_mixer`	Token mixer using separable convolution
`attention_token_mixer`	Token mixer using multi-head attention
`mlp_channel_mixer`	Channel mixer using MLP via 1x1 convolutions
`metaformer_block`	Metaformer block
`metaformer_layer`	Metaformer functional layer

Machine-readable model

patch_embeddingfunctionPatch embedding layer using 2D convolution
pool_token_mixerfunctionToken mixer using average pooling
conv_token_mixerfunctionToken mixer using separable convolution
attention_token_mixerfunctionToken mixer using multi-head attention
mlp_channel_mixerfunctionChannel mixer using MLP via 1x1 convolutions
metaformer_blockfunctionMetaformer block
metaformer_layerfunctionMetaFormer functional layer
buildfunctionBuild a MetaFormer model.

patch_embedding(
    embed_dim: int,
    patch_shape: tuple[int, int],
    stride_shape: tuple[int, int] | None = None,
    padding: str = 'same',
) -> keras.layers.Layer

Patch embedding layer using 2D convolution

Parameters

Parameters of patch_embedding
Name	Type	Default	Description
`embed_dim`	`int`	Required	Embedding dimension
`patch_shape`	`tuple[int, int]`	Required	Patch shape
`stride_shape`	`tuple[int, int]`	None	Stride shape. Defaults to None.
`padding`	`str`	'same'	Padding. Defaults to 'same'.

pool_token_mixer(pool_size: tuple[int, int] = (2, 2)) -> keras.layers.Layer

Token mixer using average pooling

Parameters

Parameters of pool_token_mixer
Name	Type	Default	Description
`pool_size`	`tuple[int, int]`	(2, 2)	Pool size. Defaults to (2, 2).

Returns

Returns of pool_token_mixer
Type	Description
`keras.layers.Layer`	keras.layers.Layer: Token mixer layer

conv_token_mixer(
    embed_dim: int,
    kernel_size: tuple[int, int] = (3, 3),
    strides: tuple[int, int] = (1, 1),
) -> keras.Layer

Token mixer using separable convolution

Parameters

Parameters of conv_token_mixer
Name	Type	Default	Description
`embed_dim`	`int`	Required	Embedding dimension
`kernel_size`	`tuple[int, int]`	(3, 3)	Kernel size. Defaults to (3, 3).
`strides`	`tuple[int, int]`	(1, 1)	Strides. Defaults to (1, 1).

Returns

Returns of conv_token_mixer
Type	Description
`keras.Layer`	keras.layers.Layer: Token mixer layer

attention_token_mixer(embed_dim: int, num_heads: int, dropout: float = 0.1) -> keras.layers.Layer

Token mixer using multi-head attention

Parameters

Parameters of attention_token_mixer
Name	Type	Default	Description
`embed_dim`	`int`	Required	Embedding dimension
`num_heads`	`int`	Required	Number of heads
`dropout`	`float`	0.1	Dropout rate. Defaults to 0.1.

Returns

Returns of attention_token_mixer
Type	Description
`keras.layers.Layer`	keras.layers.Layer: Token mixer layer

mlp_channel_mixer(embed_dim: int, ratio: int = 4, activation: str = 'gelu', dropout: float = 0) -> keras.Layer

Channel mixer using MLP via 1x1 convolutions

Parameters

Parameters of mlp_channel_mixer
Name	Type	Default	Description
`embed_dim`	`int`	Required	Embedding dimension
`ratio`	`int`	4	Expansion ratio. Defaults to 4.
`activation`	`str`	'gelu'	Activation function. Defaults to "gelu".
`dropout`	`float`	0	Dropout rate. Defaults to 0.

Returns

Returns of mlp_channel_mixer
Type	Description
`keras.Layer`	keras.layers.Layer: Channel mixer layer

metaformer_block(
    token_mixer: keras.layers.Layer | None = None,
    channel_mixer: keras.layers.Layer | None = None,
    name: str = 'mf_block',
) -> keras.layers.Layer

Metaformer block

Parameters

Parameters of metaformer_block
Name	Type	Default	Description
`token_mixer`	`keras.layers.Layer`	None	Token mixer layer. Defaults to None.
`channel_mixer`	`keras.layers.Layer`	None	Channel mixer layer. Defaults to None.
`name`	`str`	'mf_block'	Block name. Defaults to 'mf_block'.

Returns

Returns of metaformer_block
Type	Description
`keras.layers.Layer`	keras.layers.Layer: Metaformer block

metaformer_layer(x: keras.KerasTensor, params: MetaFormerParams) -> keras.KerasTensor

MetaFormer functional layer

Parameters

Parameters of metaformer_layer
Name	Type	Default	Description
`x`	`keras.KerasTensor`	Required	Input tensor
`params`	`MetaFormerParams`	Required	Model parameters.

Returns

Returns of metaformer_layer
Type	Description
`keras.KerasTensor`	keras.KerasTensor: Output tensor

build(
    params: MetaFormerParams,
    input_shape: tuple[int | None, ...],
    *,
    batch_size: int | None = None,
    name: str | None = None,
) -> keras.Model

Build a MetaFormer model.

Parameters

Parameters of build
Name	Type	Default	Description
`params`	`MetaFormerParams`	Required	Model parameters.
`input_shape`	`tuple[int \| None, ...]`	Required	Input shape without the batch axis; None for a variable axis.
`batch_size`	`int \| None`	None	Static batch size; None for a dynamic batch.
`name`	`str \| None`	None	Model name; the family when None.

Returns

Returns of build
Type	Description
`keras.Model`	keras.Model: The model, named ``metaformer`` unless ``name`` is given.