onóna: Twin Embeddings with Tied Unembeddings

AI summary

Overview

Research area: Natural Language Processing — language model architecture, specifically input/output weight tying and the expressiveness of the softmax output head.

Technical level: Intermediate. Readers should be comfortable with token embeddings, the softmax output layer, and standard pretraining loss, but the core idea is explained from first principles.

Scope: The paper proposes onóna, a two-table parameterization that keeps embedding–unembedding weight sharing while mixing two softmax distributions at the output, and evaluates it across four architectures at 142M, 330M and 1B parameters.

What This Paper Is About

Most language models reuse a single vocabulary table for both the input embedding and the output "unembedding," which saves a V×d matrix but restricts the head to one softmax. A single softmax head has a log-probability matrix of rank at most d+1, a limitation known as the softmax bottleneck, which can hurt when the next token has several plausible modes. The authors ask whether a model can gain output capacity without giving up input–output sharing, and propose tying two tables to both roles, feeding the decoder their average and mixing their two output distributions with a token-dependent gate.

Key Contributions

1. Twin tied tables for embedding and unembedding. onóna ties both vocabulary tables (A and B) to the output head, feeds their fixed average M = (A + B)/2 to the unchanged Transformer body, and mixes the two softmax distributions defined by those same tables.

2. A gate that reads only the input token. The mixture weight is a function of the current token rather than the hidden state, so all V gate values can be precomputed at serving time and the input tables fold into their average. Exact scoring and greedy decoding still need two vocabulary projections.

3. Evaluation across four architectures and three sizes. Dense, GDN hybrid, MoE and AttnRes models are trained at 142M, 330M and 1B with equal token budgets within each size. Eleven of twelve pairs improve validation loss and both external perplexity benchmarks, and ten improve the eight-task accuracy average.

4. Controls and ablations. The paper compares against an untied head, a K=2 mixture of softmaxes, input-only gating and a deeper baseline, and runs a 2×2 input–output factorial ablation showing neither role of the second table recovers the gain on its own.

Main Findings

• Validation loss at 142M: onóna lowers final validation loss for all four architectures — Dense 2.8142 to 2.7857 (−0.0285 nats), GDN hybrid 2.7600 to 2.7467 (−0.0134), MoE 2.7190 to 2.7074 (−0.0116), and AttnRes 2.7791 to 2.7564 (−0.0226).

• Validation loss at 330M: gains shrink for Dense (−0.0065, 2.5868 to 2.5803), GDN hybrid (−0.0037, 2.5440 to 2.5404) and AttnRes (−0.0119, 2.5533 to 2.5414), but the MoE gain grows from 0.0116 to 0.0144 (2.5183 to 2.5039).

• The 1B MoE exception: onóna lowers loss for 1B Dense (2.4902 to 2.4811), GDN hybrid (2.4686 to 2.4629) and AttnRes (2.4766 to 2.4673), but raises it for 1B MoE (2.4967 to 2.5238, +0.027). That MoE is also worse on LAMBADA (22.6 to 22.8) and WikiText-2 (27.7 to 28.1) while still posting a higher downstream average.

• External perplexity: both benchmarks improve in all eight 142M and 330M pairs — LAMBADA by 0.99–9.18% and WikiText-2 by 0.67–3.62%. The largest LAMBADA reduction is 330M MoE (25.84 to 23.47); the largest WikiText-2 reduction is 142M AttnRes (58.44 to 56.32).

• Both tables are needed: using only A or only B at both input and output raises validation loss in all sixteen 142M and 330M ablations, by 0.188–0.855 nats over the corresponding mixture.

• Downstream accuracy: the eight-task average improves in six of eight 142M/330M pairs. 142M Dense gains 1.38 points (47.46% to 48.84%) and AttnRes gains 0.25, while GDN hybrid and MoE decline slightly (−0.14 and −0.05). All four 330M averages improve by 0.67 to 1.17 points, and all four 1B averages improve by 0.09 to 1.02 points. Only 330M AttnRes improves on all eight tasks.

• Versus other output layers (142M Dense): onóna has the lowest validation loss at 2.7868, 0.0106 nats below MoSS at 2.7974. An untied head reaches 2.8132 and MoS (K=2) 2.8239. MoSS wins on LAMBADA perplexity (46.50 versus 58.38), and the two are level on WikiText-2 (61.36 versus 61.35). onóna also has the highest downstream average at 49.12%, 1.39 points above Standard.

• Untying does not reproduce the gain: an untied head adds the same V×d matrix but raises loss by 0.009 nats at 142M and 0.003 at 330M, relative to the respective Standard baselines.

• Input and output roles are complementary: in the 2×2 grid, input averaging alone lowers loss by 0.00729 nats, the output mixture alone raises it by 0.00716, and together they lower it by 0.02002. The measured interaction is 0.01990 nats.

• Depth is a competitive alternative: a 330M Standard deepened from 24 to 28 layers reaches 2.6258, which is 0.0071 below onóna at 2.6329, using a roughly matched budget.

• Cost: the second table adds 38.6M parameters at 142M, 51.5M at 330M and 103M at 1B. In a 1B Dense timing run on 8× H100, throughput is 443K tokens/s for onóna versus 502K for Standard, a 12% reduction. At 142M Dense the method costs 181.10M parameters and 38.42GB versus 143.65M and 33.35GB for MoSS, at similar throughput (422 versus 428 ktok/s).

Methodology in Plain English

The authors start from the standard tied-table setup: one matrix of shape V×d maps tokens into the model and, at the end, scores every vocabulary entry as the next token. They replace that single matrix with two, A and B, of the same shape.

On the input side, the model simply reads the average of the two tables. That keeps the representation width the same, so the Transformer body is untouched and no input routing is needed.

On the output side, the same hidden state is projected through each table separately, each result is turned into its own probability distribution, and the two distributions are blended. The blend weight comes from a small learned vector applied to the current input token's averaged embedding, passed through a sigmoid. Because only the token identity matters, each vocabulary item gets one learned weight that can be computed once and looked up at serving time.

Training uses ordinary next-token negative log-likelihood. The implementation works in log space with logaddexp, logsigmoid(z) and logsigmoid(-z), which stay finite when the gate saturates; an earlier version that took log σ(z) directly diverged when σ(z) underflowed.

The evaluation trains Qwen3-style decoders from scratch on ClimbMix with GPT-2 BPE and 1,024-token contexts. Four architectures (Dense, GDN hybrid, MoE, AttnRes) are trained at 142M, 330M and 1B. The 142M and 330M models share a 7.18B-token pool and the 1B models use a disjoint 23.9B-token pool. Every run uses a 40 tokens-per-parameter budget based on the Dense Standard parameter count, so variants at each size see the same number of tokens. Comparisons cover validation loss, LAMBADA last-word perplexity, WikiText-2 word perplexity, and eight downstream tasks (zero-shot except MMLU, which uses up to five shots). Ablations isolate extra vocabulary capacity, extra output computation, the input and output roles of the second table, and an alternative use of the same budget in decoder depth.

Why This Matters

The gain comes without changing the Transformer body, which makes it a drop-in modification to existing tied-embedding training recipes. It also reframes the softmax bottleneck: instead of adding more contexts under one table (as a mixture of softmaxes does), onóna adds more tables under one context while keeping every scoring table tied to the input.

Potential real-world applications:

• Small models on constrained hardware, where the V×d table is a large share of parameters and output expressiveness is a known limitation.
• Inference serving for open-ended sampling, where selecting a component by Bernoulli(g) yields an exact sample from the mixture with a single vocabulary projection at temperature 1 without truncation.
• Pretraining pipelines that already use weight tying, since the decoder and its configuration remain unchanged.
• Architecture research on mixture-style heads, where the token-conditioned gate offers a cheap alternative to contextual mixture priors.

Industry relevance: the method targets the pretraining regime that most production language models use, and the reported costs are concrete — an extra V×d table, a d-dimensional gate vector, and a 12% throughput reduction in the 1B Dense timing run. The finding that a deeper Standard model beats onóna at a similar budget at 330M is directly relevant to teams deciding where to spend parameters.

Future Directions

• Multi-seed evaluation. Every arm uses one training seed, and the authors state that differences of a few tenths of a point in downstream accuracy are not reliable without more seeds.
• A resource-matched 1B comparison. The 1B control experiments, including the depth comparison, were not run; the 330M depth result suggests the budget question remains open at larger scale.
• Explaining the 1B MoE reversal. The MoE model is the one loss exception at 1B, with worse validation loss and perplexity despite a higher downstream average, and the paper does not identify the cause.
• Separating input and output contributions more cleanly. The input-only gate control changes both the input rule and the head, so the authors note the remaining gap cannot be attributed to output gating alone.

Target Audience

Researchers and engineers working on language model architecture, particularly those interested in weight tying, output-layer design, and the softmax bottleneck. It is also useful for practitioners training small or mid-sized models who want a low-risk architectural change with measured trade-offs, and for anyone comparing mixture-of-softmaxes-style heads against alternatives under matched token budgets.

Authors’ abstract

Weight tying reduces the parameter cost of language models, but the usual tied output head is limited to a single softmax. We introduce onóna, a two-table parameterization that increases output expressiveness while preserving input-output weight sharing. The Transformer takes the average of two embedding tables as input, and the output head combines the softmax distributions defined by those same tables. A gate conditioned only on the current input token determines the…