My understanding: The Q-head predicts whether to stop early or not to reduce training times (by not always having to complete the full number of deep supervision steps).
The complication: When working in batches:
- Either, the Q-head is computed over the entire batch and averaged (which could early stopping of reasoning steps when most items in batch are complete but some aren't, meaning that those incomplete samples aren't trained on fully)
- Or, the Q-head is computed over entire batch and samples that are complete are dropped.
The first option has issues obvious as discussed.
The second option is problematic because it's complex to implement and frameworks tend to want fixed sizes for batches.
If we go with just running for the maximum number of deep supervision steps every time to avoid these issues, the Q-head becomes useless (assuming you don't want to use it during inference).
I really appreciate any help anyone can give me in understanding how they get around this, either in the original paper or in practice in more recent research.