Patchify Stem in ViT Architectures

• Patchify Stem is the initial embedding mechanism in Vision Transformers that segments images into non-overlapping patches via a large-stride convolution.
• Replacing the patchify stem with a stack of 3x3 convolutions (a convolutional stem) substantially enhances optimizer stability, hyperparameter robustness, and convergence speed.
• Empirical results show that convolutional stem-based ViTs achieve 1–2% higher top-1 accuracy on ImageNet while keeping computational costs nearly identical.

The patchify stem is the initial embedding mechanism used in Vision Transformer (ViT) architectures, wherein a raw input image is partitioned into non-overlapping patches that are linearly projected into the embedding space via a large-stride convolution. This design choice, atypical in the context of convolutional neural networks (CNNs), has been identified as a source of optimization instability and training inefficiency in vanilla ViT models. Recent work proposes replacing the patchify stem with a lightweight stack of standard convolutions—referred to as a convolutional stem—in order to substantially improve optimization stability, hyperparameter robustness, and model accuracy, all while preserving computational cost and runtime (Xiao et al., 2021).

1. Mathematical Formulation of the Patchify Stem

The canonical patchify stem in ViT operates by segmenting an input image $X \in \mathbb{R}^{H \times W \times 3}$ into non-overlapping $p \times p$ patches (typically $p=16$) and projecting each patch into a $d$-dimensional embedding space. This process is mathematically equivalent to a single 2D convolution:

%%%%4%%%%

with $i \in \{0, ..., \lfloor H/p \rfloor - 1\}$, $j \in \{0, ..., \lfloor W/p \rfloor - 1\}$, $k \in \{0, ..., d-1\}$. Here, the kernel and stride are both set to $p$, so windows of size $p \times p$ tile the image without overlap. The result $Y$ has spatial dimensions $H/p \times W/p$ and is reshaped into $N = (H/p) \cdot (W/p)$ tokens of dimension $d$ for transformer processing (Xiao et al., 2021).

2. Architectural Details: Patchify Stem vs. Convolutional Stem

The patchify stem’s use of a large stride and kernel—specifically a $p \times p$ kernel, stride $p$—differs markedly from established CNN practices, which favor smaller kernels and overlapping receptive fields. To address these issues, a convolutional stem ("ViT₍C₎") employs a stack of $3 \times 3$ convolutions (each with stride $2$), sequentially reducing resolution from $224 \times 224$ to $14 \times 14$, rather than a single immediate reduction.

A typical 4 GFLOP ViT₍C₎ stem structure consists of four layers:

• Progressive $3 \times 3$ convolutions (stride 2), with channels [48, 96, 192, 384], each followed by BatchNorm and ReLU.
• A final $1 \times 1$ (linear) convolution to match the token embedding dimension $d$.
• 4 layers for 1 GF and 4 GF models; 6 layers for 18 GF; 18 GF stem reused for 36 GF.

To maintain computational budget, one transformer block is removed when inserting the convolutional stem. This preserves overall FLOPs and model runtime (Xiao et al., 2021).

3. Impact on Optimization Dynamics

The replacement of the patchify stem yields four notable improvements:

• Training Length Stability: ViT₍P₎ (patchify) models require 300–400 training epochs on ImageNet-1k to achieve optimal accuracy. In contrast, ViT₍C₎ (convolutional) converges significantly faster; the accuracy gap between 50 and 400 epochs shrinks (from ~$10\%$ to ~$6\%$ at 1 GF), aligning training behavior with ResNets and RegNetYs.
• Optimizer Choice Stability: ViT₍P₎ suffers when trained with SGD, underperforming AdamW by up to 10 points or failing to converge. ViT₍C₎ yields nearly identical performance with either SGD or AdamW (gap $<0.2$ points).
• Hyperparameter Robustness: Error distribution for top-1 accuracy is sharply improved: at 18 GF, over 60% of runs are within 4 points of optimum with ViT₍C₎, versus less than 20% for ViT₍P₎. Effect is more pronounced with SGD.
• Gradient and Loss Stability: The convolutional stem leads to smoother loss curves and reduced gradient norm spikes in early epochs, indicating more stable dynamics at initialization.

4. Empirical Performance and Resource Parity

Under a unified training regime (AutoAugment, mixup 0.8, CutMix 1.0, label-smoothing 0.1, EMA, 400 epochs), the convolutional stem configuration demonstrates:

Model Scale	ViT₍P₎ Top-1 (%)	ViT₍C₎ Top-1 (%)	Top-1 Gap (%)
4 GFLOP	19.6	18.6	–1.0
18 GFLOP	17.9	17.0	–0.9
36 GFLOP	18.2	16.8	–1.4

On ImageNet-21k pretraining ($90$ epochs) plus ImageNet-1k fine-tuning ($20$ epochs), the gap widens: at 72 GF, ViT₍C₎ reaches 14.2%→13.6% top-1, compared to 15.1%→14.2% for ViT₍P₎. Notably, ViT₍C₎ surpasses RegNetY across all computational budgets when pretraining, a feat ViT₍P₎ cannot match (Xiao et al., 2021).

Crucially, because one transformer block is removed when inserting the convolutional stem, total FLOPs, parameter count, activations, and epoch timings are within 2% variance of the baseline—resulting in no perceptible change in training or inference throughput.

5. Ablation Studies and Isolation of Patchify Effects

Systematic ablations confirm that the patchify layer is the major source of instability.

• Intermediate “patchify+small-conv” Designs: Replacing one layer in the conv stem with a $p \times p$ convolution ($p=2,4,8,16$) leads to monotonic degradation in stability and accuracy as $p$ increases, recovering ViT₍P₎ behavior at $p=16$.
• Effect of Normalization/Activation: Adding BatchNorm and ReLU after the $16 \times 16$ patchify layer does not recover stability and marginally worsens accuracy. Replacing BN with LayerNorm slightly worsens both metrics compared to the standard convolutional stem, but considerably less than removing the convex stem.
• Deeper Networks: In a 48-block, 16 GF ViT, the convolutional stem reduces SGD/AdamW sensitivity gap by half and lowers median error by $\sim$1 point, despite constituting only 2% of total FLOPs.

6. Design Guidelines and Best Practices

For ViTs in the 1–36 GF regime (ImageNet-1k or moderate pretraining), empirical findings yield concrete recommendations:

• Replace the 16x16, stride-16 patchify convolution with a 4–6 layer stack of $3 \times 3$ convolutions with BatchNorm and ReLU, terminating in a $1 \times 1$ linear convolution to match the embedding dimension.
• Preserve training and inference cost by removing one transformer block.
• Anticipated benefits include: faster convergence, robust use of SGD and AdamW, substantial improvements in hyperparameter insensitivity (wider “good” learning rate and weight decay region), and consistent 1–2% top-1 accuracy improvements (with higher margins for large-scale pretraining).
• No observed loss of representational capacity; all advantages of ViT (e.g., transfer, scaling) are retained and realized more easily.

A small convolutional inductive bias at the model input resolves ViT training instabilities and yields improvements in final accuracy, with no increase in runtime cost. The convolutional stem is recommended as standard practice for ViT architectures in the mid-sized vision regime (Xiao et al., 2021).