Final Output Layer (LM Head) #189

Implement the final linear layer that produces logits for each token in the vocabulary. BitNet ties the word embedding matrix with this output layer (the weights are shared)
huggingface.co.

This means you can use the same embedding weight matrix (from step 5) transposed as the projection matrix to go from hidden state to vocab logits. For each position (typically one focuses on the last position for generation), apply a BitLinear using the transposed embedding weights: input is the final hidden state vector (8-bit) and weights are the quantized embedding vectors (ternary) for each vocab token. This yields a vector of raw logits of size equal to the vocabulary (128k). No bias is used in this output layer. If memory is a concern, you can compute this dot-product for only the last token’s hidden state when generating step-by-step, rather than for all 4096 positions at once. The result is a set of logits for the next token prediction.