JsT V4 — Just seismic Transformer

An 8-token conditional flow-matching generative model for three-component P-wave seismograms. Turns one earthquake record into a site-effect measurement.

Model

Architecture: 8-token condition encoder (source×3, path×4, receiver×1) + 8-layer flow-matching transformer (512 dim, 8 heads), adaLN-Zero modulation, 1-D rotary position encoding
Params: 5.3M (encoder) + transformer denoiser
Training: 800 epochs, batch 1024, AdamW, cosine annealing
Data: MLAAPDE v2.1 36-month P-wave cache (56,047 source–station pairs)
Inference: 50-step Heun ODE integration, CFG scale 1.0

Usage

import torch
from JsT import load_checkpoint_models, AblationConditionEncoder

device = torch.device("cuda")
ce, dn, ckpt = load_checkpoint_models(
    "checkpoint-last.pth", device, use_ema=True,
    sampling_method="heun", steps=50, cfg_scale=1.0,
)
ce = AblationConditionEncoder(ce, [8, 9, 10])  # remove identity tokens
dn.eval(); ce.eval()

Full code: Just-seismic-Transformer

Reference

Preprint forthcoming. When using this checkpoint, please cite the corresponding paper and the JsT code repository.

License

This checkpoint is released under CC BY 4.0. The accompanying code is released under MIT.

Downloads last month: -; Downloads are not tracked for this model. How to track