diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py index 997ed3f85..b407f0802 100644 --- a/pyhealth/models/embedding/unified.py +++ b/pyhealth/models/embedding/unified.py @@ -402,9 +402,12 @@ def forward( A dict with keys: * ``"sequence"``, ``(B, S_total, E')`` temporally-sorted events + (content + time + type embeddings) * ``"time"`` , ``(B, S_total)`` timestamps (hours) * ``"mask"`` , ``(B, S_total)`` 1=real event, 0=padding * ``"type_ids"``, ``(B, S_total)`` modality index per event + * ``"token_emb"``, ``(B, S_total, E')`` content-only event embedding + (before time/type are added); the target for masked modeling. """ all_embeddings: list[torch.Tensor] = [] all_times: list[torch.Tensor] = [] @@ -493,4 +496,10 @@ def forward( "time": cat_time, # (B, S_total) "mask": cat_mask, # (B, S_total) "type_ids": cat_types, # (B, S_total) + # Per-event content embedding BEFORE time/type are added (same sort + # order as ``sequence``). Masked-modeling pretrainers should + # reconstruct THIS rather than ``sequence``: the time/type + # components are largely recoverable from event position, so + # including them in the target dilutes the content signal. + "token_emb": cat_emb, # (B, S_total, E') } diff --git a/tests/test_unified_multimodal.py b/tests/test_unified_multimodal.py index 3efc47404..f3b35890d 100644 --- a/tests/test_unified_multimodal.py +++ b/tests/test_unified_multimodal.py @@ -223,6 +223,33 @@ def test_unified_model_code_only(): assert E == 64 +def test_unified_model_token_emb_is_content_only(): + """``token_emb`` is the content embedding BEFORE time/type are added, in the + same temporally-sorted order as ``sequence``; i.e. + ``sequence == token_emb + time_embed(time) + type_embedding(type_ids)``.""" + from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel + + processors, inputs = _make_code_processors_and_inputs() + model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=64) + model.eval() + + out = model(inputs) + + assert "token_emb" in out + assert out["token_emb"].shape == out["sequence"].shape + + # Reconstruct sequence from the content-only token_emb + the position-derived + # time and type embeddings; they must match exactly. + recomposed = ( + out["token_emb"] + + model.time_embed(out["time"]) + + model.type_embedding(out["type_ids"]) + ) + assert torch.allclose(out["sequence"], recomposed, atol=1e-5) + # And token_emb is genuinely distinct from sequence (time/type were added). + assert not torch.allclose(out["token_emb"], out["sequence"]) + + def test_unified_model_rejects_non_temporal(): from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel from pyhealth.processors import SequenceProcessor