Reproduction question: TTT update path collapses during long-context inference under public recipe

[Android-Login-Analysis]

Hi, thanks for open-sourcing the In-Place TTT codebase.

We are trying to reproduce a small-scale long-context training run using the public code and public Qwen3 checkpoints. We are not using short-text data, and our current setup uses 32K-token long-document sequences.

We ran several diagnostic attempts and consistently found that the base model remains healthy when ttt_mode=False, but the TTT update path collapses during long-context inference when ttt_mode=True.

Summary of our setup:

• Base model: Qwen3-4B-Base
• devices: 2*H100 PCIe 80G
• Sequence length: 32768
• Data: long-document 32K-token sequences
• Training: full-parameter training
• Optimizer: AdamW with FP32 parameters and FP32 Adam moments
• Forward: BF16 autocast
• Loss: fused CE, verified numerically against standard torch CE
• ttt_mode=True
• ttt_layers: [0, 6, 12, 18, 24, 30, 36]
  • On Qwen3-4B / Qwen3-8B, this effectively activates 6 layers because layer 36 does not exist.

At around step 500, we observe:

• ttt_conv grows from zero, but remains small.
• ttt_proj norm appears almost unchanged from initialization.
• The base model remains healthy if the checkpoint is loaded with ttt_mode=False.
• With ttt_mode=True, long-context inference collapses: NIAH-style probes produce empty / repeated / newline outputs.
• The collapse appears to come from accumulated present_w drift over multiple TTT chunks, not from base-weight corruption.

We would appreciate guidance on whether this matches a known failure mode.

Specific questions:

1. Should ttt_proj receive a separate learning rate or warmup schedule?
2. In your successful runs, how quickly do ttt_conv and ttt_proj norms move during the first 500–1000 steps?
3. Do you have an expected training-loss curve for a small public reproduction run?
4. Is PG19-style 32K long-document training sufficient, or is ProLong 64K strongly required?
5. Is there a recommended sequence-length curriculum, e.g. 32K → 64K → 128K?
6. Is ttt_layers=[0,6,12,18,24,30,36] intended for 36-layer Qwen models, even though layer index 36 is out of range?
7. Are there any required parameter groups, optimizer settings, or initialization details not reflected in the public config?
8. Is there a minimal known-good public setup, such as Qwen3-4B + ProLong subset + exact command, that should show non-collapsing TTT behavior within 500–1000 steps?

We are happy to provide more exact logs if useful. Our main goal is to determine whether the public recipe is sufficient to reproduce a functional TTT update path on public models and long-context data.

Thanks!

[GuhFeng]

Hi, thank you so much for the detailed and well-structured write-up!

Here are our answers to each of your questions — hope they help!

Q1: Should ttt_proj receive a separate learning rate or warmup schedule?
No, ttt_proj does not need a separate learning rate or warmup schedule — the default settings should be fine.

Q2 & Q3: Expected training-loss behavior in the first 500–1000 steps?
What you're observing sounds roughly in line with expectations. In the first ~500 steps, the training loss with TTT enabled is typically about 0.02–0.05 higher than the baseline (ttt_mode=False). This gap gradually narrows and should converge to around 0.002–0.005 by step 3000 and beyond. So the behavior at step 500 alone may not indicate a fundamental problem — it likely just needs more training time. 500 steps is still quite early for the TTT path to fully stabilize.

Q4 & Q5: Is PG19 sufficient, or is ProLong required? Is there a sequence-length curriculum?
We would recommend switching to the ProLong dataset with the mixing ratios described in the ProLong paper. PG19 is composed of pre-1919 novels, so the data is quite narrow and outdated — this could very well be a contributing factor to the issues you're seeing. As for sequence length, there is no strict curriculum required — you can directly use either 32K or 64K without needing to progressively scale up. Both should work fine.

Q6: Is ttt_layers=[0,6,12,18,24,30,36] intended for 36-layer models?
Yes — layer index 36 will simply be automatically ignored when it is out of range, so this is the expected behavior. On Qwen3-4B (36 layers), 6 TTT layers are effectively activated, which is correct.

Q7: Any hidden optimizer settings or initialization details?
No, there are no additional required settings beyond what is already in the public config.

Q8: Tips for faster stabilization
If you'd like to speed things up a bit during experimentation, you could try reducing the TTT learning rate (e.g., to 0.1 or 0.3). In our experience, a smaller ttt_lr (e.g., 0.3) tends to stabilize more quickly, though the final long-context performance may be somewhat lower compared to using a higher ttt_lr with more training steps. It's a trade-off worth exploring depending on your goals.

In summary: Based on your description, we believe the collapse you're seeing is likely caused by two factors combined:

1. Data: PG19 is a relatively narrow and outdated corpus (pre-1919 novels). We recommend switching to the ProLong dataset with the recommended mixing ratios, which provides much more diverse and representative long-context training signal.
2. Insufficient training steps: 500 steps is still very early for the TTT update path. The loss gap at this stage (0.02–0.05 above baseline) is expected. The TTT parameters typically need around 3000+ steps to properly converge, at which point the gap should narrow to 0.002–0.005. Evaluating long-context inference quality at step 500 is likely premature.

We'd suggest addressing both — switching the data and training for more steps — and we expect the TTT path should stabilize nicely. If you'd like quicker feedback loops during experimentation, lowering ttt_lr as mentioned above can also help.

Please don't hesitate to follow up if you run into any further issues or have more questions — we're happy to help. Thanks again for your interest in the project!