PaperlayerView exact revision on arXiv

Source-linked AI summary

TriPLU: Bypassing the Gate with Direct Trilinear Product FFNs in Tiny Language Models

He Zhang

arXiv:2608.20360v1cs.CLcs.LG

TL;DR

The paper asks whether tiny decoder-only language models benefit from directly multiplying learned feature projections instead of relying on conventional gated FFNs. It evaluates TriPLU, a product-only degree-3 branch, in controlled character-level and Byte-BPE studies. TriPLU improves selected validation and heldout language-modeling metrics, but the evidence remains optimization-sensitive and limited to small low-compute settings.

  • Problem

    The paper asks whether direct product FFNs add value beyond strong gated FFNs for feature interactions in tiny decoder-only language models.

  • Method

    TriPLU replaces the usual gated FFN branch with a product-only degree-3 branch that multiplies three learned projections coordinatewise.

  • Results

    TriPLU improves best validation loss in the TinyStories prefix comparison and validation and heldout BPB in low-learning-rate Byte-BPE comparisons on TinyStories and WikiText-2 raw.

  • Takeaways & Limitations

    Direct product FFNs can improve fixed-budget small-model loss in specific low-compute regimes, with gains aligning with seen co-occurrence slices.

  • Takeaways & Limitations

    The evidence does not establish FLOP-normalized efficiency, scaling behavior, or broad LLM performance, and product branches remain learning-rate sensitive.

Abstract

from arXiv · show

We study whether tiny decoder-only language models benefit from feed-forward layers that directly multiply learned feature projections. TriPLU, a Trilinear Product Linear Unit, replaces the usual gated FFN branch with a product-only degree-3 branch that multiplies three projected streams coordinatewise. In a character-level TinyStories 1M-byte prefix study, TriPLU reaches a mean best validation loss of 1.0637, compared with 1.1017 for closely matched SwiGLU, 1.0780 for a degree-4 product control, and 1.1026 for a degree-2 control. In train-only Byte-BPE experiments, TriPLU also lowers validation and heldout bits per byte on TinyStories and WikiText-2 raw under low-learning-rate settings, with PMI-slice evidence suggesting gains on seen middle- and high-PMI adjacent-token pairs. Constant-learning-rate diagnostics show that product-branch normalization can reduce the high-learning-rate best-checkpoint gap, although final BPB still degrades under hot schedules. The resulting claim is deliberately narrow: direct product FFNs can improve fixed-budget small-model loss in specific low-compute regimes, but the branch is optimization-sensitive and does not establish FLOP-normalized efficiency, scaling behavior, or broad LLM performance.

1 Introduction

The paper tests whether direct products of learned projections can improve tiny decoder-only language models beyond a closely matched SwiGLU baseline. Its controlled study reports stronger character-level TinyStories results while framing the contribution as a low-compute language-modeling study.

  • Motivation: Direct-product FFNs target feature interactions that ordinary Transformer FFNs may learn only through additional width, depth, or data.The motivation is to express multiplicative co-occurrence structure more directly.
  • Research question: SwiGLU is the necessary primary baseline because gated GLU variants have improved Transformer FFNs over ReLU or GELU in prior work.The comparison asks whether a more explicit product branch adds value under matched parameters and training tokens.
  • Method: TriPLU replaces the usual activation-and-gate form with a direct path that multiplies three learned projections coordinatewise.The product branch provides the nonlinearity without applying a unary activation to one stream.
  • Scope: The contribution is a controlled low-compute language-modeling study rather than a new arithmetic-unit proposal.The paper positions direct products as the tested architectural choice, not as a novel general multiplication mechanism.
  • Main result: 1.0637 mean best validation loss gives TriPLU a 0.0380 reduction over closely matched SwiGLU on the TinyStories 1M-byte prefix.The reduction is about 3.4% relative.

2 Related Work

Related work places TriPLU among gated Transformer FFNs, product units, multiplicative architectures, and neural arithmetic modules. The paper distinguishes its contribution by testing one direct-product FFN in a constrained tiny-language-model setting, while recognizing optimization and stability challenges.

  • Transformer FFNs: SwiGLU and GEGLU provide the closest mainstream precedent because GLU-style FFNs can improve over ReLU or GELU FFNs.This prior result motivates SwiGLU as the primary baseline for TriPLU.
  • Direct products: TriPLU replaces SwiGLU’s nonlinear gating emphasis with a product of three learned projections.The comparison isolates a more explicit direct-product bias against a strong gated FFN.
  • Product and arithmetic units: Product units and neural arithmetic modules show that explicit multiplicative structure can model nonlinear interactions but is difficult to train.The paper treats optimization, sign, zero-handling, and stability as relevant challenges rather than introducing a new arithmetic unit.
  • Multiplicative interactions: Gating, attention, hypernetworks, and dynamic layers are established forms of multiplicative interaction, so TriPLU tests a specific placement of multiplication rather than introducing multiplication to neural networks.The contribution is intentionally confined to direct-product FFNs in tiny language models.
  • Evidence standard: Arithmetic diagnostics remain secondary because arithmetic loss is mechanism evidence, whereas language-model loss determines the paper’s claim.This separates evidence about arithmetic behavior from the central language-modeling evaluation.

3 Method

The method replaces the FFN hidden activation with elementwise products of learned projections, focusing on a scaled degree-3 branch. Experiments compare this branch with matched baselines and product-order controls while probing optimization difficulty through negative ablations and scale settings.

  • Model and training: The models are compact decoder-only Transformers trained from scratch with causal attention, residual blocks, layer normalization, and configurable FFNs.The shared character-level setup uses 2 layers, 2 heads, width 96, dropout 0.0, AdamW, learning rate 0.0003, weight decay 0.1, and batch size 32.
  • Direct-product branches: Direct product FFNs project the hidden state into two, three, or four streams, multiply them elementwise, and apply an output projection.TriPLU uses the degree-3 branch, with degree 2 and degree 4 as product-order controls.
  • Scale control: The scalar α is a stability parameter that keeps product-branch activations numerically comparable to standard FFN activations.Character-level runs use fixed gains selected in exploratory screening; later Byte-BPE diagnostics use distinct learnable or normalized scale settings.
  • Negative ablations: Log-domain power units, hard-sign variants, integer-power variants, and attention-side product ablations were weaker and harder to optimize than direct products.These experiments are treated as negative diagnostics rather than headline language-modeling evidence.
  • Matching: Main TinyStories comparisons are close parameter-matched rather than exactly parameter- or FLOP-matched.Discrete branch-width choices leave a 576-parameter spread across the matched variants.

4 Experimental Setup

The experiments combine deterministic TinyStories and WikiText-2 checks with character-level and train-only Byte-BPE protocols. Evaluation uses validation-selected checkpoints, heldout metrics, PMI-based co-occurrence slices, and controlled follow-up runs, while arithmetic exact match remains diagnostic.

  • Datasets: Public experiments use deterministic TinyStories prefixes, with the main character-level setting requesting 1,000,000 training bytes and 200,000 validation bytes.The Byte-BPE extension adds TinyStories and WikiText-2 raw, while arithmetic_lite is secondary.
  • Byte-BPE protocol: Train-only Byte-BPE comparisons use 8-layer, width-256 decoder-only models with four heads, block size 128, vocabulary size 512, and seeds 1–3.The matched SwiGLU hidden size is 768; unnormalized TriPLU uses product branch size 576 with a learnable scale initialized at 5.0.
  • Mechanism slices: PMI slices bin training-observed adjacent-token pairs by PMI and report unseen pairs separately to test whether products help where co-occurrence structure is available.The slices do not identify individual hidden units.
  • Metrics and selection: Primary checkpoints are the lowest-validation-loss checkpoints, with final validation loss retained as a stability diagnostic; Byte-BPE additionally evaluates heldout data.Character-level results remain validation-prefix evidence.
  • Diagnostics: Arithmetic exact match is measured by greedy decoding answers to addition and multiplication prompts but is not used as the central paper claim.This keeps the arithmetic task diagnostic rather than decisive.
  • Controlled follow-up: The corrected follow-up freezes widths, gains, seeds, data prefix, schedule, bias-free projections, and metrics before comparison.The design is intended to make the main comparison controlled and reproducible.

5 Results

Across matched tiny-model experiments, TriPLU improves validation and heldout loss over SwiGLU in low-learning-rate settings, with evidence concentrated in observed co-occurrence contexts. Optimization diagnostics show that direct-product performance is sensitive to branch scale and learning rate, limiting the scope of the result.

  • 5.1 Main TinyStories Prefix Result: TriPLU wins all three matched TinyStories seeds, with a 0.0380 validation-loss advantage over SwiGLU.The paired seed improvements are 0.0322, 0.0259, and 0.0558; formal significance tests are not run with only three seeds.
  • 5.3 Tokenizer-Neutral Heldout BPB: TriPLU improves validation-selected and heldout BPB on TinyStories, with the same direction on WikiText-2 raw.The reported effect sizes are modest, so the emphasis is consistency across matched seeds and datasets.
  • 5.4 Co-occurrence Slices: TriPLU improves several seen adjacent-token PMI bins more than unseen pairs, especially higher-PMI seen bins.The slice evidence is compatible with learned co-occurrence benefits, but omitted token counts and intervals make it directional.
  • 5.5 Optimization Sensitivity and Normalization: Normalization and fixed branch scale improve high-learning-rate best-checkpoint BPB, while final BPB still degrades under hot schedules.Normed TriPLU reaches the best validation BPB at 10^-3, but final means at that rate range from 2.0337 to 2.1363; these diagnostics use shared AdamW settings rather than product-specific tuning.

5.6 Prefix Depth Ablation

The depth ablation tests TriPLU against matched SwiGLU across layer counts under a shared 5k-step schedule, while additional controls examine product order, gating, and arithmetic behavior.

  • Depth ablation: TriPLU has lower best validation loss than matched SwiGLU at every tested depth from 1 to 8 layers.The comparison uses the public TinyStories 1M-byte prefix and seeds 1–3.
  • Depth ablation: The depth-dependent gap is non-monotonic, so the ablation is a stress test rather than a scaling law.
  • Sample efficiency: TriPLU is the only main-run variant reaching validation-loss target 1.08 in all three seeds.Because throughput is lower, this indicates token-sample-efficiency gains rather than wall-clock or FLOP efficiency.
  • Gated products: The gated-product ablation asks whether unary gates explain TriPLU’s gain, finding scaled gated products better than matched SwiGLU while product-only TriPLU is numerically best.The small gap to scaled triple_swiglu remains comparable to seed variation, so the result does not prove unary gates hurt.
  • Arithmetic diagnostic: On arithmetic_lite, quad_prod has the lowest validation-minimum loss and best exact-match diagnostics, with TriPLU second among product-only rows.This supports the mechanism intuition but not systematic arithmetic extrapolation.
  • Product degree: Degree 4 is competitive but worse than degree 3, while degree 2 is roughly tied with or slightly worse than SwiGLU.The strict comparison varies product order and matched branch width.

6 Discussion

The evidence supports TriPLU in narrow small-model settings, with gains linked to observed co-occurrences but constrained by optimization sensitivity and efficiency limits.

  • TriPLU beats degree 2, degree 4, and closely matched SwiGLU on character-level TinyStories best validation loss.Byte-BPE experiments also improve validation and heldout BPB on TinyStories and WikiText-2 raw under low-learning-rate settings.
  • PMI slices associate gains with seen adjacent-token co-occurrences, while unseen-pair regressions argue against universal improvement.The slice evidence is directional because token counts and intervals are omitted.
  • The comparison is gated pairwise modulation versus direct higher-order products, and performance depends on product order, branch width, scale, tokenization, and optimizer setting.Normalization is presented as a repair direction validated only in a 50M-byte constant-learning-rate diagnostic.
  • TriPLU is slower per token in character-level timing, so the evidence supports validation-loss and token-sample-efficiency gains rather than wall-clock or FLOP-normalized superiority.The paper frames the work as a small-model architecture and benchmark study.

7 Broader Impact Statement

The work is a controlled small-language-model study rather than a deployed or safety-critical system, so its main risks concern overgeneralizing narrow efficiency findings.

  • The study does not introduce a deployed system, user-facing application, human-subject dataset, or safety-critical decision procedure.
  • The main indirect risks are overgeneralizing small-scale efficiency claims, wasting compute on unstable variants, and treating lower loss as deployment readiness.
  • The paper mitigates these risks by bounding the claim, reporting optimization sensitivity and negative settings, and avoiding broad scaling, FLOP-normalized, or production-suitability claims.

8 Reproducibility and Data Availability

The paper documents its experimental implementation, corrected reruns, data settings, checkpoint procedures, and diagnostic run families for reproducibility.

  • Experiments are implemented in PyTorch under repro/, with settings, data prefixes, seeds, checkpoint rules, parameter counts, and run-family identifiers recorded.
  • Corrected public-prefix reruns are families 047–053 with seeds 1–3.
  • The repository summarizes Byte-BPE heldout and PMI evidence, WikiText-2 transfer, fixed-α diagnostics, and the α = 1 constant-learning-rate sweep.

9 Limitations

The evidence is confined to small decoder-only models, limited seeds, and comparisons that are not exactly parameter- or FLOP-matched. Product branches also remain sensitive to learning-rate schedules, while broader scaling and efficiency questions are untested.

  • Character-level evidence covers only 2-layer, 2-head, width-96 decoder-only Transformers, while Byte-BPE uses roughly 7M-parameter models and does not support broad LLM claims.
  • Three seeds expose seed variance but are insufficient for formal significance testing, so robustness is assessed through per-seed paired wins.
  • Comparisons are close parameter-matched rather than exactly parameter- or FLOP-matched, so wall-clock and FLOP-normalized superiority are not claimed.
  • TriPLU is more learning-rate sensitive than SwiGLU in Byte-BPE, and normalization narrows the best-checkpoint gap under hot schedules without preventing final BPB degradation.

10 Conclusion

The conclusion is that directly multiplying learned FFN projections improves several small-model loss metrics in specific low-learning-rate comparisons. The evidence remains narrow and does not establish universal benefits across multiplicative variants or broader scaling and efficiency claims.

  • Directly multiplying learned FFN projections improves best validation loss and tokenizer-neutral validation/heldout BPB in low-learning-rate comparisons across TinyStories and WikiText-2 raw.
  • TriPLU reaches the 1.10 and 1.08 character-level loss targets in all seeds, whereas closely matched SwiGLU does not.
  • The result does not hold for every multiplicative variant, establish arithmetic extrapolation, prove scaling, or remove learning-rate sensitivity.
Loading 2608.20360v1…