Question about carry state management across different samples during training

Hello! Thank you for sharing this excellent implementation. The architecture is very elegant and I'm learning a lot from your code.

I have a question about the carry state management during training that I'd like to discuss.

Observation

I noticed a difference in how carry states are initialized between training and evaluation:

Training (in train_batch function):

# Init carry if it is None
if train_state.carry is None:
    with torch.device("cuda"):
        train_state.carry = train_state.model.initial_carry(batch)

Link to source: train_batch function - line 217-220

The carry is initialized only once and preserved across different batches.

Evaluation (in evaluate function):

with torch.device("cuda"):
    carry = train_state.model.initial_carry(batch)

Link to source: evaluate function - line 280-281

The carry is reset for each new batch.

Potential Issue

Looking at the forward pass, carry states are reset when carry.halted=True:

new_inner_carry = self.inner.reset_carry(carry.halted, carry.inner_carry)

Link to source: forward function - line 242

However, halted becomes False when:

• (new_steps < self.config.halt_max_steps) AND (q_halt_logits <= q_continue_logits)
• OR (new_steps < min_halt_steps)

This means that when halted=False, the z_H and z_L states from solving one specific Sudoku instance could be carried over to a completely independent Sudoku instance in the next batch, mixing the reasoning states between unrelated puzzles.

This situation likely occurs frequently during early training when:

• Q-values are not yet well-trained (q_halt ≈ q_continue)
• Exploration probability is high (larger min_halt_steps)

Question

Is this behavior intentional? I'm wondering if samples should maintain their own independent thought processes rather than inheriting states from unrelated problems.

If my understanding of the code is incorrect or if I'm missing something important about how the carry states are managed, I would greatly appreciate any clarification. I want to make sure I'm interpreting the implementation correctly.

Potential Solution

If this is unintended, a simple fix might be to remove the condition in train_batch:

# Always initialize carry for each batch
# if train_state.carry is None:  # <- Comment out this line
with torch.device("cuda"):
    train_state.carry = train_state.model.initial_carry(batch)

Link to source: train_batch function - line 217-220

This would ensure each batch starts with fresh carry states, similar to the evaluation behavior.

Additional Context

I understand this might be related to memory efficiency or a specific training strategy. If this is intentional, I'd really appreciate learning about the rationale behind this design choice!

Thank you again for your great work and for any clarification you can provide. I'm excited to better understand this implementation.

Best regards