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.

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.

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?

• 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