CorrectorNet Architecture

The CorrectorNet is a U-Net convolutional neural network that maps a 101-channel ocean state to a 101-channel additive correction field. It is designed to be computationally efficient while capturing multi-scale spatial structures in ocean dynamics.

U-Net Architecture

The network follows the classic encoder-decoder U-Net design with skip connections:

Input (B, 101, H, W)
        │
    ┌───▼──── ResBlock (101 → 256) ────────────────────────┐
    │                                                       │
    │   Down (256 → 512)                                    │
    │        │                                               │
    │   Down (512 → 1024)                                   │
    │        │                                               │
    │   Down (1024 → 2048)  ← bottleneck                    │
    │        │                                               │
    │   Up (2048+1024 → 1024) ──── skip from Down2 ──────   │
    │        │                                               │
    │   Up (1024+512 → 512) ────── skip from Down1 ──────   │
    │        │                                               │
    │   Up (512+256 → 256) ─────── skip from inc ────────   │
    │        │                                               │
    └────────▼───────────────────────────────────────────────┘
        │
    Conv2d (256 → 101, kernel=1)  ← zero-initialized
        │
   Output (B, 101, H, W)

Channel Dimensions

Stage	Module	Input Channels	Output Channels
Input	—	101	—
Encoder 1	inc (ResBlock)	101	256
Encoder 2	down1 (Down)	256	512
Encoder 3	down2 (Down)	512	1024
Bottleneck	down3 (Down)	1024	2048
Decoder 1	up1 (Up)	2048 + 1024 (skip)	1024
Decoder 2	up2 (Up)	1024 + 512 (skip)	512
Decoder 3	up3 (Up)	512 + 256 (skip)	256
Output	outc (Conv2d 1×1)	256	101

Building Blocks

ResBlock

The ResBlock is a residual convolutional block with two 3×3 convolutions:

x ──┬── Conv2d(3×3, in→mid) ── BatchNorm ── ReLU ── Conv2d(3×3, mid→out) ── BatchNorm ── [+ shortcut] ── ReLU ──→ out
    │
    └── [1×1 Conv + BatchNorm if in_channels ≠ out_channels, else identity]

$$\text{ResBlock}(x)=\text{ReLU}(\text{BN}({\text{Conv}}_{3\times 3}(\text{ReLU}(\text{BN}({\text{Conv}}_{3\times 3}(x)))))+\text{shortcut}(x))$$

The shortcut connection is an identity when in_channels == out_channels, and a 1×1 convolution followed by batch normalization when they differ. This design ensures smooth gradient flow and makes training stable even at depth.

Down (Down-sampling)

class Down(nn.Module):
    def __init__(self, in_channels, out_channels):
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            ResBlock(in_channels, out_channels),
        )

Each Down block halves the spatial resolution via $2\times 2$ max pooling, then applies a ResBlock that doubles the channel dimension. The spatial dimension at the bottleneck is $H/8\times W/8$.

Up (Up-sampling)

class Up(nn.Module):
    def __init__(self, in_channels, out_channels, bilinear=True):
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
            self.conv = ResBlock(in_channels, out_channels, in_channels // 2)
        else:
            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
            self.conv = ResBlock(in_channels, out_channels)

Each Up block doubles the spatial resolution (via bilinear interpolation by default) and halves the channel dimension. The upsampled feature map is concatenated with the corresponding encoder skip connection along the channel axis. Padding is applied as needed to handle odd spatial dimensions. The default setting uses bilinear interpolation (bilinear=True), which is more memory-efficient and avoids checkerboard artifacts common with transposed convolutions.

Zero-Initialized Output

A critical design choice: the final 1×1 output convolution is zero-initialized:

self.outc = nn.Conv2d(hidden_dim, out_channels, kernel_size=1)
nn.init.constant_(self.outc.weight, 0.0)
nn.init.constant_(self.outc.bias, 0.0)

This ensures that at initialization, the CorrectorNet outputs exactly zero. The initial correction is the identity — the physics model runs unchanged. The network then gradually learns to produce non-zero corrections as training progresses. This trick, originally from the diffusion model literature, prevents the network from disrupting the physics solution early in training and leads to more stable convergence.

Input Channel Breakdown

The 101 input channels correspond to 12 prognostic fields as follows:

#	Field	Description	Vertical Levels	Channels
1	u	Baroclinic zonal velocity	16	16
2	ub	Barotropic zonal velocity	16	16
3	v	Baroclinic meridional velocity	16	16
4	vb	Barotropic meridional velocity	16	16
5	t	Temperature	16	16
6	s	Salinity	16	16
7	el	Sea surface height	—	1
8	ua	Depth-averaged zonal velocity	—	1
9	uab	Barotropic depth-avg. zonal vel.	—	1
10	va	Depth-averaged meridional velocity	—	1
11	vab	Barotropic depth-avg. meridional vel.	—	1
			Total	101

The constant SUPERVISED_KB = 16 determines how many vertical levels are supervised and thus how many channels each 3D field contributes. The dimension transform for 3D fields is:

$${\mathbb{R}}^{H\times W\times K}\stackrel{\text{permute}(0,3,1,2)}{\to }{\mathbb{R}}^{K\times H\times W}$$

SyncBatchNorm for DDP

For distributed data-parallel training, the network uses SyncBatchNorm to synchronize batch statistics across all GPUs:

model = nn.SyncBatchNorm.convert_sync_batchnorm(model)

This is critical when training on 8 GPUs with a per-GPU batch size that might be small (e.g., 1-4). Without synchronized batch statistics, the batch norm layers would see insufficient samples and produce noisy estimates, degrading training stability.

Computational Cost

For a typical grid size of $H\times W\approx 360\times 240$:

Metric	Approximate Value
Parameters	~18 million
FLOPs (forward pass)	~100 GFLOPs per $360\times 240$ input
GPU memory (inference)	~200 MB
GPU memory (training with gradients)	~600 MB per GPU
Forward pass latency	~5 ms on A100

The U-Net architecture is deliberately lightweight. The physics core (POM) is the dominant computational cost, not the neural corrector. The CorrectorNet overhead is approximately 1-2% of the total per-step computation.

Design Rationale

Several design choices merit explanation:

1. U-Net over alternatives: U-Nets are well-suited for spatial field transformations because the skip connections preserve fine-grained spatial information that would otherwise be lost in the bottleneck. Pure convolutional networks without skip connections or pure MLPs would struggle to capture the multi-scale spatial structures in ocean dynamics.

2. Symmetric encoder-decoder: The CorrectorNet uses a fully symmetric channel progression ($101\to 256\to 512\to 1024\to 2048$ and back). This ensures the network can represent corrections that are spatially localized (through the encoder path) while maintaining global context (through the decoder path).

3. No temporal modeling: The CorrectorNet operates purely on instantaneous spatial snapshots — it has no recurrent state, no attention over time, and no memory of previous states. This is intentional: the physics core handles temporal integration. The network's job is only to provide per-timestep spatial corrections.