My previous ViT-SR model handled broad structure decently well but lost fine detail. Scaling parameters alone didn't fix it β a larger model with the same spatial compression through patch merging scored worse than the baseline. The problem wasn't capacity. It was architecture.
The reconstruction head was being asked to recover spatial detail from a heavily compressed token grid β information that was already gone before upsampling began. Fixing that required a different design, not just more parameters.
What changed: Constant 16Γ16 token grid across all 4 transformer stages. No patch merging, no token downsampling. Every token maps to the same 4Γ4 pixel region from input through to reconstruction (DenseNet's Principle to Transformers)
Hierarchical embed dims [192 β 256 β 288 β 384]. These architectural changes and the parameter count are not separable β GDFN, dense concatenation, and hierarchical dims each require sufficient channel depth to function. A 786K version of this architecture would be a fundamentally different, weaker model.
Dense concatenation of all 4 stage outputs directly to the reconstruction head β simultaneously accessing low-level edge maps and high-level semantic context.
Gated Depthwise Feed-Forward Network (GDFN) β replaces standard MLPs with gated depthwise 3Γ3 convolutions, injecting local spatial context into every attention step.
Bilinear skip connection β model learns residual correction only, not full reconstruction from scratch.
Results:
DIV2K validation: 25.20 dB PSNR, SSIM 0.8298 ~23.8M parameters, trained from scratch on LSDIR (84,991 images)
The model handles smooth regions and structured content well. High-frequency fine detail remains the harder case β likely a property of global attention at this token count rather than anything specific to this architecture.
Built a ViT for Γ4 image super-resolution from scratch in PyTorch β sharing the model.
No pretrained weights. Every component implemented from scratch: strided Conv2d patch embedding, multi-head self-attention across 1,024 tokens, 6 pre-norm transformer blocks, and a PixelShuffle reconstruction head for learned upsampling.
Trained on real-images from LSDIR dataset with fp16 AMP on a laptop GPU. Tiled inference handles arbitrary input sizes.
Current architecture: patch size 2, embed dim 64, 4 attention heads, 6 transformer blocks, ~786K parameters β test PSNR 23.30 dB.
The model handles broad structure well β fine textures and sharp edges need more capacity. Working on a larger configuration next.