Experimenting with large kernels in our ConvNeXt processor - #727
Closed
alxmrs wants to merge 18 commits into
Closed
Conversation
- Added LayerNorm and TrueConvNeXt blocks - Added ability to configure kernel sizes - updated unet backbone to support adjusting the kernel.
After PR #669, the OM4 data has lev-dim prognostics but masks remain split (mask_0..mask_18). The mask names match the level-encoding regex, so _is_compact wrongly returned False, sending validate_data down the non-compact branch where with_level_index_vars only handles "var_lev_<depth>" string names — never expanding the lev dim. Result: Normalize.filter(["uo_0", ...]) hit KeyError "uo_0" at startup. Excluding mask_* from the check restores compactness detection so the filter's compact branch decodes "uo_0" -> data["uo"].isel(lev=0). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`Dataset.to_array().reshape(-1)` broadcasts non-lev variables (zos) over the lev dim, producing 5*19=95 elements instead of the expected 4*19+1=77. This caused `assert data.shape[-3] == self._prognostic_mean_np.shape[0]` to fire in `unnormalize_tensor_prognostic` during validation. `_flatten` (defined in the same file) uses conditional_rearrange to handle mixed lev/non-lev variables correctly, producing the right per-channel order for both compact (lev-dim) and split data layouts. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The v1 run with decreasing kernels showed a clear surface-vs-deep trade-off against the Samudra-2 paper baseline (Yuan et al. 2026): upper-ocean R² 0.66 vs 0.87, but deep-ocean R² -6.56 vs -16.14. The surface degradation is most consistent with shallow-stage 31x31 smoothing the high-frequency input. Reversing the schedule keeps the shallowest stage at 7x7 (preserving surface fidelity) and pushes the largest kernels to the bottleneck where global-scale context is most useful for slow, low-frequency deep-ocean signals. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
v3 (`model_kernel_control.yaml`): true_conv_next_block + uniform 3x3 kernels. Isolates the block-type contribution so we can attribute the v1 large-kernel run's deep-ocean gain to either the depthwise block swap or the kernel-size schedule. v4 (`model_kernel_peakmid.yaml`): true_conv_next_block + peak-in-middle kernels [7, 21, 21, 13]. Adapts RepLKNet's actual shape (peak middle, drop at deepest) to our U-Net's feature-map sizes; keeps shallowest at 7x7 to preserve surface fidelity and avoids the >50% padding waste a 31x31 would incur on the 22x45 deepest stage. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
7x7 is the canonical ConvNeXt depthwise baseline (Liu et al. 2022). With v3=[7,7,7,7] and v1=[31,21,13,7] sharing the same deepest-stage kernel, v1 = v3 + (large kernels at shallow stages), so any v1-vs-v3 difference attributes cleanly to the shallow-stage large kernels rather than being confounded by the deepest-stage size. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
v6 (model_dense_dilated.yaml): paper's dense conv_next_block with kernel=3 and aggressive dilation [1, 8, 16, 32], pushing ERF from paper's [3,5,9,17] to [3,17,33,65] without the param explosion of large dense kernels. v7 (model_kernel_peakmid_wide.yaml): identical to v4 but with ~25% wider channels [320,480,640,720] vs [280,380,480,520]. Tests whether v4's upper-R² loss vs paper is from depthwise expressivity vs undercapacity. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Structural re-parameterization for ConvNeXt-style depthwise large-kernel
blocks, matching the canonical RepLKNet recipe (Ding et al. 2022 §5,
DingXiaoH/RepLKNet-pytorch:replknet.py):
* Each branch is Conv (bias=False) -> BatchNorm.
* Parallel small (3x3) branch alongside the main k×k branch when k>3.
* Per-branch BN sums into a single output before the ConvNeXt FFN.
* fold_reparam() merges Conv+BN per branch via standard BN folding,
then sums the two folded convs into a single equivalent k×k conv
(test asserts numerical equivalence to 1e-6).
Targets v4's surface/upper-R² regression. The reparam benefit is in
training dynamics; fold is purely an inference optimization and is not
invoked in the eval pipeline.
Also adds SPDX-License-Identifier headers to the kernel-experiment
config / script / prose files (Apache-2.0 for code+configs,
CC-BY-4.0 for exp_kernel.md), matching the repo's REUSE convention.
New YAML: configs/samudra_om4_v2/model_kernel_peakmid_reparam.yaml
mirrors v4 exactly except for block_type='rep_conv_next_block'.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
PyTorch's mode='circular' F.pad hard-fails when pad > input_size, and the deepest-stage feature map at 1° has height 22. With kernel=3 and dilation=32, N_pad=32 > 22 → 'Padding value causes wrapping around more than once' raised on the first forward, killing E1 train at step 0. Cap at dilation=16 (N_pad=16 < 22 height): still saturates the deepest feature map's 22-pixel height with the kernel taps spanning the full extent, while staying inside PyTorch's circular-pad constraint. ERFs [3, 17, 33, 33] vs paper's [3, 5, 9, 17] — still substantially more aggressive RF expansion than the published baseline. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mirrors the 1° E1 architecture (dilation [1, 8, 16, 16] in the dense conv_next_block) but keeps the v2_highres recipe's instance norm and bfloat16 (small-batch-friendly + memory). At 1° this config beat the paper baseline on Niño 3.4 R²/RMSE and reduced deep-ocean R² from -16.14 to -1.88 — testing whether the gain transfers to higher resolutions. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ee `exp_multiscale.md`.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
https://ar5iv.labs.arxiv.org/html/2203.06717.