onóna: Twin Embeddings with Tied Unembeddings

onóna. Twin Embeddings with Tied Unembeddings. By Cheng Luo and six others, from Independent Researcher, Michigan State University, and four other institutions. Published on the seventh of November.

Weight tying is a familiar economy in language models. The same table that turns tokens into vectors on the way in is reused to score the vocabulary on the way out, and that sharing cuts the parameter count. The price is on the output side. A tied head of the usual kind is limited to a single softmax, one distribution drawn from one table, and that ceiling on expressiveness comes bundled with the saving.

onóna keeps the sharing and lifts the ceiling by using two tables instead of one. Both serve as input embeddings and both serve as output heads. Going in, the Transformer receives the average of the two embedding tables. Coming out, the head combines the two softmax distributions defined by those very same tables. Input and output still share weights; there are simply two sets of them now, and the output becomes a mixture rather than a single softmax.

How much each table counts in the mixture is decided by a gate, and the gate looks at one thing only: the current input token. Because it depends on nothing else, its values can be precomputed for inference. The Transformer body itself is left unchanged. All the new machinery sits at the two ends of the network.

The evaluation spans four architecture families: Dense, a Gated DeltaNet hybrid, mixture-of-experts, and Attention Residual. Each is tested at three scales, from one hundred forty-two million parameters, through three hundred thirty million, up to one billion. With token budgets matched, onóna lowers final validation loss and external perplexity in eleven of twelve comparisons. It raises average accuracy across eight downstream tasks in ten of twelve. The single exception on loss is the one billion parameter mixture-of-experts model.

On the ledger, the addition costs one extra embedding table and one extra vocabulary projection. In exchange, a tied model gets a mixed output distribution in place of a lone softmax, a gate that can be computed ahead of time, and a Transformer body nobody had to touch.