Skip to content

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.

Classes:

Functions:

Classes

NameArgs

Name and arguments

Attributes:

  • name (str) –

    Name

  • args (dict) –

    Arguments

MetaFormerBlockParams

MetaFormer block parameters

Attributes:

  • layers (int) –

    Number of layers

  • patch_embed (dict) –

    Patch embedding

  • token_mixer (NameArgs) –

    Token mixer

  • channel_mixer (NameArgs) –

    Channel mixer

MetaFormerParams

MetaFormer parameters

Attributes:

  • blocks (list[MetaFormerBlockParams]) –

    MetaFormer blocks

  • output_filters (int) –

    Output filters

  • output_activation (str | None) –

    Output activation

  • include_top (bool) –

    Include top

  • dropout (float) –

    Dropout rate

  • drop_connect_rate (float) –

    Drop connect rate

  • name (str) –

    Model name

MetaFormerModel

Helper class to generate model from parameters

Functions

layer_from_params staticmethod 
layer_from_params(inputs: keras.Input, params: MetaFormerParams | dict, num_classes: int | None = None)

Create layer from parameters

Source code in neuralspot_edge/models/metaformer.py 
472
473
474
475
476
477
@staticmethod
def layer_from_params(inputs: keras.Input, params: MetaFormerParams | dict, num_classes: int | None = None):
    """Create layer from parameters"""
    if isinstance(params, dict):
        params = MetaFormerParams(**params)
    return metaformer_layer(x=inputs, params=params, num_classes=num_classes)
model_from_params staticmethod 
model_from_params(inputs: keras.Input, params: MetaFormerParams | dict, num_classes: int | None = None)

Create model from parameters

Source code in neuralspot_edge/models/metaformer.py 
479
480
481
482
483
@staticmethod
def model_from_params(inputs: keras.Input, params: MetaFormerParams | dict, num_classes: int | None = None):
    """Create model from parameters"""
    outputs = MetaFormerModel.layer_from_params(inputs=inputs, params=params, num_classes=num_classes)
    return keras.Model(inputs=inputs, outputs=outputs)

Functions

patch_embedding 

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:

  • embed_dim (int) –

    Embedding dimension

  • patch_shape (tuple[int, int]) –

    Patch shape

  • stride_shape (tuple[int, int], default: None ) –

    Stride shape. Defaults to None.

  • padding (str, default: 'same' ) –

    Padding. Defaults to 'same'.

Source code in neuralspot_edge/models/metaformer.py 
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def 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

    Args:
        embed_dim (int): Embedding dimension
        patch_shape (tuple[int, int]): Patch shape
        stride_shape (tuple[int, int], optional): Stride shape. Defaults to None.
        padding (str, optional): Padding. Defaults to 'same'.

    """
    # SHAPE (B, F, T, C) -> (B, F//S, T//S, E)
    # SHAPE (B, H, W, C) -> (B, H//S, W//S, E
    return keras.layers.Conv2D(
        filters=embed_dim,
        kernel_size=patch_shape,
        strides=stride_shape or patch_shape,
        padding=padding,
        use_bias=False,
    )

pool_token_mixer 

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

Token mixer using average pooling

Parameters:

  • pool_size (tuple[int, int], default: (2, 2) ) –

    Pool size. Defaults to (2, 2).

Returns:

  • Layer –

    keras.layers.Layer: Token mixer layer

Source code in neuralspot_edge/models/metaformer.py 
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
def pool_token_mixer(
    pool_size: tuple[int, int] = (2, 2),
) -> keras.layers.Layer:
    """Token mixer using average pooling

    Args:
        pool_size (tuple[int, int], optional): Pool size. Defaults to (2, 2).

    Returns:
        keras.layers.Layer: Token mixer layer
    """

    def layer(x: keras.KerasTensor) -> keras.KerasTensor:
        y = x
        y = (
            keras.layers.AveragePooling2D(
                pool_size=pool_size,
                strides=(1, 1),
                padding="same",
            )(y)
            - y
        )
        return y

    # END DEF
    return layer

conv_token_mixer 

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:

  • embed_dim (int) –

    Embedding dimension

  • kernel_size (tuple[int, int], default: (3, 3) ) –

    Kernel size. Defaults to (3, 3).

  • strides (tuple[int, int], default: (1, 1) ) –

    Strides. Defaults to (1, 1).

Returns:

  • Layer –

    keras.layers.Layer: Token mixer layer

Source code in neuralspot_edge/models/metaformer.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
def 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

    Args:
        embed_dim (int): Embedding dimension
        kernel_size (tuple[int, int], optional): Kernel size. Defaults to (3, 3).
        strides (tuple[int, int], optional): Strides. Defaults to (1, 1).

    Returns:
        keras.layers.Layer: Token mixer layer
    """

    def layer(x: keras.KerasTensor) -> keras.KerasTensor:
        y = x
        y = keras.layers.SeparableConv2D(
            filters=embed_dim,
            kernel_size=kernel_size,
            strides=strides,
            padding="same",
            use_bias=False,
        )(y)
        return y

    # END DEF
    return layer

attention_token_mixer 

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

Token mixer using multi-head attention

Parameters:

  • embed_dim (int) –

    Embedding dimension

  • num_heads (int) –

    Number of heads

  • dropout (float, default: 0.1 ) –

    Dropout rate. Defaults to 0.1.

Returns:

  • Layer –

    keras.layers.Layer: Token mixer layer

Source code in neuralspot_edge/models/metaformer.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
def attention_token_mixer(
    embed_dim: int,
    num_heads: int,
    dropout: float = 0.1,
) -> keras.layers.Layer:
    """Token mixer using multi-head attention

    Args:
        embed_dim (int): Embedding dimension
        num_heads (int): Number of heads
        dropout (float, optional): Dropout rate. Defaults to 0.1.

    Returns:
        keras.layers.Layer: Token mixer layer
    """

    def layer(x: keras.KerasTensor) -> keras.KerasTensor:
        input_shape = keras.ops.shape(x)
        height = input_shape[1]
        width = input_shape[2]
        channels = input_shape[3]
        y = x  # Shape is B, H, W, C
        y = keras.layers.Reshape((height * width, channels))(y)
        y = keras.layers.MultiHeadAttention(
            num_heads=num_heads,
            key_dim=embed_dim,
            dropout=dropout,
        )(y, y)
        # Reshape back to B, H, W, C
        y = keras.layers.Reshape((height, width, channels))(y)
        return y

    # END DEF
    return layer

mlp_channel_mixer 

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:

  • embed_dim (int) –

    Embedding dimension

  • ratio (int, default: 4 ) –

    Expansion ratio. Defaults to 4.

  • activation (str, default: 'gelu' ) –

    Activation function. Defaults to "gelu".

  • dropout (float, default: 0 ) –

    Dropout rate. Defaults to 0.

Returns:

  • Layer –

    keras.layers.Layer: Channel mixer layer

Source code in neuralspot_edge/models/metaformer.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
def mlp_channel_mixer(
    embed_dim: int,
    ratio: int = 4,
    activation: str = "gelu",
    dropout: float = 0,
) -> keras.Layer:
    """Channel mixer using MLP via 1x1 convolutions

    Args:
        embed_dim (int): Embedding dimension
        ratio (int, optional): Expansion ratio. Defaults to 4.
        activation (str, optional): Activation function. Defaults to "gelu".
        dropout (float, optional): Dropout rate. Defaults to 0.

    Returns:
        keras.layers.Layer: Channel mixer layer
    """

    def layer(x: keras.KerasTensor) -> keras.KerasTensor:
        y = x
        y = keras.layers.Conv2D(
            filters=int(embed_dim * ratio),
            kernel_size=(1, 1),
            activation=activation,
        )(y)

        if 0 < dropout < 1:
            y = keras.layers.Dropout(dropout)(y)

        y = keras.layers.Conv2D(
            filters=embed_dim,
            kernel_size=(1, 1),
        )(y)

        if 0 < dropout < 1:
            y = keras.layers.Dropout(dropout)(y)

        return y

    # END DEF
    return layer

metaformer_block 

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:

  • token_mixer (Layer, default: None ) –

    Token mixer layer. Defaults to None.

  • channel_mixer (Layer, default: None ) –

    Channel mixer layer. Defaults to None.

  • name (str, default: 'mf_block' ) –

    Block name. Defaults to 'mf_block'.

Returns:

  • Layer –

    keras.layers.Layer: Metaformer block

Source code in neuralspot_edge/models/metaformer.py 
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
def 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

    Args:
        token_mixer (keras.layers.Layer, optional): Token mixer layer. Defaults to None.
        channel_mixer (keras.layers.Layer, optional): Channel mixer layer. Defaults to None.
        name (str, optional): Block name. Defaults to 'mf_block'.

    Returns:
        keras.layers.Layer: Metaformer block
    """

    def layer(x: keras.KerasTensor) -> keras.KerasTensor:
        # Shape is B, T, F, C
        y = x

        if token_mixer:
            # Apply layer normalization
            y1 = keras.layers.LayerNormalization(name=f"{name}.token_ln", axis=-1)(y)

            # Apply token mixer
            y1 = token_mixer(y1)

            # Add residual connection
            y = y + y1
        # END IF

        if channel_mixer:
            # Apply layer normalization
            y1 = keras.layers.LayerNormalization(name=f"{name}.ch_ln", axis=-1)(y)

            # Apply channel mixer
            y1 = channel_mixer(y1)

            # Add residual connection
            y = y + y1

        return y

    # END DEF

    return layer

metaformer_layer 

metaformer_layer(x: keras.KerasTensor, params: MetaFormerParams, num_classes: int | None = None) -> keras.KerasTensor

MetaFormer functional layer

Parameters:

  • x (KerasTensor) –

    Input tensor

  • params (MetaFormerParams) –

    Model parameters.

  • num_classes (int, default: None ) –

    Number of classes.

Returns:

  • KerasTensor –

    keras.KerasTensor: Output tensor

Source code in neuralspot_edge/models/metaformer.py 
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
def metaformer_layer(
    x: keras.KerasTensor,
    params: MetaFormerParams,
    num_classes: int | None = None,
) -> keras.KerasTensor:
    """MetaFormer functional layer

    Args:
        x (keras.KerasTensor): Input tensor
        params (MetaFormerParams): Model parameters.
        num_classes (int, optional): Number of classes.

    Returns:
        keras.KerasTensor: Output tensor
    """
    y = x

    for b, block in enumerate(params.blocks):
        print("Adding block", b + 1)
        # Apply patch embedding
        y = patch_embedding(**block.patch_embed)(y)
        for lyr in range(block.layers):
            token_mixer = None
            if block.token_mixer.name == "conv":
                token_mixer = conv_token_mixer(**block.token_mixer.args)
            elif block.token_mixer.name == "pool":
                token_mixer = pool_token_mixer(**block.token_mixer.args)
            elif block.token_mixer.name == "attention":
                token_mixer = attention_token_mixer(**block.token_mixer.args)
            # END IF
            channel_mixer = None
            if block.channel_mixer.name == "mlp":
                channel_mixer = mlp_channel_mixer(**block.channel_mixer.args)
            # END IF
            y = metaformer_block(
                token_mixer=token_mixer,
                channel_mixer=channel_mixer,
                name=f"B{b+1}.L{lyr+1}",
            )(y)
        # END FOR
    # END FOR

    if params.include_top:
        name = "top"
        y = keras.layers.GlobalAveragePooling2D(name=f"{name}.gap")(y)
        if 0 < params.dropout < 1:
            y = keras.layers.Dropout(params.dropout, name=f"{name}.dropout")(y)
        if num_classes is not None:
            y = keras.layers.Dense(num_classes, name=name)(y)
        if params.output_activation:
            y = keras.layers.Activation(params.output_activation)(y)