Mix-FFN in SegFormer: Adding Local Context to Transformer MLPs

Mix-FFN is SegFormer's spatially aware replacement for the ordinary Transformer feed-forward network. It inserts a depthwise $3\times3$ convolution between the two linear projections:

The convolution operates on a 2D feature map, so the token sequence must be reshaped from [batch, tokens, channels] to [batch, channels, height, width] and then flattened again.

Why add a convolution?

An ordinary token-wise MLP processes every spatial position independently. Self-attention mixes global information, but the MLP itself has no explicit local neighborhood. The depthwise convolution gives each expanded channel a local receptive field with relatively little additional computation.

SegFormer does not use positional embeddings in its encoder. The paper argues that the zero-padded $3\times3$ convolution in Mix-FFN leaks enough location information for this role. Avoiding fixed positional embeddings also avoids interpolating them when the test resolution differs from the training resolution; in the paper's comparison, Mix-FFN was less sensitive to that change.

The order matters

SegFormer expands the channel width, reshapes the tokens into their 2D grid, applies the depthwise convolution, and then applies GELU. Afterward it flattens the grid and projects the channels back down. Changing that order defines a different block and will not reproduce a pretrained SegFormer checkpoint.

Practical cautions

• The number of tokens must equal height × width; rectangular inputs must retain both dimensions.
• A depthwise convolution processes each expanded channel separately rather than mixing channels.
• Padding affects the positional signal and boundary behavior.
• Mix-FFN is an encoder component, not the lightweight multiscale MLP decoder that gives SegFormer its name.

Primary sources

• Enze Xie et al., SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers, NeurIPS 2021.
• NVIDIA, official SegFormer implementation.