Short Convolutions in Transformers

Summary: Lightweight depthwise convolutions (kernel-4) inserted at key points in transformer blocks — after K/V projections and on attention/MLP outputs — to provide explicit local token mixing and short-range inductive bias complementary to global attention. Used in Inkling (975B MoE) as architectural "surprises" beyond standard MoE recipe.

Overview

Pure transformers lack explicit local inductive bias — attention can learn local patterns but must discover them from data. Adding small convolutions gives "free" local mixing, reducing attention's burden for short-range dependencies. This idea appears in various forms (ConvNeXt, Megatron-BERT, CCT) but Inkling applies it systematically in a massive MoE.

Key Components

Placement Points (Inkling)

  1. After K/V projections: K = Conv1d(K_proj(x)), V = Conv1d(V_proj(x))
    • Mixes keys/values locally before attention
    • Like local attention preprocessing
  2. After attention output: attn_out = Conv1d(attn_out)
    • Smoothes attention output locally
  3. After MLP output: mlp_out = Conv1d(mlp_out)
    • Local mixing in FFN pathway

Convolution Specs

  • Kernel size: 4 (small — local only)
  • Groups: d_model (depthwise) or d_head (per-head)
  • Stride: 1 (preserves length)
  • Padding: Same (causal or symmetric)
  • Activation: None (linear) or GeLU

Computational Cost

  • Depthwise conv: O(B * L * d_model * k) — negligible vs attention O(B * L² * d_model) or MLP O(B * L * d_model * d_ffn)
  • At k=4: ~4× the cost of a bias add
  • Parameter-free if no learnable weights (just fixed smoothing) — but Inkling's are learnable

Mathematical Formulation

In plain English: A depthwise conv with kernel size 4 looks at each token and its 3 neighbors, learns a 4-tap filter, and mixes them. It's like giving attention a "local blur" preprocessing step.

# Standard attention block
Q = x @ W_Q
K = x @ W_K
V = x @ W_V
attn = softmax(Q @ K.transpose / sqrt(d)) @ V

# With short convs (Inkling style)
K = conv1d_depthwise(x @ W_K, kernel=4, groups=d_head)
V = conv1d_depthwise(x @ W_V, kernel=4, groups=d_head)
attn = softmax(Q @ K.transpose / sqrt(d)) @ V
attn = conv1d_depthwise(attn, kernel=4, groups=d_model)
mlp_out = conv1d_depthwise(mlp(x), kernel=4, groups=d_model)

Variants / Extensions

Variant Placement Purpose
Pre-Attention (K/V) After K/V proj Local key/value smoothing
Post-Attention After attn output Smooth attention aggregation
Post-MLP After FFN Local mixing in FFN path
Pre-Embedding After token embed Early local structure (CCT)
Hybrid Conv-Attn Replace some attn layers Local layers = conv (ConvNeXt, MetaFormer)

Applications

  • Inkling: 66-layer MoE, convs at all 3 points in every block
  • ConvNeXt / MetaFormer: Conv stem + transformer body
  • Megatron-BERT: Convolutional relative position
  • CCT (Compact Conv Transformers): Conv token embedding + transformer
  • Local Attention alternatives: Conv as O(L) local mixer instead of sliding window

Historical Context

Year Work Contribution
2019 CCT Conv tokenization + transformer
2020 Megatron-BERT Conv relative position
2021 ConvBERT Conv + attention hybrid
2022 ConvNeXt Pure conv matches transformer
2022 MetaFormer Conv as "token mixer" in MetaFormer
2023 Various Depthwise conv in attention (RetNet, etc.)
2026 Inkling Systematic depthwise convs in 975B MoE at 3 points/block

Related Concepts

Sources

  • inkling-raschka-blog: Architecture observations #1 — Short kernel-4 convolutions after K/V projections and on attention/MLP outputs; cheap local token mixing + explicit short-range inductive bias