DeepHealthNew

这个目录包含 DeepHealth 的数据准备、数据集、模型、readout 和 loss 代码。当前版本的核心设计是：

疾病序列 stream + 统一的额外信息 token stream

疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 basic_info；其他体检、暴露、生活方式等信息统一整理成 (type, value, value_kind, time) token。

数据准备

运行：

python prepare_data.py

输入文件：

• ukb_data.csv
• field_ids_enriched.csv
• icd10_codes_mod.tsv
• labels.csv

输出文件：

• ukb_event_data.npy

  • 形状为 (N, 3)
  • 每行是 (eid, days, label)
  • 包含疾病、死亡、checkup 事件

• ukb_basic_info.csv

  • index 为 eid
  • 当前只保留 sex

• ukb_other_info.npy

  • 形状为 (M, 5)
  • 每行是：

    (eid, type, value, value_kind, time)
    

  • type=0 预留给 padding
  • value_kind=1 表示连续变量
  • value_kind=2 表示分类变量
  • 缺失值不会生成 token
  • 当前 UKB 额外信息只有一个时间点，所以 time 暂时都是 date_of_assessment

• cate_types.csv

  • 分类变量元信息
  • 字段：

    type,name,n_categories
    

  • ukb_other_info.npy 里的分类 value 是变量内部的局部 id；global category id 在 dataset 中根据实验选择的变量动态计算。

Dataset

dataset.py 提供两个 dataset：

• NextStepHealthDataset

  • 用于 next-token / next-time-point 监督
  • 对应 Delphi2MLoss 和 UniqueTimeSetExponentialLoss

• AllFutureHealthDataset

  • 用于 query-conditioned all-future 监督
  • 对应 ExponentialLoss、WeibullLoss、MixedLoss

为了兼容旧训练入口：

HealthDataset = NextStepHealthDataset
collate_fn = next_step_collate_fn

dataset 会输出：

event_seq
time_seq
padding_mask
sex
other_type
other_value
other_value_kind
other_time

其中 other_type=0 表示 padding，不额外传 other-token mask。

可以通过 extra_info_types 选择纳入哪些额外信息变量：

dataset = NextStepHealthDataset(extra_info_types=[1, 3, 7])

如果不传，则使用全部可用 other-info type。

dataset 会暴露模型初始化需要的元信息：

dataset.n_types
dataset.n_cont_types
dataset.n_categories
dataset.cont_type_ids
dataset.vocab_size

模型

模型主体定义在 models.py，通用网络模块定义在 backbones.py。

BaselineEncoder

BaselineEncoder 编码统一的 other-info token：

other_type        # (B, K)
other_value       # (B, K)
other_value_kind  # (B, K)

它暂时不直接使用 other_time。时间信息保留给后续 CrossAttention，用于建模疾病/query 与 other-info token 的相对时间关系。

连续值使用 TokenAutoDiscretization：

type_id -> continuous type index -> soft bin embedding

分类值使用 dataset 动态计算后的 global category id：

selected type offsets + local category id

CrossAttention

CrossAttention 让 disease-side hidden state 注意到 other-info token：

h_disease  # (B, L, D)
t_disease  # (B, L)
h_token    # (B, K, D)
t_token    # (B, K)

时间信息通过两种方式进入注意力：

• TimeRoPE

  • 使用 query time 和 key time 旋转 q/k
  • 让 q-k 相似度带有时间位置信息

• GaussianRBFTimeBasis

  • 对 t_disease - t_token 做 RBF 编码
  • 投影成每个 attention head 的时间 bias

注意力是时间因果的：

other_info_time <= disease_or_query_time

如果某个 disease/query 位置没有任何可见 other-info token，则该位置保持原 hidden 不变。

DeepHealth

DeepHealth 的统一路径是：

disease-side sequence
-> disease temporal backbone
-> CrossAttention 到 other-info tokens
-> risk head

两种目标模式共用同一套语义：

• next_token

  • h_disease 长度为 L
  • 输出 (B, L, D)

• all_future

  • 在 disease-side 序列末尾拼接一个 query token
  • query token 的时间是 t_query
  • h_disease 长度为 L + 1
  • 输出最后一个 query hidden，形状为 (B, D)

模型初始化示例：

model = DeepHealth(
    vocab_size=dataset.vocab_size,
    n_embd=120,
    n_head=10,
    n_hist_layer=12,
    n_tab_layer=4,
    n_types=dataset.n_types,
    n_cont_types=dataset.n_cont_types,
    n_categories=dataset.n_categories,
    cont_type_ids=dataset.cont_type_ids,
)

Loss

losses.py 中保留：

next-token 监督：

• Delphi2MLoss
• UniqueTimeSetExponentialLoss

all-future / query-conditioned 监督：

• ExponentialLoss
• WeibullLoss
• MixedLoss

UniqueTimeSetExponentialLoss 的 observed term 固定使用 sum reduction，不再暴露旧的 observed_reduction 参数。

训练

当前 train.py 是 next-step 训练入口，使用：

HealthDataset = NextStepHealthDataset

示例：

python train.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --target_mode uts \
  --n_embd 120 \
  --n_head 10 \
  --n_hist_layer 12 \
  --n_tab_layer 4

选择额外信息变量：

python train.py --extra_info_types 1 3 7

如果不传 --extra_info_types，默认使用全部 other-info type。

主要文件

• prepare_data.py

  • UKB 原始数据到模型输入文件的 ETL

• dataset.py

  • next-step 和 all-future dataset
  • 动态选择 other-info type
  • 动态计算 categorical global id

• models.py

  • DeepHealth

• backbones.py

  • TimeRoPE
  • GaussianRBFTimeBasis
  • TemporalAttention
  • GPTBlock
  • TokenAutoDiscretization
  • BaselineEncoder
  • CrossAttention
  • AgeSinusoidalEncoding

• losses.py

  • next-token 和 all-future losses

• readouts.py

  • token readout
  • same-time group readout
  • last-valid readout