Relations between sequence, segment, batch and chunk

[embarrassing-noob-questions]

Thanks for putting this repo together! This is great work!
There is something I think I am completely missing somehow...
I am confused wrt to the relations between sequences, segments, batches and chunks

At inference, suppose I want to feed the model with a long text (>> context window size) and then ask it questions about the whole thing -
So IIUC I do the following:

I initialize the model and the memory to the learned M0
I break the long text into window-size segments, feet it the first segment, and do a forward pass once, letting the memory be updated
Then, without reinitializing the memory I feed it the next segment, and so on
Finally I can ask it questions wrt to facts spread across all segments, and hopefully it retrieves them successfully (NIHS-style)
If I than want to go on to work on another long text then I need to init the memory back to the learned M0 and go again
Is that correct?

If so - does training work the same?
How do you do the next-token prediction on the long context?
Do you do the same? Feed it very long context and update the memory?
Is the memory initialized to the learned M0 value after each long text? When is M0 updated?
How do batch and chunk fit in here?

Finally - after training - do we expect the model to scale to texts longer then those it has seen in training?

I know this is all a noob set of questions but I have been through the papers and through the code and I can't figure it out...

[hustcselwb]

I carefully reviewed this codebase and found that the current implementation is incomplete, not a replica of Google's NL. In particular, the implementation of Titians is quite simple, as is the optimizer. Therefore, your confusion likely stems from the code itself being somewhat incomplete.

[kmccleary3301]

Your understanding is mostly correct. Here is the concise mapping I use:

• sequence: token length processed per forward call
• segment: one long-document slice fed sequentially at inference/eval
• chunk: a training subdivision used for online updates inside one sequence
• batch: number of sequences processed together

Inference on long text:

• I feed segments in order.
• I keep memory state across segments for the same document/task.
• I reset memory state before a new independent document/task.

Training:

• I train with standard next-token loss.
• With train.online_updates=true, each sequence is chunked so later chunks can read effects of earlier writes.
• Memory initialization is learned; runtime state resets per new training context.

I document the exact implementation boundaries in docs/PAPER_COMPLIANCE.md, including where I intentionally use update-pass writes instead of full differentiable-through-write semantics.

If helpful, I can add a README diagram with one concrete sequence/segment/chunk/batch example.

[embarrassing-noob-questions]

Thanks! adding a diagram would be super helpful!