Kairos: A Native World Model Stack for Physical AI

algorithm 2606.16533
world-modelflow-matchinglinear-attentionworld-action-modelcross-embodiment-curriculum

Kairos: A Native World Model Stack for Physical AI — L2 #

1. TL;DR #

Kairos is a 4B native world-action model that trains a video-diffusion backbone from scratch on a passive→human→robot data curriculum, replaces quadratic attention with a hybrid SWA+DSWA+GLA design (proven to bound long-horizon error via a contractive gated-delta memory), and co-designs distillation + FP8 inference for real-time edge rollout.

2. Q1 / Q2 / Q3 #

Q1 — 痛点. A useful world model for Physical AI must simultaneously (a) learn unified world knowledge from heterogeneous sources (open-world video is abundant but action-free; robot data is grounded but scarce), (b) maintain persistent state across long horizons despite dense temporal attention being prohibitively expensive and autoregressive rollout accumulating drift, and (c) run inside real observation–action–feedback loops under latency/memory budgets. Prior work optimizes one axis: pixel generators (Cosmos), latent predictors (V-JEPA), or interactive sims (Genie 3). Crucially, Kairos argues (and proves) that a bounded recent-attention window incurs an irreducible excess risk once causally-relevant events fall outside it — a structural limitation, not a capacity one.

Q2 — 方法. A native pre-training stack organized around learn/maintain/run: (1) a Cross-Embodiment Data Curriculum (CEDC) — a data pyramid Observation→Imitation→Embodiment; (2) a unified Understanding-Generation-Prediction Mixture-of-Transformers where a Video DiT and a ~5×-smaller Action DiT share a world-state, built on a hybrid linear-attention backbone (SWA local + DSWA mid-range + GLA global memory); (3) deployment-aware co-design (4-step distillation + FP8/INT4 attention quantization + tiled GatedDeltaNet). Trained with flow matching. 核心技术壁垒: the GLA global-memory branch instantiated via a gated delta update (GatedDeltaNet) whose contractivity ($\rho<1$) is what makes the theoretical long-horizon error bound hold and what enables constant-memory recurrent rollout — the theory, the architecture, and the inference kernel are all the same object. Replicating Kairos means reproducing this one mechanism as a stable, hardware-efficient, contractive memory, not merely stacking linear-attention layers.

Q3 — 结果. With only 4B params and 2.3 PFLOPs/inference, Kairos tops WorldModelBench-robot (9.30), DreamGen (AVG_Score 0.618), PAI-Bench-robot (82.57), VideoPhy (45.55, beating Cosmos-14B), and is SOTA/near-SOTA on LIBERO-Plus (90.8) and RoboTwin 2.0 (96.1, 2nd). It is 28×–85× faster than Cosmos-14B and scales linearly with duration where baselines grow exponentially; 480P generation is real-time on an A800.

3. 架构 / 方法图 #

Figure 2: Kairos framework overview

Paper's Figure 2 ("Framework of Kairos"). The top-level stack couples three pillars — learn (CEDC), maintain (unified U-G-P + hybrid memory), run (deployment co-design). Note that these are not stacked stages but a single endogenous backbone whose shared world-state threads through understanding, generation, and prediction.

Figure 4: Model architecture of Kairos

Paper's Figure 4 ("Model Architecture of Kairos"). The generation path is a conditional diffusion: high-compression video VAE → multimodal conditioning encoder → temporally-scalable DiT, with T2V/I2V/TI2V injected via cross-attention. The prediction path formulates a World-Action Model as a MoT: a Video DiT (initialized from the pretrained generator) plus an Action DiT (~1/5 scale). The masking is asymmetric — history video tokens attend only to history; future video tokens use sparse spatiotemporal attention; future action tokens use full attention and, critically, do not depend on future video tokens, enabling action-only inference.

Figure 5: DiT block with hybrid linear attention

Paper's Figure 5 ("DiT block architecture of the proposed hybrid linear attention"). The backbone is organized into $M$ groups; each interleaves SWA ($d{=}1$, local), DSWA ($d\in\{6,12\}$, mid-range), and GLA (the single global pathway). RoPE handles local positional reasoning; global positional reasoning is delegated to the linear-attention path. This is the load-bearing structural diff from a standard full-softmax DiT.

Figure 6: Gated linear attention module (GDN)

Paper's Figure 6 ("Architecture of the gated linear attention module GDN"). GDN maintains an associative memory matrix $\mathbf{S}_t$ updated by a delta rule that removes stale key-value associations and writes new ones, with an added decay gate $\alpha_t$ for global forgetting. This is the mechanism the §4 theory proves contractive.

Method core — the one novel mechanism (before/after diff):

Before (standard DiT / linear attn)After (Kairos hybrid)
Global temporal mixingfull softmax, $\mathcal{O}(N^2)$GLA (GatedDeltaNet), $\mathcal{O}(N)$
Memory across horizonKV grows with length (OOM)fixed-size contractive state $\mathbf{S}_t$
Key collisionsvanilla linear attn overwritesdelta-rule remove-then-write
Forgettingnone / uniformlearned decay gate $\alpha_t$
Positional reasoningone schemeRoPE local + linear-path global

4. 作者证明 #

Kairos includes formal necessity + sufficiency theorems for its hybrid memory (main §2.3, full proofs App-2).

Notation table:

SymbolMeaning
$Y$square-integrable long-horizon target ($Y_t^{(\tau)}\in L^2$)
$\mathcal{H}_t$full-history σ-algebra; $\mathcal{W}_t^{(w)}$ recent $w$-window
$m_t,\,m_t^{(w)}$$\mathbb{E}[Y\mid\mathcal{H}_t]$, $\mathbb{E}[Y\mid\mathcal{W}_t^{(w)}]$
$R_{\mathrm{full}}^\star,\,R_w^\star$optimal full-history / recent-window risk
$\rho<1$contraction factor of the gated delta update
$\bar\xi$max one-step perturbation error
$L,L_G$Lipschitz constants of the decoder
$\varepsilon$max per-branch approximation error

方程物理意义. The learning objective is flow matching: interpolate $\bm{z}_\sigma=(1-\sigma)\bm{z}_0+\sigma\bm{\epsilon}$ with constant target velocity $\bm{u}_\sigma=\bm{\epsilon}-\bm{z}_0$, minimizing $\mathcal{L}_{\mathrm{FM}}=\mathbb{E}\|\mathcal{V}_\theta(\bm{z}_\sigma,\sigma,\bm{c})-\bm{u}_\sigma\|_2^2$. The memory recursion $\mathbf{S}_t=\alpha_t\mathbf{S}_{t-1}+\beta_t(\mathbf{v}_t-\mathbf{v}_t^{\text{old}})\mathbf{k}_t^\top$ is one SGD step on the online regression loss $\|\mathbf{v}_t-\mathbf{S}\mathbf{k}_t\|^2$ with $\alpha_t$ as a global forget gate — associative correction plus adaptive long-term control.

Necessity — Theorem 1 + Corollary 1. The excess-risk identity

$$R_{w}^{\star}-R_{\mathrm{full}}^{\star}=\mathbb{E}\bigl[(m_{t}-m_{t}^{(w)})^{2}\bigr]=\mathbb{E}\bigl[\mathrm{Var}(m_{t}\mid\mathcal{W}_{t}^{(w)})\bigr]$$

is strictly positive iff $m_t$ is not $\mathcal{W}_t^{(w)}$-measurable, and Corollary 1 gives an explicit floor $R_{w}^{\star}-R_{\mathrm{full}}^{\star}\geq\mathbb{P}(E)\alpha(1-\alpha)(\mu_1-\mu_2)^2$. Assumption: the target genuinely depends on supra-window history (breaks for purely Markov/local dynamics, where a window suffices). The bound is information-theoretic, so scaling params/compute cannot close it.

Sufficiency — Theorem 2 + Remark 3. If the Bayes predictor factorizes into shared/local/dilated/global components and each branch is learned to error $\varepsilon$ with the global branch contractive ($\rho<1$):

$$\mathcal{R}_{t}(\hat{\mu}_{t})-\mathcal{R}_{t}^{\star}\leq\left(L\varepsilon+\frac{L_{G}\bar{\xi}}{1-\rho}\right)^{2}\quad\text{as }t\to\infty.$$

Proof sketch (expert-fillable): (1) Lemma 1 — conditional expectation is the $L^2$-optimal predictor; (2) decompose excess risk into per-branch approximation + global-memory propagation; (3) show the gated-delta recursion satisfies $e_t\leq\rho^t e_0+\frac{1-\rho^t}{1-\rho}\sup_i\xi_i$, so $e_t\leq\bar\xi/(1-\rho)$ (geometric damping); (4) combine via Lipschitz decoder. Corollary 2: with $\varepsilon=0,\bar\xi=0$ the predictor is exactly optimal.

6 minimum checks:

  1. Statement fidelity — theorems reproduced as stated (Eq. 9–12); ✓.
  2. Assumptions — supra-window dependence (necessity); component factorization + contractivity (sufficiency). Fails for Markov dynamics / non-contractive updates.
  3. Bound dependence — sufficiency bound scales as $(L\varepsilon+L_G\bar\xi/(1-\rho))^2$; blows up as $\rho\to1$. Necessity floor scales with event probability × future-mean gap.
  4. Loss decomposition — total loss $\mathcal{L}=\mathcal{L}_{video}+\lambda\mathcal{L}_{action}$; video term enforces environment dynamics, action term enforces executable control.
  5. Tightness — Corollary 2 shows the sufficiency bound is achievable (0 in the realizable case); Corollary 1 gives an explicit non-vacuous lower bound.
  6. Empirical corroboration — contractivity claim ↔ minimal 15s degradation (Table 18, §5).
  7. 5. 实验与数据 #

    Figure 3: SOTA + linear-scaling efficiency

    Paper's Figure 3. (a) world-action benchmarks, (b) embodied world-model benchmarks, (c) per-step DiT inference time. The load-bearing panel is (c): the zoom window shows Kairos's per-step time is flat vs sequence length, so total latency scales linearly while baselines grow super-linearly — the efficiency claim rests on this, not just absolute FLOPs.

    The main efficiency table (Table 5, TI2V 720P/5s, A800) makes the parameter-efficiency case concrete:

    ModelMem (GB)Complexity (PFLOPs)1-GPU (s)4-GPU (s)
    Lingbot-28B46.1347.455251436
    Cosmos-Predict2.5-14B70.2156.52526687
    Wan2.2-5B23.416.620185
    Kairos-4B23.52.3439

    Kairos has the lowest complexity (2.3 PFLOPs, ~68× below Cosmos-14B) and fastest latency (43s/9s), which is what backs the 28×–85× speedup headline.

    Figure 19: Human evaluation win rates

    Paper's Figure 19 ("Human evaluation results"). Blind ranking by 10 volunteers. Kairos-4B wins 74.1% vs Wan2.2-5B and 60.2% vs Cosmos-14B on PAI-Bench-robot, and 86.7%/65.0% on WorldModelBench — the subjective results track the automatic scores, mitigating single-metric concern.

    Figure 24: Long-horizon 15s samples

    Paper's Figure 24 ("Kairos samples on the PAI-Bench-15s dataset"). This is the empirical counterpart of the §4 contractivity theory: over 15s, object appearance and scene structure stay consistent (dog + floating bubbles; forest mist evolving smoothly). Quantitatively (Table 18) Kairos scores best overall (79.9) while Cosmos-2B/14B drop to 77.2/76.2 and Wan2.2 to 77.8 — baselines degrade on autonomous-driving/robot domains at 15s where Kairos holds.

    Ablations (load-bearing):

    AblationMetricw/o → w/
    Human-centric data (Table 9)WMBench total9.08 → 9.25
    Stronger VLM encoder (Table 9)WMBench total9.25 → 9.30
    Human-centric pretrain (Table 13)LIBERO-Plus avg83.0 → 89.0 (+6.0)
    Joint gen+action training (Table 14)LIBERO-Plus avg65.8 → 89.0 (+23.2)
    Joint denoising (Kairos-joint)LIBERO-Plus avg89.0 → 90.8

    The +23.2 from adding the video-generation objective to action prediction is the strongest single ablation: it directly evidences the "video supervision grounds action" claim.

    6. 论证链 #

    StepClaimPaper-internal support
    1A bounded attention window cannot solve long-horizon predictionTheorem 1 + Corollary 1: strictly positive, information-theoretic excess risk (Eq. 9–11)
    2Therefore a persistent global memory is requiredDefinitions 3–4 (persistent/sufficient state); Remark 2
    3A hybrid SWA+DSWA+GLA factorization can recover near-optimal riskTheorem 2 sufficiency bound (Eq. 12), contingent on contractivity
    4The GLA gated-delta update is contractive, so error does not accumulateRemark 3 geometric-damping bound $e_t\leq\bar\xi/(1-\rho)$
    5This memory also enables constant-memory recurrent rollout + action-only inference§2.2 expandability; asymmetric masking (action branch independent of future video)
    6Empirically this yields SOTA quality at 4B with linear scalingTables 5–8, 11–18; Figures 3, 19, 24

    7. 实现 cross-reference #

    Code is stated as open at https://github.com/kairos-agi/kairos-sensenova (weights on HuggingFace/ModelScope), but the L1 source does not include file:line-level pointers, so specific module locations are [实现未公开] at the L1 granularity available here.

    核心技术壁垒 (dedicated). The single hardest-to-replicate insight is the contractive gated-delta global memory (GatedDeltaNet-based GLA) used simultaneously as (a) the theoretical device that makes the sufficiency bound finite (the $1/(1-\rho)$ term), (b) the architectural sole-global-pathway that keeps the backbone linear, and (c) the inference kernel ("Tiled Gated DeltaNet with Streaming Access") that hides DRAM latency and avoids OOM for arbitrary-length rollout. A reimplementation that stacks generic linear attention but does not preserve the delta-rule remove-then-write with a learned decay gate — and does not keep the update contractive ($\rho<1$) — will lose the long-horizon consistency that the entire results section rests on. The action branch's independence from future video tokens (asymmetric masking) is the second non-obvious lever, since it is what makes deployment-time action-only inference cheap.

    关键实现细节 (easy-to-miss tricks):

    1. Shape-aware exponential timestep shifting — the flow-matching scheduler is remapped per latent shape, $s=\exp(f(L))\sqrt{F}$ (Eq. 16–19), so larger/longer latents get a stronger upward shift; without this the fixed scheduler's effective timestep distribution drifts as resolution/length scale during progressive training.
    2. Dilation interleaving pattern — SWA ($d{=}1$) and DSWA ($d\in\{6,12\}$) share the same window size; only the rearrange stride differs (Eq. 7–8), and the ActionDiT uses a fixed timestep shift (unlike VideoDiT's dynamic shift) because it optimizes stably without the dynamic schedule.