Update docs

README.md

@@ -81,12 +81,12 @@ Group Relative Policy Optimization (GRPO) was applied on top of the SFT checkpoi
 - **LoRA rank:** 256 (RSLoRA with α=16)
 - **Trainable parameters:** 239M of 1.2B (20%)
 - **Generations per prompt:** 4
-- **Batch size:** 24
-- **Learning rate:** 2e-6 with cosine decay
+- **Batch size:** 16
+- **Learning rate:** 5e-6 with cosine decay
 - **Warmup:** 3% of training steps
 - **Optimizer:** AdamW (8-bit)
 
-The reward function combined five components: Token Error Rate (TER), repetition penalty, face marker accuracy, line structure accuracy, and cuneiform character ratio. The adapter was merged back into the base model at 16-bit precision.
+The reward function combined five components: weighted Token Error Rate using glyph visual similarity and curriculum learning, length deviation penalty, repetition penalty, line structure accuracy, and cuneiform character ratio. The adapter was merged back into the base model at 16-bit precision.
 
 ![grpo-reward](./assets/grpo-reward.png)
 

STORY.md

@@ -58,12 +58,12 @@ For GRPO, I switched from full fine-tuning to LoRA to reduce memory pressure dur
 - **LoRA alpha:** 16 (using RSLoRA scaling: \\(\alpha = \sqrt{r}\\))
 - **Trainable parameters:** 239M of 1.2B (20%)
 - **Generations per prompt:** 4
-- **Batch size:** 24
-- **Learning rate:** 2e-6 with cosine decay (10× lower than SFT)
+- **Batch size:** 16
+- **Learning rate:** 5e-6 with cosine decay
 - **Warmup:** 3% of training steps
 - **Optimizer:** AdamW (8-bit), \\(\beta_1 = 0.9\\), \\(\beta_2 = 0.99\\)
 - **Weight decay:** 0.03
-- **Max gradient norm:** 0.8
+- **Max gradient norm:** 1.5
 - **Loss type:** DR-GRPO
 
 I made three significant mistakes along the way that taught me some lessons. I've included them here in hopes that others can avoid making these same mistakes.
@@ -106,27 +106,49 @@ The initial design tried to encourage correct "shape" (line count, characters pe
 
 I simplified to five reward functions that aligned with how we actually judge transcription quality. All use smooth mappings of the form \\((1 - r) / (1 + r)\\) to avoid sharp cliffs and provide continuous gradients.
 
-**TER Reward** is the primary accuracy signal, based on Token Error Rate (TER), which measures edit distance normalized by ground truth length:
+#### Visual Similarity for Cuneiform Glyphs
 
-$$\text{TER} = \frac{S + D + I}{N}$$
+A major challenge in cuneiform OCR is that many signs are visually similar, differing by only one or two small wedge marks. Standard edit distance treats all substitution errors equally, which doesn't reflect the reality that confusing 𒐉 with 𒀀 (nearly identical shapes) is more forgivable than confusing 𒐉 with 𒀭 (completely different).
 
-where \\(S\\) is substitutions, \\(D\\) is deletions, \\(I\\) is insertions, and \\(N\\) is the number of tokens in the ground truth. The reward transforms this into:
+To address this, I built a visual similarity matrix using the Dice coefficient. For each cuneiform token in the vocabulary, I rendered it as a 64×64 binary image using the NotoSansCuneiform font, then computed pairwise similarities:
 
-$$R_{\text{ter}} = 3 \cdot \frac{1 - 2 \cdot \text{TER}}{1 + 2 \cdot \text{TER}} \in [-3, 3]$$
+$$\text{Dice}(A, B) = \frac{2 |A \cap B|}{|A| + |B|}$$
+
+where \\(A\\) and \\(B\\) are binary masks of the rendered glyphs, and \\(|\cdot|\\) denotes pixel count. The Dice coefficient ranges from 0 (completely different) to 1 (identical). This precomputed matrix of 424,452 pairwise similarities is cached to disk.
+
+These similarities are used in a weighted Levenshtein distance where substitution costs are reduced for visually similar tokens:
+
+$$\text{cost}_{\text{sub}}(a, b) = 1.0 - (\text{Dice}(a, b) \times 0.7) \in [0.3, 1.0]$$
+
+Insertion and deletion costs remain 1.0. This gives partial credit for "close" mistakes, which is crucial for stable reinforcement learning when the model is still learning to distinguish similar signs.
+
+**Annealed Weighted TER Reward** is the primary accuracy signal. It uses the weighted Levenshtein distance described above, normalized by ground truth length to produce Token Error Rate (TER):
+
+$$\text{TER}_{\text{weighted}} = \frac{\text{weighted\_distance}(\text{pred}, \text{truth})}{N}$$
+
+where \\(N\\) is the number of tokens in the ground truth. To implement curriculum learning, the similarity scaling is annealed over training:
+
+$$\text{sharpness} = 0.3 + 0.7 \times \frac{\text{step}}{\text{total\_steps}}$$
+
+$$\text{similarity}_{\text{scaled}}(a, b) = \text{Dice}(a, b) \times (1.0 - \text{sharpness})$$
+
+Early in training (sharpness = 0.3), the model receives generous partial credit for visually similar substitutions. Late in training (sharpness = 1.0), the similarity scaling approaches zero, and the reward converges to standard TER. The reward transforms TER into:
+
+$$R_{\text{ter}} = 3 \cdot \frac{1 - 2 \cdot \text{TER}_{\text{weighted}}}{1 + 2 \cdot \text{TER}_{\text{weighted}}} \in [-3, 3]$$
 
 At 0% error this gives +3, at 50% error it gives 0, and it approaches -3 as errors increase. This dominates the reward signal because accurate transcription is the ultimate goal.
 
-**Repetition Reward** penalizes excess token repetitions beyond what the ground truth contains:
+**Length Reward** penalizes length deviation from the expected output:
 
-$$R_{\text{rep}} = \frac{-r}{1 + r} \in [-1, 0], \quad r = \frac{\max(0, \text{pred\_reps} - \text{truth\_reps})}{N}$$
+$$R_{\text{len}} = \frac{-0.5 \cdot d}{1 + d} \in [-0.5, 0], \quad d = \left| 1 - \frac{|\text{pred}|}{N} \right|$$
 
-where \\(N\\) is the ground truth length. This specifically targets the repetition loops that plagued earlier training runs without penalizing legitimate repeated characters that appear in the ground truth.
+where \\(N\\) is the ground truth length. This works with the weighted TER reward to ensure symmetric penalties for insertions and deletions, since the weighted edit distance alone is normalized by truth length.
 
-**Faces Reward** captures correct identification of tablet face markers (@obverse, @reverse, etc.):
+**Repetition Reward** penalizes excess token repetitions beyond what the ground truth contains:
 
-$$R_{\text{faces}} = 0.5 \cdot \frac{1 - r}{1 + r} \in [-0.5, 0.5], \quad r = \frac{e}{\max(1, n)}$$
+$$R_{\text{rep}} = \frac{-r}{1 + r} \in [-1, 0], \quad r = \frac{\max(0, \text{pred\_reps} - \text{truth\_reps})}{N}$$
 
-where \\(e\\) is the count of incorrect face markers (missing, duplicated, or unexpected) and \\(n\\) is the expected count.
+where \\(N\\) is the ground truth length. This specifically targets the repetition loops that plagued earlier training runs without penalizing legitimate repeated characters that appear in the ground truth.
 
 **Lines Reward** captures structural correctness by measuring deviation from expected line count and line lengths:
 
@@ -136,15 +158,15 @@ where \\(d\\) is the average deviation across line count and per-line character
 
 **Character Reward** encourages using cuneiform script over other characters:
 
-$$R_{\text{char}} = 0.2 \cdot \frac{c}{c + u} - 0.1 \in [-0.1, 0.1]$$
+$$R_{\text{char}} = \frac{c}{c + u} - 0.5 \in [-0.5, 0.5]$$
 
-where \\(c\\) is the count of cuneiform characters and \\(u\\) is unwanted characters (excluding whitespace, punctuation, and face markers).
+where \\(c\\) is the count of cuneiform characters (U+12000–U+1254F) and \\(u\\) is unwanted characters (excluding whitespace, punctuation, structural markers, and face markers).
 
 **Total Reward**
 
-$$R_{\text{total}} = R_{\text{ter}} + R_{\text{rep}} + R_{\text{faces}} + R_{\text{lines}} + R_{\text{char}} \in [-4.85, 3.85]$$
+$$R_{\text{total}} = R_{\text{ter}} + R_{\text{len}} + R_{\text{rep}} + R_{\text{lines}} + R_{\text{char}} \in [-4.25, 3.75]$$
 
-The TER reward dominates by design, with the other components providing shaping signals that help the model avoid common failure modes during early training.
+The weighted TER reward dominates by design, with curriculum learning gradually increasing its strictness. The other components provide shaping signals that help the model avoid common failure modes during training.
 
 With this design, training became stable and outputs became coherent.