NiNformer: A Network in Network Transformer with Token Mixing Generated Gating Function

TL;DR

The paper tackles the high computational cost of attention in vision transformers by introducing NiNformer, a two-level Network in Network block that replaces standard attention with a dynamic gating mechanism generated by an inner MLP-Mixer token mixer. The outer network applies a token-wise MLP to the gated representations, achieving $O(n)$ complexity through element-wise operations while preserving global context. Empirical results on MNIST, CIFAR-10, and CIFAR-100 show NiNformer outperforms ViT, MLP-Mixer, and Local-ViT baselines, with substantial accuracy gains and only marginal inference-time overhead compared to MLP-Mixer. This work demonstrates that a dynamic gating signal learned via token mixing can recover dynamic token interactions efficiently, offering a practical alternative to attention-heavy transformers in computer vision.

Abstract

The attention mechanism is the primary component of the transformer architecture; it has led to significant advancements in deep learning spanning many domains and covering multiple tasks. In computer vision, the attention mechanism was first incorporated in the Vision Transformer ViT, and then its usage has expanded into many tasks in the vision domain, such as classification, segmentation, object detection, and image generation. While the attention mechanism is very expressive and capable, it comes with the disadvantage of being computationally expensive and requiring datasets of considerable size for effective optimization. To address these shortcomings, many designs have been proposed in the literature to reduce the computational burden and alleviate the data size requirements. Examples of such attempts in the vision domain are the MLP-Mixer, the Conv-Mixer, the Perciver-IO, and many more attempts with different sets of advantages and disadvantages. This paper introduces a new computational block as an alternative to the standard ViT block. The newly proposed block reduces the computational requirements by replacing the normal attention layers with a Network in Network structure, therefore enhancing the static approach of the MLP-Mixer with a dynamic learning of element-wise gating function generated by a token mixing process. Extensive experimentation shows that the proposed design provides better performance than the baseline architectures on multiple datasets applied in the image classification task of the vision domain.

Figures (4)

• Figure 1: A diagrammatic comparison of NiNformer architecture with ViT, MLP-Mixer and Local-ViT.
• Figure 2: A diagrammatic comparison of NiNformer mechanism with the attention mechanism.
• Figure 3: An illustration of the accuracy curves for NiNformer architecture.
• Figure 4: An illustration of the loss curves for NiNformer architecture.