A convolutional layer that modulates its own output using a mask made from its own weights, with a residual skip connection.
Problem
Standard convolutions apply the same filter uniformly across the entire spatial domain. In image-to-image tasks (autoencoders, U-Nets), this spatial uniformity contributes to over-smoothed, blurry outputs — the model cannot selectively sharpen or attenuate features at different spatial locations.
Solution
ReflexConv2d lets each layer modulate its own output spatially using a mask extracted from its convolution weights, then adds a residual connection to preserve the input signal.
| Step | Operation | Shape |
|---|---|---|
| 1 | Depthwise-like C→C convolution |
(B, C, H, W) |
| 2 | Sum weights across input channels × learnable squash |
(C, k, k) |
| 3 | Tile each k×k mask to H×W |
(C, H, W) |
| 4 | Elementwise multiply mask onto conv output | (B, C, H, W) |
| 5 | 1×1 pointwise projection |
(B, C', H, W) |
| 6 | Add residual skip (1×1 conv or identity) | (B, C', H, W) |
The spatial mask is globally learned (via squash and the conv weights) but shared across all input positions — the tiling preserves the kernel's internal structure without interpolation artifacts.
Result
In a U-Net autoencoder benchmark (both paths use residual skip connections), ReflexConv2d outperforms standard convolutions at every round of recursive encoding/decoding. Trained on Flick8k samples.
| Rounds | Standard | Reflex | Improvement | ||||
|---|---|---|---|---|---|---|---|
| L1 ↓ | PSNR ↑ | L1 ↓ | PSNR ↑ | L1 | PSNR | SSIM ↑ | |
| 1 | 0.1134 | 16.23 | 0.0933 | 17.97 | 17.7% | +1.7 dB | 0.64 → 0.74 |
| 2 | 0.1458 | 13.62 | 0.1114 | 16.49 | 23.6% | +2.9 dB | 0.54 → 0.66 |
| 4 | 0.2196 | 8.85 | 0.1416 | 14.37 | 35.5% | +5.5 dB | 0.40 → 0.53 |
| 8 | 0.4441 | 1.40 | 0.1887 | 11.86 | 57.5% | +10.5 dB | 0.22 → 0.37 |
The weight-derived mask preserves structure through repeated encoding — the model retains detail where standard convolutions degrade. The advantage grows with each recursive pass.
Ablation
What happens when we remove components?
| Config | L1 @ R8 | PSNR @ R8 | SSIM @ R8 |
|---|---|---|---|
| Standard + Residual | 0.4441 | 1.40 dB | 0.22 |
| Reflex − Residual | 0.2718 (−39%) | 8.50 dB | 0.28 |
| Reflex − Squash | 0.1896 (−57%) | 12.02 dB | 0.32 |
| Reflex Full | 0.1887 (−57%) | 11.86 dB | 0.37 |
- Residual helps — Full reflex (0.1887) beats no-residual reflex (0.2718), but even without residual, reflex still beats standard (0.4441)
- Squash has minimal impact — Removing it (0.1896) performs nearly identically to full reflex (0.1887). The raw kernel sum is the real signal
Install
pip install git+https://github.com/singam96/ReflexConv2D.git
Usage
import torch from reflex_conv2d import ReflexConv2d layer = ReflexConv2d(in_channels=64, out_channels=128, kernel_size=3) x = torch.randn(4, 64, 32, 32) y = layer(x) # (4, 128, 32, 32)
Drop it into any model:
nn.Sequential( ReflexConv2d(3, 64, 3), nn.ReLU(), ReflexConv2d(64, 64, 3), nn.ReLU(), ... )
Reproduce
# Run full ablation study (trains 4 models, generates images + metrics) python benchmark.py # Run main comparison only python demo_comparison.py
Outputs:
comparison.jpg— standard vs reflex visual comparisonablation_grid.jpg— all ablation configs side by sideablation_*.jpg— individual ablation result images- Console metrics (L1, PSNR, SSIM per round)
Test
pip install pytest python -m pytest test_reflex_conv2d.py
Notes
- The
squashparameter (Cscalars initialized to 1) lets the network selectively disable self-modulation on any channel. - Residual skip uses
nn.Identitywhenin_channels == out_channels, otherwise a1×1conv to match dimensions. - Odd kernel sizes only (1, 3, 5, 7, ...). Even kernels shift spatial dimensions due to asymmetric padding.
- Negligible parameter increase over standard
Conv2d(C, C', k):C(squash) + skip conv when channels differ.
License
Apache 2.0


























