Update README.md

README.md

@@ -68,9 +68,9 @@ Input Text (Reasoning Trace)
     ↓
 [Frozen Base LM Encoder]  ← Pre-trained, frozen during training
     ↓
-[Mean Pooling]
+[Final Token (EOS) Pooling]
     ↓
-[Lightweight MLP Head]    ← Only these parameters are trained
+[Lightweight Linear Head]    ← Only these parameters are trained
     ↓
 Scalar Reward Score
 ```
@@ -129,14 +129,14 @@ Each pair contains:
 ### Training Configuration
 
 **Hyperparameters:**
-- **Base Model**: facebook/OPT1.3b
+- **Base Model**: facebook/opt-1.3b
 - **Trainable Parameters**: Scoring head only (~500K-1M params)
 - **Optimizer**: AdamW
-  - Learning rate: 1e-4
+  - Learning rate: 2e-5
   - Betas: (0.9, 0.999)
   - Weight decay: 0.01
-- **Learning Rate Schedule**: Cosine decay with 50-step warmup
-- **Batch Size**: 32 pairs
+- **Learning Rate Schedule**: Linear warmup (50 steps) + constant
+- **Batch Size**: 8 pairs
 - **Gradient Clipping**: Max norm 1.0
 - **Training Steps**: 800
 - **Mixed Precision**: FP16
@@ -249,8 +249,9 @@ def score_trace(trace_text: str) -> float:
     with torch.no_grad():
         # Get base model embeddings
         encoder_outputs = base_model(**inputs)
-        # Pool final token (EOS)
-        pooled = encoder_outputs.last_hidden_state[:, -1, :]
+        # Pool at actual sequence end (accounts for padding)
+        seq_lengths = inputs["attention_mask"].sum(dim=1) - 1
+        pooled = encoder_outputs.last_hidden_state[torch.arange(seq_lengths.size(0)), seq_lengths]
         # Get reward score
         score = scoring_head(pooled).squeeze(-1).cpu().item()