Concatenation joins tensors along one chosen axis. Every other axis must already have the same size, so the operation stacks values without changing their alignment across those axes. section 6.3.2 section 6.3
For tensors shaped (B, D) and (B, H), concatenating on dim=1 produces (B, D+H). In PyTorch, torch.cat([x, h], dim=1) expresses this directly: dim=1 is the feature axis, while the batch axis remains unchanged. section 6.3
Check these conditions before concatenating:
- Choose the axis whose size is allowed to grow. section 6.3.2
- Match every other axis exactly, including spatial dimensions in image tensors. section 6.3.2
- Account for the expanded input size in the following layer; concatenating channels changes the next layer’s channel count. section 6.3 section 8.3.1
Skip connections and U-Net
A convolutional skip connection can join (B, C, W, H) with (B, C/2, W, H) on dim=1, producing (B, 3C/2, W, H). It cannot join tensors across a pooling step when the spatial shapes differ, such as (B, C, W, H) and (B, C, W/2, H/2). U-Net uses the same idea: it concatenates a full-resolution result with the bottleneck result before applying its output layers, preserving information from multiple resolutions. section 6.3.2 section 8.3.1
def forward(self, x):
full_scale_result = self.in_model(x) #1
if self.bottleneck is not None: #2
bottle_result = self.bottleneck(full_scale_result) #3
full_scale_result = torch.cat( #4
[full_scale_result, bottle_result], dim=1)
return self.out_model(full_scale_result) #5
Inspect the U-Net forward pass to see torch.cat combine full-scale and bottleneck results on the channel axis. section 8.3.1
