onóna Uses Twin Embeddings to Expand Language Model Output Expressiveness

Summary:onóna introduces a twin-embedding parameterization that retains input-output weight sharing while moving beyond a single-softmax output head. The Transformer receives the average of two embedding tables, while the output head combines the Softmax distributions produced by those same tables. A lightweight gate conditioned only on the current input token selects the mixture weight and can be precomputed for inference, leaving the Transformer body unchanged. The method was evaluated at 142M, 330M,

---

Core Idea: Escaping the Limits of a Single Softmax

Language models often tie the input embedding table to the output projection to reduce parameter count and memory usage. This saves space, but the resulting output head is constrained by a single Softmax distribution. onóna introduces a twin-table parameterization that preserves input-output weight sharing while increasing output expressiveness.

The Transformer receives the average of the two embedding tables:

$$x=\frac{E_1[t]+E_2[t]}{2}$$

The output head separately computes two vocabulary distributions, $p_1(y\mid t)$ and $p_2(y\mid t)$, and mixes them using a gate conditioned only on the current input token $t$:

$$p(y\mid t)=\alpha(t)p_1(y\mid t)+(1-\alpha(t))p_2(y\mid t)$$

This is not the addition of two independent output heads. The gate does not inspect the hidden state, and the Transformer body remains unchanged. Its coefficients can be precomputed for inference, avoiding an additional network call, recurrent control flow, or changes to the KV Cache path.

Parameterization and Computation

For a vocabulary of size $V$ and hidden dimension $d$, the main overhead is a second $V\times d$ embedding matrix and one additional vocabulary projection. With a 50K-token vocabulary and $d=4096$, the extra matrix occupies roughly 400 MiB in FP16. This is a vocabulary-centric cost rather than an increase in the depth or width of the main network.

During training, both Softmax distributions can be computed from the same set of vocabulary logits and mixed using the token-specific coefficient. During inference, the mixed distribution can be materialized, or two projected outputs can be cached and combined. The memory impact depends on whether the implementation reuses vocabulary projections and whether it materializes full probability tensors for both tables.

The method's main benefit is structural: the input sequence still consumes the shared average of the two tables, while the output side gains two independently learnable vocabulary geometries. It targets the expressiveness limitation of a tied Softmax without replacing the Transformer backbone.

Cross-Architecture Evaluation

The paper evaluates onóna at 142M, 330M, and 1B parameter scales across Dense, Gated DeltaNet hybrid, mixture-of-experts, and Attention Residual architectures under matched token budgets.

• Final validation loss and external perplexity improve in 11 of 12 comparisons.
• Average accuracy across eight downstream tasks improves in 10 of 12 comparisons.
• The only loss exception is the 1B MoE model.

The results suggest that the method is not tied to a single Transformer implementation and can transfer across state-space hybrids, sparse experts, and residual architectures. The 1B MoE exception also shows that the benefit is not guaranteed to grow monotonically with scale. Expert routing, capacity limits, and shared embeddings may alter optimization dynamics, so results from other configurations should not be extrapolated without testing.

Engineering Trade-offs

The dual-table design is not free. Standard tying avoids duplicating vocabulary-related parameters; onóna deliberately retains two such tables to expand the output distribution.

The principal deployment concerns are:

1. Vocabulary projection cost: Each token may require reading two $V\times d$ matrices or performing two matrix multiplications unless projections are fused.
2. Memory bandwidth: Gate coefficients are cheap to precompute, but both vocabulary tables must remain resident. In long-context serving, model-weight bandwidth may matter more than gate computation.
3. Training stability: The two Softmax branches may develop different scales and gradient patterns. Useful diagnostics include branch collapse, a near-constant gate, and domination by one table.

For models with moderate vocabularies, the additional weight footprint is manageable. For extremely large vocabularies or constrained edge deployments, it may outweigh the benefit of tying. The abstract does not report end-to-end throughput, time to first token, training-step time, or fused-kernel performance.

How to Evaluate the Benchmarks

The results across 11 of 12 loss comparisons and 10 of 12 downstream comparisons are more persuasive than a single benchmark win. However, the disclosed summary does not provide effect sizes, confidence intervals, random-seed variance, detailed baseline configurations, or the models used for external perplexity.

Averaged downstream accuracy can also conceal task-level regressions. Independent replication should report per-task accuracy, validation loss, added parameters and FLOPs, wall-clock training time, inference throughput, and matched-budget comparisons against standard tying and untied single-table output heads.

Practical Adoption Path

onóna is most relevant when output expressiveness matters but the input-output sharing benefit remains desirable, such as multilingual vocabularies, code-heavy token distributions, or domains with strongly divergent subword statistics. A practical integration can begin by:

• Initializing the second table as a perturbed copy of the first rather than fully random weights;
• Regularizing the gate against early collapse or near-constant behavior;
• Computing numerically stable normalizations for both vocabulary branches;
• Precomputing token-to-gate coefficients as a lookup table;
• Fusing or chunking vocabulary projections to avoid repeated large-table reads;
• Comparing tied, untied, and onóna variants under identical token counts, training steps, and data order.

For production systems, the decision should not rest on perplexity alone. Teams should calculate additional parameter storage, inference bandwidth, training cost, and throughput. Standard tying remains preferable when parameter economy is the primary objective; onóna offers a more targeted alternative when the tied Softmax is the binding constraint.

Conclusion

onóna combines two Softmax distributions through a token-conditioned gate while leaving the Transformer body unchanged and retaining a form of input-output weight sharing. Its multi-architecture experiments show consistent gains, but the additional embedding table and vocabulary projection introduce meaningful memory and bandwidth costs. Fused implementations and independent performance measurements are needed to confirm whether those costs are justified in production.

---

Source: Hugging Face Daily Papers (2026-11-07)

— END —

Tags: #onóna Follow #Weight Tying Follow #Twin Embeddings Follow #Language Model Architecture Follow