Prepare full one-degree spatial-grid Perceiver experiment - #848
Conversation
alxmrs
left a comment
alxmrs
left a comment
There was a problem hiding this comment.
Overall, this looks really solid.
| gradient_accumulation_steps: 8 | ||
| learning_rate: 0.0006 | ||
| scheduler: {type: cosine} | ||
| loss: mse |
There was a problem hiding this comment.
Samudra 2 typically uses a dynamic mse with a limit of 20 -- we can experiment with loss configs later. I agree, this loss is good to start with.
There was a problem hiding this comment.
Agreed. I retained plain normalized MSE so this run stays comparable to the architecture searches and isolates the transport intervention. Dynamic MSE is a worthwhile later loss ablation after we have this architectural baseline.
| return resized[..., scale_width : scale_width + target_width] | ||
|
|
||
|
|
||
| class DirectCrossAttentionIO(nn.Module): |
There was a problem hiding this comment.
Should we run a search to test how many layers worth of cross and self attentions we should try?
This is surprisingly simple, but maybe that is the point.
There was a problem hiding this comment.
The simplicity is intentional for the decoder. The encoder is already a complete local Perceiver IO with input cross-attention, two latent self-attention stages, and output-query cross-attention. The decoder retains only one query-to-processor cross-attention plus FFN because the processor already mixes spatial state and Jesse’s localization rejected a second decoder latent/self-attention stack. I documented a focused follow-up over encoder depths {1,2,4} and direct cross-attention depths {1,2} after this baseline. I do not recommend delaying the full run for that search, and I would keep output-query self-attention out initially because it is quadratic within each output window and lacks supporting evidence.
6d1baff to
d3650b8
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
1 participant
Summary
This is stacked on #842 so its review diff excludes the native-SDPA implementation. After #842 lands, this branch can be rebased onto main.
Intended run
The preset assumes two GPUs, per-rank batch 2, and eight accumulation steps for effective global batch 32. It uses the shared OM4 data preset pointed at the one-degree data root. A real-data one-degree optimizer probe is required before releasing the 70-epoch allocation.
Validation