ReflexConv2d

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 comparison
• ablation_grid.jpg — all ablation configs side by side
• ablation_*.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 squash parameter (C scalars initialized to 1) lets the network selectively disable self-modulation on any channel.
• Residual skip uses nn.Identity when in_channels == out_channels, otherwise a 1×1 conv 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