Allow extra outputs from `GenerationMixin.generate`

[jood-canva]

Feature request

Hi all, first of all if this feature already exists I apologise!

With the rise of multimodal LLMs if would be great if we could add extra outputs to GenerationMixin.generate results. For instance if we implement a model like Janus from DeepSeek, there are two output heads. One lm_head, and one image_head. The outputs of forward method have extra attributes that can't be passed to the generate results.

I know these multimodal models are not common within this repo so this is pretty bleeding edge, but I'm working on research in this domain and it would be great if we could forward all model outputs to the generate result. Maybe through an attribute like kwarg_outputs in classes like GenerateDecoderOnlyOutput?

Motivation

As far as I understand it's possible to feed the extra output during the autoregressive loop, through prepare_inputs_for_generation and _update_model_kwargs_for_generation where we can forward model outputs to the next forward call.

But when it comes to forward these outputs to the result of generate, it doesn't seem possible? I know the generation mixin is geared towards text generation, but it would be great to be able to forward extra model outputs

Your contribution

Happy to have a try but not sure how big of a PR it would be, especially if it touches the pytorch / tf / flax implementations.

[zucchini-nlp]

Hey @jood-canva !

Can you give an example of which outputs are needed for these models? Indeed, we can overwrite _update_model_kwargs_for_generation to pass over all kwargs from model outputs to the next inputs. Do you mean a possibility to obtain image_hidden_states or similar like

out = model.generate(**inputs, return_dict_in_generate=True)
generated_ids = out.sequences
hiddens = out.image_hidden_states

[jood-canva]

Hey @zucchini-nlp sorry for the delay I was on a long break. Thanks for your reply!

You are right that we would like to access each head's output after generate is done, but my request goes a bit deeper than that.
Let's say you have a model like Janus with two output heads, lm_head and image_head. The model is also autoregressive so when you implement the generate function you need to pass both head outputs as arguments to the next forward call. Like, when it starts generating image tokens you need to pass that as a pixel_values to the next forward. From what I understand the current generate implementation doesn't allow that, it focuses heavily on input_ids only?

What I'm suggesting is maybe extending Output classes so that we can add extra arguments. Right now the generate logic only works well with input_ids but it shouldn't take too much efforts to forward extra kwargs?

Tell me if that still doesn't make sense!

[zucchini-nlp]

You are right, Janus-like models are a special case and the generate() method cannot handle it well, which is why we had to write a whole new generation loop for the model. TBH it doesn't work with both lm and image head, instead the model can do either text or image-only generation. Not sure what type of output is needed in this case

I agree that returning extra kwargs in GenerationOutput would be useful for multimodal models (probably MoE can also benefit) and it will not be hard.