How to handle hidden state reset？

[Babylonehy]

1. I have a tensor like BLD input to Mamba, the output is also BLD. This is for training. I wanna to know every new Batch is given, does the hidden stats reset?
2. for inference, I pass throught 1 image at once, how to keep the hidden state until the end of sequence. for new seq, how to manual reset hidden states?

[gkianfar]

@Babylonehy I have the same question about how the hidden stats reset is handled. Have you found any answers?

[Hprairie]

You can pass it in as different samples in a batch and the hidden state for each one will be kept different. You can think of it as essentially just doing a different scan for each sample in B for BLD, starting with a hidden state initialized to 0 for each of the samples in the batch.

[Babylonehy]

You can pass it in as different samples in a batch and the hidden state for each one will be kept different. You can think of it as essentially just doing a different scan for each sample in B for BLD, starting with a hidden state initialized to 0 for each of the samples in the batch.

The training is fine, but the issue is that during online inference with a video sequence, I can’t obtain the full data for dimension B; I only have 1. However, I still want the hidden state to be passed across different frames within the same video sequence.

[Hprairie]

They have an inference mode, which will cache the hidden state for inference. I would take a look at InferenceParams, which will call this function in Mamba.