diff --git a/notebooks/laya_finetune_typed_decisions_mps.py b/notebooks/laya_finetune_typed_decisions_mps.py
index ff465f4..d2a39e4 100644
--- a/notebooks/laya_finetune_typed_decisions_mps.py
+++ b/notebooks/laya_finetune_typed_decisions_mps.py
@@ -223,6 +223,9 @@ def train(args, model_dir, items_path, device):
         )
         model.head_checkpointing = True
     model.to(device).train()
+    if args.head_only:
+        model.encoder.requires_grad_(False)
+        model.encoder.eval()
 
     all_items = torch.load(items_path, map_location="cpu", weights_only=False)
     order = list(range(len(all_items)))
@@ -236,8 +239,11 @@ def train(args, model_dir, items_path, device):
 
     encoder_params = [p for n, p in model.named_parameters() if "encoder." in n]
     head_params = [p for n, p in model.named_parameters() if "encoder." not in n]
+    groups = [{"params": head_params, "lr": 1e-4}]
+    if not args.head_only:
+        groups.insert(0, {"params": encoder_params, "lr": 2.5e-5})
     optimizer = torch.optim.AdamW(
-        [{"params": encoder_params, "lr": 2.5e-5}, {"params": head_params, "lr": 1e-4}],
+        groups,
         weight_decay=0.01,
     )
     updates = max(1, math.ceil(len(train_items) / args.micro_batch / args.grad_accum) * args.epochs)
@@ -265,7 +271,7 @@ def train(args, model_dir, items_path, device):
             target = target.to(device)
             qtype = qtype.to(device)
 
-            logits, activation = model(ids, attention, positions, mask, qtype)
+            logits, activation = model(ids, attention, positions, mask, qtype, detach_encoder=args.head_only)
             logits = logits.float()
             k = mask.sum(-1, keepdim=True).float()
 
@@ -355,6 +361,7 @@ def main():
     parser.add_argument("--device", choices=["auto", "mps", "cpu"], default="auto")
     parser.add_argument("--force-preprocess", action="store_true")
     parser.add_argument("--no-checkpointing", action="store_true")
+    parser.add_argument("--head-only", action="store_true", help="freeze the encoder; train head/scorer/type_emb/act_head")
     args = parser.parse_args()
 
     if args.model_id != MODEL_ID: