"""Original reduced-width ResNet-18 reference with the original stage pattern. 7x7 stride-2 stem, max pool, [2,2,2,2] basic blocks, widths [8,16,32,64]. Post-addition ReLU, batch normalization and projection shortcuts are explicit. Untrained 64x64 RGB input; ten output classes. Standard model uses wider stages. """ import torch import torch.nn as nn import torch.nn.functional as F class IdentityBlock(nn.Module): def __init__(self, channels): super().__init__() self.conv1 = nn.Conv2d(channels, channels, 3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(channels) self.relu1 = nn.ReLU() self.conv2 = nn.Conv2d(channels, channels, 3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(channels) self.relu_out = nn.ReLU() def forward(self, x): branch = self.bn2(self.conv2(self.relu1(self.bn1(self.conv1(x))))) return self.relu_out(x + branch) class ProjectionBlock(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.conv1 = nn.Conv2d( in_channels, out_channels, 3, stride=2, padding=1, bias=False ) self.bn1 = nn.BatchNorm2d(out_channels) self.relu1 = nn.ReLU() self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(out_channels) self.shortcut = nn.Conv2d(in_channels, out_channels, 1, stride=2, bias=False) self.shortcut_norm = nn.BatchNorm2d(out_channels) self.relu_out = nn.ReLU() def forward(self, x): branch = self.bn2(self.conv2(self.relu1(self.bn1(self.conv1(x))))) shortcut = self.shortcut_norm(self.shortcut(x)) return self.relu_out(shortcut + branch) class IdentityStage(nn.Module): def __init__(self, channels): super().__init__() self.first = IdentityBlock(channels) self.second = IdentityBlock(channels) def forward(self, x): return self.second(self.first(x)) class DownsampleStage(nn.Module): def __init__(self, in_channels, out_channels): super().__init__() self.first = ProjectionBlock(in_channels, out_channels) self.second = IdentityBlock(out_channels) def forward(self, x): return self.second(self.first(x)) # tensorviz: input=1,3,64,64 class ReducedResNet18(nn.Module): def __init__(self): super().__init__() self.stem = nn.Conv2d(3, 8, 7, stride=2, padding=3, bias=False) self.stem_norm = nn.BatchNorm2d(8) self.stem_relu = nn.ReLU() self.pool = nn.MaxPool2d(3, stride=2, padding=1) self.stage1 = IdentityStage(8) self.stage2 = DownsampleStage(8, 16) self.stage3 = DownsampleStage(16, 32) self.stage4 = DownsampleStage(32, 64) self.average = nn.AdaptiveAvgPool2d(1) self.flatten = nn.Flatten(1) self.head = nn.Linear(64, 10) def forward(self, x): x = self.pool(self.stem_relu(self.stem_norm(self.stem(x)))) x = self.stage1(x) x = self.stage2(x) x = self.stage3(x) x = self.stage4(x) return self.head(self.flatten(self.average(x))) def verify(): torch.manual_seed(17) block = IdentityBlock(8).double().eval() with torch.no_grad(): block.conv2.weight.zero_() x = torch.randn(1, 8, 6, 6, dtype=torch.float64, requires_grad=True) torch.testing.assert_close(block(x), x.relu(), rtol=0, atol=0) block(x).sum().backward() torch.testing.assert_close(x.grad, (x > 0).to(x.dtype), rtol=0, atol=0) projection = ProjectionBlock(8, 16).double().eval() with torch.no_grad(): projection.conv2.weight.zero_() torch.testing.assert_close( projection(x), projection.shortcut_norm(projection.shortcut(x)).relu() ) model = ReducedResNet18().double().eval() image = torch.randn(1, 3, 64, 64, dtype=torch.float64) hidden = model.pool(model.stem_relu(model.stem_norm(model.stem(image)))) for stage, shape in [ (model.stage1, (1, 8, 16, 16)), (model.stage2, (1, 16, 8, 8)), (model.stage3, (1, 32, 4, 4)), (model.stage4, (1, 64, 2, 2)), ]: hidden = stage(hidden) assert hidden.shape == shape torch.testing.assert_close( model.average(hidden).flatten(1), hidden.mean((-2, -1)) ) assert model(image).shape == (1, 10) return [ "A zero residual branch leaves ReLU(x), including the original post-addition activation", "The shortcut carries unit gradients on positive inputs when the residual branch is zero", "Stride-two projection shortcuts match the residual branch channel and spatial dimensions", "Four [2,2,2,2] stages produce 16/8/4/2-pixel maps before global averaging", ] def experiment(): torch.manual_seed(17) original = IdentityBlock(8).double().eval() x = torch.linspace(-1, 1, 64, dtype=torch.float64).reshape( 1, 1, 8, 8 ) + torch.linspace(-0.4, 0.4, 8, dtype=torch.float64).reshape(1, 8, 1, 1) cases = [] for scale in [0.0, 0.5, 1.0, 2.0]: block = IdentityBlock(8).double().eval() block.load_state_dict(original.state_dict()) with torch.no_grad(): block.conv2.weight.mul_(scale) branch = block.bn2(block.conv2(block.relu1(block.bn1(block.conv1(x))))) joined = x + branch output = block(x) cases.append( { "id": f'scale-{str(scale).replace(".","-")}', "label": f"Residual output scale {scale:g}", "target": "stage1.first", "note": "The first basic-block form is tested on fixed synthetic features while its final convolution weights are scaled. The shortcut remains fixed. Original ResNet basic blocks apply ReLU after addition, so a zero branch returns ReLU(x), not every signed input unchanged. This demonstrates arithmetic and a gradient path, not an empirical training-stability result.", "matrices": [ {"label": "Shortcut · channel 0", "values": x[0, 0].tolist()}, { "label": "Residual branch · channel 0", "values": branch[0, 0].tolist(), }, { "label": "After addition + ReLU · channel 0", "values": output[0, 0].tolist(), }, ], "vectors": [ { "label": "Before addition · eight channels at pixel (4,4)", "values": x[0, :, 4, 4].tolist(), }, { "label": "After addition · before ReLU", "values": joined[0, :, 4, 4].tolist(), }, {"label": "After ReLU", "values": output[0, :, 4, 4].tolist()}, ], "metrics": [ {"label": "Residual branch scale", "value": scale}, { "label": "Elements clipped by output ReLU", "value": (joined < 0).sum().item(), }, ], } ) return { "kind": "matrices", "title": "A shortcut around two convolutions.", "description": "Recorded intervention on an untrained identity-shortcut basic block. The full graph preserves ResNet-18's stage pattern, while this comparison isolates a single residual addition.", "controlLabel": "Residual branch intervention", "cases": cases, }