Understanding Trajectory Heterogeneity in Federated World Model Learning

1 Introduction

• World-model learning requires enough trajectories, diverse conditions, and temporal context; isolated observations cannot provide the same examples as ordered histories and future windows.Federated learning lets clients train locally and exchange model updates while raw records remain local.Source: §1 · pp. 1–2
• Cross-time federated world modeling assigns distinct segments of the same trajectory to different clients, limiting each learner’s historical context and future training window.The study instantiates this setting with simulated severity-stage clients within MIMIC-IV; each disease is a separately trained task.Source: §1 · p. 2
• The evaluation asks how cross-time ownership shapes trajectory heterogeneity, how partitioning and forecast horizon affect prediction, and how existing algorithms respond across eight disease tasks.It connects available local windows with participation, optimization behavior, and multi-step forecast errors.Source: §1 · p. 2
• Only 7.55%–21.36% of pooled-available H32 windows had a locally complete anchor visited during the five-round runs.FedProx’s small aggregate gain was concentrated at shorter horizons.Source: §1 · p. 2
• The study defines ownership and local-access rules, compares ten federated algorithms, and connects forecast scores to temporal access, participation, and recorded update behavior.These measurements identify what must be assessed when evaluating world models trained from distributed trajectories.Source: §1 · pp. 2–3

2.1 World Models

• World-model research spans latent dynamics and imagined-trajectory planning, action-conditioned driving-scene generation, and clinical disease progression and treatment-response modeling [1] [2] [12] [13] [14] [15] [18] [19] [20] [21] [22] [25] [26] [27] [28].Examples include Delphi-2M’s future health trajectories and EHRWorld’s long-horizon patient simulation [26].Source: §2.1 · p. 3

2.2 Federated Learning

• Federated-learning methods address client differences through objective regularization, drift correction, representation alignment, server adaptation, update caching, and personalization [31] [32] [33] [34] [35] [36] [37] [38] [40] [41] [42].The cited approaches include FedProx, SCAFFOLD, FedVARP, and Ditto.Source: §2.2 · p. 3
• Federated evaluation resources include LEAF, FedML, and FedScale, while FLamby supplies healthcare tasks with natural institutional partitions [6] [7] [8] [44].Medical applications include multicenter imaging and clinical-outcome prediction.Source: §2.2 · pp. 3–4
• Temporal-data research includes heterogeneous device time series, continual forecasting, and collaborative trajectory-fragment preparation [10].Table 1 compares representative studies; this work jointly evaluates clinical dynamics, federated training, and temporal access within trajectories.Source: §2.2 · p. 4, Table 1 · p. 4
• This study focuses on action-conditioned world models whose trajectory segments belong to different clients, linking temporal ownership to local windows, participation, updates, and multi-step prediction.It characterizes histories and future windows available to each client.Source: §2.1 · p. 3, §2.2 · p. 4

3 Benchmark Construction

• The benchmark combines cross-time learning, eight clinical trajectory tasks, and a reference world model.Source: §3 · p. 4

3.1.1 Conditional trajectory prediction

• Each hourly prediction uses same-stay physiological history, treatment conditions, context, and elapsed time to predict a future state along the recorded action path.Source: §3.1.1 · pp. 4–5

3.1.2 Ownership along a trajectory

• Severity-band assignment and patient sharding let worsening or recovery move a patient’s transitions between clients, defining which temporal relationships each client can learn.Source: §3.1.2 · p. 5

3.1.3 Local access and the learning objective

• Local history retains only same-client positions within the lookback, and local future supervision stops at an ownership boundary, while evaluation rollouts may cross stages.Source: §3.1.3 · p. 5

3.2 Dataset Construction and Client Partitioning

• Eight clinical tasks are constructed through hourly alignment, patient separation, and severity-based partitioning.Source: §3.2 · p. 5

3.2.1 Cohorts and temporal cleaning

• The eight cohorts contain 40,871,490 transition memberships and 284,141 patient memberships; cohorts overlap, so these are membership counts rather than globally unique totals.Source: §3.2.1 · p. 6, Table 2 · p. 6

3.2.2 Patient separation and normalization

• Within each disease, patients are split into Train80 and Test20, while construction and normalization use Train80 and valid, fresh values define supervision and evaluation targets.Source: §3.2.2 · pp. 6–7

3.2.3 Severity partitions and observed movement

• Under SOFA, 68.47% and 86.18% of training patients span multiple groups with three and five bands; the five-band partition has a 3.94-fold group-size ratio, and severity groups overlap while their latent centroids and distributions shift.The reported multi-client proportions for OASIS are 91.68% and 97.52% with three and five bands.Source: §3.2.3 · p. 7, Table 3 · p. 7

• The reference architecture encodes visible history, predicts conditional latent dynamics, and decodes future states, with cached field semantics supporting its numerical model.Source: §3.3 · p. 7, Figure 4 · p. 9
• Action tokens and elapsed time condition a gated latent transition, with the encoded state determining action relevance and the resulting condition modulating the next transition.Source: §3.3 · p. 9
• Training combines future prediction across available horizons with reconstruction and EMA latent alignment; evaluation uses recursive prediction throughout.Source: §3.3 · p. 9

• The experiments ask how dividing trajectories across clients changes their structure, prediction quality, and algorithm updates.Source: §4 · p. 10
• SOFA/OASIS and three/five severity bands vary the temporal partition, while H1–H32 test prediction at different lengths.Source: §4 · p. 10
• The evaluation links partition granularity and forecast horizon to prediction, and examines how federated updates respond to trajectory segments.Source: §4 · p. 10

4.1.1 Algorithms and participation

• The main matrix compares ten federated algorithms across eight diseases and four partitions, using five rounds, one local epoch per participation, and 10% client participation per round.Round-5 performance is reported; FedAvg weights selected models by training-anchor counts.Source: §4.1.1 · p. 10

4.1.2 Scoring and comparison rules

• The primary score is Fresh Feature–Patient Macro nMAE, averaging time points within patient–feature pairs, patients within features, and then features across six horizons.The study also reports nMSE using the same valid, fresh targets and feature–patient weighting; its squared errors give greater weight to large prediction errors.Source: §4.1.2 · pp. 10–11

4.1.3 Reference-model comparison

• In all eight cohorts, the semantic latent model achieves lower overall error than the best structured candidate in the earlier centralized comparison.That comparison uses the original model configuration and patient-macro metric, unlike the federated experiments’ current model and feature–patient metric.Source: §4.1.3 · p. 11, Table 4 · p. 11, Figure 5 · p. 11

4.2 Characterizing trajectory heterogeneity (RQ1)

• The study characterizes trajectory heterogeneity through state-distribution shifts, unequal temporal context, and changes in which severity groups supply long-window examples.It compares local with pooled access across SOFA/OASIS partitions and measures windows reached under sampled participation.Source: §4.2 · p. 11

4.2.1 Temporal context under local ownership

• Local ownership reduces mean visible history from approximately 14.8 pooled positions to 12.94 and 12.14 for SOFA-30 and SOFA-50, and 11.89 and 10.44 for OASIS-30 and OASIS-50.Between 28.37% and 61.34% of anchors lose at least one history position; H32 retention also falls with horizon, reaching 21.71% for OASIS-50.Source: §4.2.1 · pp. 11–12, Table 5 · p. 12, Figure 6 · p. 12

4.2.2 Long-window support changes the severity composition

• In five-band OASIS, the highest-severity group supplies 31.16% of anchors but 64.26% of complete local H32 windows, while the 31–35 group supplies 21.94% of anchors but 8.03% of H32 windows.The highest group retains 39.73% of H32 windows, versus 7.60% for the 31–35 group; SOFA shows a similar concentration.Source: §4.2.2 · p. 13, Figure 7 · p. 13

4.2.3 Coverage after sampled participation

• Only 21.36%, 15.76%, 14.80%, and 7.55% of pooled-available H32 windows have a locally complete anchor visited during training in SOFA-30, SOFA-50, OASIS-30, and OASIS-50, respectively.The coverage decomposition separates locally complete windows visited during participation from unvisited locally complete windows and windows truncated by ownership.Source: §4.2.3 · pp. 13–14, Figure 8 · p. 14

4.3.1 Finer severity partitions and forecasting error

• FedAvg error increases in 15 of 16 finer-partition comparisons, alongside mean H32-retention drops of 11.17 percentage points for SOFA and 14.17 for OASIS.The paired configurations hold eligible-anchor counts, pooled H32-window counts, pooled mean history length, and centralized references constant.Source: §4.3.1 · p. 15, Table 7 · p. 15
• The finer-partition error penalty is largest at short horizons, falling from 1.70% at H1 to 0.73% at H32 under SOFA and from 4.26% to 2.11% under OASIS.The absolute penalty also falls across horizons, and the gap is already visible at H1 despite complete future-window retention.Source: §4.3.1 · p. 15

4.3.2 Existing federated methods under the benchmark protocol

• Across paired configurations, FedProx lowers mean nMAE by 0.56%, while FedAvgM, FedAdam, SCAFFOLD, and FedMuon raise it by 13.04%, 19.24%, 22.21%, and 26.78%.FedProx improves by more than 1% in seven of 32 settings; its other 25 changes are within ±1%.Source: §4.3.2 · pp. 15–16, Figure 10 · p. 16, Table 8 · p. 17

4.3.3 Rollout horizon and centralized references

• FedProx’s mean paired nMAE gain falls from 0.989% at H1 to near zero at H32; at H32, 16 settings improve and 16 worsen.Fourteen H1 settings improve by more than 1%, compared with one at H32.Source: §4.3.3 · p. 17, Figure 13 · p. 18, Table 9 · p. 18

4.4.1 Update scales under the same temporal partition

• FedAvg’s round-5 update norm is 5.74% of its round-1 norm, while FedAvgM and FedAdam retain 131.09% and 128.36%; high errors occur at both update-scale extremes.FedAvgM’s and FedAdam’s round-5 ratios are 23.06-fold and 18.41-fold the FedAvg norm, respectively.Source: §4.4.1 · pp. 19–20, Figure 15 · p. 20

4.4.2 Mechanism activation across trajectory partitions

• FedExP’s extrapolation factor equals one in 119 of 160 rounds (74.4%); all 40 OASIS-30 rounds are at that bound.In ten runs, the factor is one every round and final nMAE exactly matches paired FedAvg; in the other 22 runs it activates at least once, with final changes within ±1%.Source: §4.4.2 · p. 20, Figure 16 · p. 21, Table 12 · p. 21

4.4 Algorithm behavior under fixed trajectory partitions (RQ3)

• Trajectory heterogeneity elicits distinct algorithm responses: update persistence varies, FedExP often leaves extrapolation inactive, and FedVARP’s degradation depends strongly on partition.The analyses examine update magnitude, extrapolation activation, and cached-update sensitivity.Source: §4.4 · p. 19, §4.4.3 · p. 21

4.4.3 Partition sensitivity of cached updates

• FedVARP’s mean paired degradation is 32.89% under SOFA-30 and 5.15% under SOFA-50, and 46.70% under OASIS-30 and 3.12% under OASIS-50.FedVARP reuses stored updates from previously participating clients, so a round can combine updates from different client mixtures and global-model states.Source: §4.4.3 · p. 21, Figure 17 · p. 22

4.5 Discussion and research opportunities

• The discussion identifies temporal access, participation coverage, update behavior, and horizon-specific prediction as complementary dimensions for evaluating federated world models.It motivates methods that preserve temporal relationships across ownership boundaries and track which trajectory segments server-side memory represents.Source: §4.5 · pp. 22–23
• Only 7.55%–21.36% of pooled-available H32 windows have a locally complete anchor visited during training, and accessible examples also shift in severity composition.Finer partitions reduce retention in every cohort and accompany higher FedAvg error in 15 of 16 paired comparisons.Source: §4.5 · p. 22
• The benchmark uses credentialed-access MIMIC-IV data; construction specifications, field dictionaries, protocols, aggregate results, and analysis scripts support reconstruction and comparison.Source: §4.5 · p. 23

5 Conclusion

• The results show limited long-window coverage, small horizon-dependent gains, and partition-dependent algorithm behavior, motivating methods that jointly account for temporal access, participation, and trajectory-specific updates.Source: §5 · p. 23