A question about Figure 6 (left)

[tyxqc]

The embedding and lm-head of Qwen3-4B share the same parameter, why there is still such gap between the distribution of Input and Output without alignment.

[xiyuanyang45]

Hi @tyxqc,

Thanks for the question! Qwen3-4B uses tied input/output embeddings, meaning the LM head weight is tied to the token embedding matrix (typically $W_{\text{out}} = W_{\text{in}}^{\top}$, up to conventions). In standard text generation, the model’s next-step input is always obtained by sampling/choosing a token and then looking up its embedding $e_{t+1}$ from $W_{\text{in}}$. So the inputs stay on the learned token-embedding manifold, regardless of how the internal hidden states $h_t$ are distributed.

In LatentMAS, we bypass the discrete token decode→re-embed loop and instead generate “hidden thoughts” autoregressively by conditioning directly on previous hidden representations. Since these hidden vectors are contextual features produced by the transformer (not embedding-table rows), their statistics can drift away from the token-embedding distribution. This leads to a mismatch between the input embedding distribution and the generated hidden-thought distribution.

We hope our explanation here could help clarify your question. Thanks!

[tyxqc]

Thank you very much for your reply! I have one more question based on your reply. If $W_{out} = W_{in}^{T}$ , then the Equation 6 in Appendix A.1 becomes $W_{out}^{T}W_{out}W_{a} - W_{out}^{T}W_{out}^{T}=0$. But $W_{out}^{T}W_{out}^{T}$ is illegal. So I think we should use $W_{out} = W_{in}$ in Qwen3-4B for practical computing. Then the Equation 6 in Appendix A.1 becomes $W_{out}^{T}W_{out}W_{a} - W_{out}^{T}W_{out}=0$ and $W_{a}$ should be identity matrix, theoretically. Then $W_{a}$ loss the realignment ability for Qwen3-4B. But why LatentMAS still effects in Qwen3-4B?