Hi, thank you for releasing WAM-Flow — it's an impressive piece of work.
We are reproducing Stage 3. Inference with the released weights matches the paper on our harness (PDMS 90.3 on navtest, Table 5 row 6), and our training port produces bit-identical per-step loss to this repo's train.py (HEAD 2ed9f4d) under the same config and seed. However, training Stage 3 from FUDOKI + navsim_668k.jsonl with the Table 8 recipe (global batch 64, 2 epochs, lr 5e-6) reaches PDMS 66.8 vs 86.7 (Table 5 row 4), and 49.3 vs 76.2 for the "w/o numerical tokenizer" ablation (row 1).
While investigating, we found that row 1 cannot be trained with the released code. With use_quantize: false, training_step fails unconditionally:
train.py L116: x_1 >= vl_chat_processor.num_start_id is evaluated whenever stage == "s2", but num_start_id is only set inside the use_quantize: true branch (L189–199) → AttributeError at the first step.train.py L108–110: masked_select(...).view(b, -1, c) assumes a fixed answer length, which does not hold for text-vocabulary numbers (variable token count per sample) → RuntimeError.
We would greatly appreciate clarification on two points:
Q1. Row 1 training. How was the "w/o numerical tokenizer" ablation (Table 5 row 1, PDMS 76.2) trained, given the two failures above? Was a different training_step used internally?
Q2. Scope of the Stage 3 CE loss (Eq. 8). The released code applies CrossEntropyLoss() with no ignore_index over everything after Assistant: in a fixed 500-token canvas, so ~91% of the supervised positions are pad tokens. In your internal Stage 3 training, (a) were pad tokens included in the CE, and (b) was the loss computed on trajectory answer tokens only, or was any understanding/VQA data (or auxiliary loss) mixed into Stage 3?
Thank you very much for your time!
Hi, thank you for releasing WAM-Flow — it's an impressive piece of work.
We are reproducing Stage 3. Inference with the released weights matches the paper on our harness (PDMS 90.3 on navtest, Table 5 row 6), and our training port produces bit-identical per-step loss to this repo's
train.py(HEAD2ed9f4d) under the same config and seed. However, training Stage 3 from FUDOKI +navsim_668k.jsonlwith the Table 8 recipe (global batch 64, 2 epochs, lr 5e-6) reaches PDMS 66.8 vs 86.7 (Table 5 row 4), and 49.3 vs 76.2 for the "w/o numerical tokenizer" ablation (row 1).While investigating, we found that row 1 cannot be trained with the released code. With
use_quantize: false,training_stepfails unconditionally:train.pyL116:x_1 >= vl_chat_processor.num_start_idis evaluated wheneverstage == "s2", butnum_start_idis only set inside theuse_quantize: truebranch (L189–199) →AttributeErrorat the first step.train.pyL108–110:masked_select(...).view(b, -1, c)assumes a fixed answer length, which does not hold for text-vocabulary numbers (variable token count per sample) →RuntimeError.We would greatly appreciate clarification on two points:
Q1. Row 1 training. How was the "w/o numerical tokenizer" ablation (Table 5 row 1, PDMS 76.2) trained, given the two failures above? Was a different
training_stepused internally?Q2. Scope of the Stage 3 CE loss (Eq. 8). The released code applies
CrossEntropyLoss()with noignore_indexover everything afterAssistant:in a fixed 500-token canvas, so ~91% of the supervised positions are pad tokens. In your internal Stage 3 training, (a) were pad tokens included in the CE, and (b) was the loss computed on trajectory answer tokens only, or was any understanding/VQA data (or auxiliary loss) mixed into Stage 3?Thank you very much for your time!