Residual head connection

[contentis]

Hello World!

Upfront disclaimer: I'm no LLM researcher.

Looking at the additional heads, I'm wondering if the model could benefit from having a residual connection from head N to N+1. Given that token N+1 strongly depends on token N, I expect the accuracy to improve, especially for an increasing number of heads.

In its easiest form:

medusa_logits.append(self.base_model.lm_head(hiden_states))
for i in range(self.medusa):
     medusa_logits.append(self.medusa_head[i](medusa_logits[i] + hidden_states))

[leeyeehoo]

Agree. We have ongoing discussions inside the team about this. However, it might take some time for us to figure out a better structure. You can check the discussion in the closed PR