Our paper onóna: Twin Embeddings with Tied Unembeddings is now on arXiv. It asks whether a language model can keep the parameter efficiency of tied input and output embeddings while giving its output head more expressive power. The research code is available.
Why change a tied output head?
Weight tying uses the same vocabulary table to turn token IDs into input vectors and to score candidate next tokens. It saves a large matrix. Its output head, however, still produces one softmax distribution from one table. Simply untying the output matrix adds parameters while losing the input–output sharing we wanted to preserve.
onóna keeps the sharing but gives the head two distributions. Let and be vocabulary tables. The model reads their average, , as its input embedding. At output, it computes a softmax from each table and mixes the two distributions using a gate determined by the current input token. The Transformer body is unchanged. Because the gate depends on the token rather than the hidden state, its values can be precomputed for inference.
This is a compact way to add output flexibility. The cost is real: one extra embedding table and one extra vocabulary projection. Exact output scoring still needs both projections; the method does not claim a free inference speedup.
What the experiments show
The paper compares onóna with the standard tied head under matched token budgets across 142M, 330M, and 1B models, including Dense, Gated DeltaNet hybrid, mixture-of-experts, and Attention Residual architectures. Final validation loss and external perplexity improve in 11 of 12 comparisons. The average accuracy over eight downstream tasks improves in 10 of 12. The loss exception is the 1B MoE model.
These counts summarize the reported settings rather than guaranteeing improvement for every model or workload. Figure 1 shows the architecture, the 142M and 330M validation-loss differences, and an early training curve. The figure’s bars represent one seed per setting, so they should not be read as statistical error bars.
When it may be useful
The design is most interesting when a model already uses tied embeddings and output expressiveness matters enough to justify another vocabulary projection. It offers a specific architecture tradeoff: keep input–output sharing and the existing Transformer body, while spending more parameters and output computation on the head. Evaluating the tradeoff on larger models, different tokenizers, and production inference workloads remains open.
Paper: arXiv:2611.10556 · Project: FutureMLS · Code: FutureMLS-Lab/onona
**中文摘要：**onóna 用两张共享的词表构造输入与输出：输入取平均，输出通过由当前 token 决定的门控混合两个 softmax。它保留 Transformer 主体及输入输出权重共享；代价是一张额外词表与一次额外投影。在论文的匹配训练预算实验中，12 组比较有 11 组改善验证损失和外部困惑度。完整论文和代码见上方链接。