Junu Choi
GitHub

Publications

Efficient probabilistic weather forecasting at quarter-degree resolution with two-sided transport attention

Astra forecasts on the 0.25° grid with Kantorovich Attention, in which each region chooses both where it reads from and where it writes to, so every region sends and receives information at token resolution. Trained on one node of eight GPUs, it matches GenCast with 9.8% of its hardware-normalised training compute.

2026 Junu Choi, Youngjoon Hong · arXiv:2610.35545 arXiv PDF Code
Scorecard: the CRPS of Astra is lower than that of the ECMWF ensemble, IFS ENS, for all eight variables at all 30 lead times from 12 hours to 15 days, by up to about 30% at short lead times; every cell is blue and significant.
CRPS improvement over the ECMWF ensemble (IFS ENS) in WeatherBench 2 at 0.25°: lower on every variable and lead time.
Kantorovich Attention in motion

Each region (tile) of the grid chooses both the tiles it reads from (pull) and the tiles it writes to (push). Together the two choices form a transport plan in which every tile sends and receives at token resolution.

  • Pull Each tile picks the K tiles it reads from. Alone, it leaves some tiles unread (60 of 180 on the trained 2° globe).
  • Push Each tile picks the K tiles it writes to, so every tile is read by at least K others.
  • Two-sided plan Without a linear-programming solver, the union of pull and push leaves 0 unread tiles by construction and reaches 94–98% of the optimal transport plan’s gain.

Two-sided 3D tile routing

Animation cycling through Pull alone (leaving 2 senders unread), Push alone (leaving 2 receivers reading nothing), and Kantorovich Attention (Pull union Push) on 3D raised Earth plates with twisting transport tubes.
Pull gathers K senders into query tile i; push scatters i to K receivers. Their union satisfies both marginal floors (≥ K) and achieves total score 21.8 vs. the LP optimum 21.9.

Global coverage on trained scores

Animation of the 180 global 20-degree raised Earth tiles on a trained L2 attention head: one-sided Pull leaves 60 of 180 tiles unread (hatched), whereas Two-sided top-K rescues every unread tile (0 unread by construction).
Across 4,608 trained score matrices, one-sided pull leaves 28–38% of global tiles unread. Two-sided top-K guarantees 0% unread tiles by construction and captures 94–98% of the optimal transport gain.
Figure 2: a, one tile at the top of a cut-away Earth reads from, writes to, and exchanges with selected tiles through transport tubes, while every other tile reaches it through a pooled term. b, plans on the trained network's scores: pull alone leaves 60 of 180 tiles unread, the two-sided top-K plan none, close to the optimal plan. c, a toy plan with at least K partners per tile on both sides.
Figure 2. Full architecture overview: 3D cut-away globe (a), trained-network score plans (b), and bipartite 3D plate decomposition (c).

Posterior Transport Matching for Inverse Problems

Posterior Transport Matching pairs numerical transport maps with minimum-energy stochastic rematching so each map’s inputs remain posterior-compatible using only a frozen denoiser and the observation operator.

2026 Junu Choi, Youngjoon Hong · Under review, ICLR 2027
Posterior Transport Matching empirical results: (a) image restoration on ImageNet and FFHQ with 48x48 zoom insets and PSNR gains up to +5.83 dB over SITCOM and DAPS, and (b) InverseBench scientific inverse problems across black-hole interferometry, inverse scattering, multi-coil knee MRI, and acoustic full waveform inversion.
Image restoration with 48×48 zoom insets (a; ImageNet & FFHQ 256×256, PSNR in dB) and scientific inverse problems on InverseBench (b; single-particle PTM vs. training-free diffusion baselines).
Geometry in motion

A map changes volume. An observation changes the shape of uncertainty.

Local volume

A smooth map deforms a square grid. Two initially equal cells change area by factors of 0.48 and 1.35 at the endpoint.
Equal source cells acquire different areas under a smooth transport map. The counters track each cell’s area relative to its starting value.

Conditioning

As observation precision increases, the rust Gaussian kernel narrows along the observed direction. Its variance falls from 0.8 to 0.073, while the unobserved variance stays at 0.8. Blue contours show the prior convolved with that kernel.
Stronger conditioning narrows the Gaussian kernel along the observed direction. Blue contours show the same prior convolved with that kernel.
Figure 2: numerical transport maps alternate with rematching. Blue current laws are compared with dashed required input laws and their common posterior target.
Figure 2. Rematching supplies a compatible input law.

Illustrative 2D geometry. The blue contours describe a smoothed prior, rather than posterior marginals.

2026