diff --git a/README.md b/README.md
index 2e4fa5d..4ffd62c 100644
--- a/README.md
+++ b/README.md
@@ -73,6 +73,8 @@ HealthDataset = NextStepHealthDataset
 collate_fn = next_step_collate_fn
 ```
 
+all-future 训练入口会显式使用 `AllFutureHealthDataset` 和 `all_future_collate_fn`。
+
 dataset 会输出：
 
 ```python
@@ -227,18 +229,21 @@ all-future / query-conditioned 监督：
   - 使用 `NextStepHealthDataset`
   - `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
   - `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
+  - 当前 next-token 训练只支持 `--dist_mode exponential`
 - `--model_target_mode all_future`
   - 使用 `AllFutureHealthDataset`
   - 不使用 readout，直接对 query hidden 计算风险
   - `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
 
-模型结构组合由 `model_target_mode × time_mode × dist_mode` 决定：
+当前 `train.py` 支持所有已有训练目标定义的组合：
 
-| 维度 | 可选项 |
-| --- | --- |
-| `model_target_mode` | `next_token`, `all_future` |
-| `time_mode` | `relative`, `absolute` |
-| `dist_mode` | `exponential`, `weibull`, `mixed` |
+| 训练模式 | 时间模式 | 分布/监督 | 默认 loss/readout |
+| --- | --- | --- | --- |
+| `next_token` | `relative`, `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` |
+| `next_token` | `relative`, `absolute` | `target_mode=uts`, `dist_mode=exponential` | `UniqueTimeSetExponentialLoss` + `same_time_group_end` |
+| `all_future` | `relative`, `absolute` | `dist_mode=exponential` | `ExponentialLoss`，无 readout |
+| `all_future` | `relative`, `absolute` | `dist_mode=weibull` | `WeibullLoss`，无 readout |
+| `all_future` | `relative`, `absolute` | `dist_mode=mixed` | `MixedLoss`，无 readout |
 
 示例：
 
diff --git a/evaluate_auc.py b/evaluate_auc.py
index 568d129..0c9265a 100644
--- a/evaluate_auc.py
+++ b/evaluate_auc.py
@@ -632,13 +632,15 @@ def compute_logits_for_disease_chunk(
         device.type == "cuda" and use_amp) else torch.float32
     weight = model.risk_head.weight[disease_ids].detach().to(
         device=device, dtype=compute_dtype)
+    bias = model.risk_head.bias[disease_ids].detach().to(
+        device=device, dtype=compute_dtype)
 
     parts: List[np.ndarray] = []
     for start in tqdm(range(0, n, logit_batch_size), desc="Risk-head projection", leave=False, dynamic_ncols=True):
         end = min(start + logit_batch_size, n)
         h = torch.from_numpy(hidden_all[start:end]).to(
             device=device, dtype=compute_dtype, non_blocking=True)
-        logits = torch.matmul(h, weight.t())
+        logits = torch.matmul(h, weight.t()) + bias
         parts.append(logits.float().cpu().numpy().astype(
             np.float32, copy=False))
         del h, logits
diff --git a/train.py b/train.py
index 366f7f2..c0501ab 100644
--- a/train.py
+++ b/train.py
@@ -733,6 +733,11 @@ def normalize_training_config(args: argparse.Namespace) -> None:
         raise ValueError(f"Unknown model_target_mode: {args.model_target_mode}")
     if args.dist_mode not in {"exponential", "weibull", "mixed"}:
         raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
+    if args.model_target_mode == "next_token" and args.dist_mode != "exponential":
+        raise ValueError(
+            "next_token training currently supports dist_mode='exponential' only. "
+            "Use model_target_mode='all_future' for weibull or mixed distributions."
+        )
     if args.all_future_min_history_events < 1:
         raise ValueError("all_future_min_history_events must be >= 1")
     if args.all_future_min_future_events < 1:
@@ -833,7 +838,7 @@ def main():
                         help="Model forward/training mode")
     parser.add_argument("--dist_mode", type=str, default="exponential",
                         choices=["exponential", "weibull", "mixed"],
-                        help="Event-time distribution for model heads and all-future loss")
+                        help="Event-time distribution. next_token requires exponential; all_future supports exponential, weibull, and mixed")
     parser.add_argument("--dropout", type=float, default=0.0,
                         help="Dropout rate")
     parser.add_argument("--extra_info_types_file", type=str, default=None,