PaperlayerView exact revision on arXiv

Source-linked AI summary

Wiola 13M, a Gated Spiral Attention Architecture for Parameter Efficient Small Language Models

Aryuemaan Kumar Chowdhury, Praveen Oosa, Vineesha Reddy

arXiv:2608.14604v1cs.CLcs.AI

TL;DR

Small language models need architecture choices adapted to tight positional, attention, and parameter budgets. Wiola introduces three drop-in layer components and shows exact equivalence between full-sequence training and cached autoregressive decoding, with a 12,915,888-parameter Nano configuration.

  • Problem

    Small language models face tight positional, attention-head, and feed-forward parameter budgets that standard transformer choices were designed for much larger models.

  • Method

    Wiola modifies each decoder layer with parameter-free Spiral Rotary encoding, causally gated attention, and a multiplicative Butterfly feed-forward block.

  • Results

    12,915,888 parameters define the Nano configuration, while gated attention matches full-sequence training and cached decoding to floating-point tolerance without approximation.

  • Takeaways & Limitations

    Wiola provides a clean, open-source baseline for studying architecture in the small language model regime.

  • Takeaways & Limitations

    The reference implementation uses explicit attention rather than fused kernels, forgoing their memory benefits at long context.

Abstract

from arXiv · show

Small language models in the ten to one hundred million parameter range are attractive for on device inference, rapid experimentation, and controlled scientific study, yet most of them reuse the standard transformer block without adaptation to the small scale regime. We present Wiola, a decoder only language model whose novelty is concentrated in three drop in components of every layer. First, Spiral Rotary Positional Encoding perturbs the standard rotary frequencies by a slowly growing per dimension factor so that phase trajectories fan outward, improving long range discrimination while adding no parameters. Second, Gated Spiral Attention introduces a per head, content adaptive scalar gate derived from a causal cumulative statistic of the query stream, providing an implicit and differentiable form of soft head selection at negligible cost. Third, the Butterfly feed forward block replaces the conventional expansion layer with a multiplicative interaction and an intra block bypass path, matching the parameter count of a four times gated linear unit block while improving gradient flow in shallow stacks. We formalize each component, derive exact parameter and computation budgets, and prove that the gated attention admits an exact and numerically verified equivalence between full sequence training and cached autoregressive decoding, so that no approximation is introduced at inference time. We also describe a fully reproducible training and evaluation protocol on a standard tiny story corpus. The reference implementation is released as an open source package with weights ready publishing support.

I. INTRODUCTION

Wiola targets the ten-to-one-hundred-million-parameter regime with three drop-in transformer-layer modifications addressing positional efficiency, limited attention-head specialization, and feed-forward parameter costs. The paper also provides exact budgets and verifies that gated attention preserves equivalence between parallel training and cached decoding.

  • Motivation: Small language models offer local deployment and modest-compute training, but commonly inherit transformer designs tuned for much larger models.The target regime is roughly ten to one hundred million parameters.
  • Contribution: Wiola introduces three targeted transformer-layer modifications while leaving the surrounding training recipe standard.The modifications address positional encoding, attention-head specialization, and feed-forward parameter efficiency.
  • Spiral Rotary Positional Encoding: Spiral Rotary Positional Encoding scales rotation frequencies by a slowly increasing factor across dimension pairs, adding no parameters and recovering standard rotary encoding in a limiting case.Its phase trajectories fan outward with depth in the frequency spectrum.
  • Gated Spiral Attention: Gated Spiral Attention computes a per-head scalar gate from a causal cumulative mean of query projections and applies it multiplicatively to pre-softmax scores.The gate enables unhelpful heads to suppress themselves as an implicit soft head selection and adds only a few hundred parameters per layer.
  • Butterfly Feed-Forward Block: The Butterfly feed-forward block uses multiplicative expansion and an intra-block bypass path while matching a four-times gated linear unit block in parameter count.The bypass path improves gradient propagation in shallow networks.
  • Verification: Wiola provides exact parameter and computation budgets and proves and numerically verifies exact equivalence between parallel training and cached decoding for gated attention.This prevents approximation from being introduced at inference time by the gate definition.

II. RELATED WORK … D. Small Language Models

The related work situates Wiola among positional encoding, gated feed-forward, and attention-efficiency methods, while targeting the reproducible small-model regime. Its components adapt these ideas toward long-range discrimination, gradient flow in shallow stacks, and useful computation across a small number of heads.

  • A. Positional Encoding: Fixed sinusoidal embeddings inject order, whereas Rotary Position Encoding rotates queries and keys by position so attention depends only on relative offset.Rotary encoding requires no learned parameters and has become standard in modern decoder-only models.
  • A. Positional Encoding: Rotary encoding is favored in modern decoder-only models because it extrapolates gracefully without learned parameters.Alternative approaches include additive linear biases for length extrapolation.
  • B. Feed Forward and Gated Linear Units: Gated linear units improve feed-forward blocks through multiplicative interactions between two projections, typically using three matrices with reduced inner width.This preserves the parameter budget while changing the feed-forward computation.
  • B. Feed Forward and Gated Linear Units: The Butterfly block extends gated feed-forward designs with an explicit intra-block bypass that supports gradient flow when the layer count is small.Its design follows the parameter-preserving tradition of gated linear unit variants.
  • C. Attention Efficiency and Gating: Key-value sharing and restructured attention reduce attention cost, targeting large models and long-context memory bandwidth.These approaches differ from Wiola’s aim of improving the usefulness of each head rather than compressing attention.
  • C. Attention Efficiency and Gating: Wiola uses a content-adaptive gate to modulate each of a small number of heads, aligning more closely with conditional computation than attention compression.The gate is intended to make individual heads more useful without reducing attention cost.
  • D. Small Language Models: The tiny story setting showed that small models can generate coherent text from a controlled distribution.This setting provides the experimental context for Wiola’s small-model target.
  • D. Small Language Models: Compute-optimal scaling analysis clarified how parameters and tokens should be balanced, motivating Wiola as a clean and reproducible baseline in this regime.Wiola targets the small-language-model setting represented by the tiny story corpus.

III. MODEL ARCHITECTURE · A. Overview and Notation · B. Spiral Rotary Positional Encoding

Wiola is a pre-norm decoder-only transformer built from repeated residual attention and Butterfly feed-forward sub-blocks, with tied input and output embeddings. Its Spiral Rotary Positional Encoding applies a parameter-free, monotonically increasing frequency perturbation that preserves rotary relative-offset behavior and reduces to standard rotary encoding when α = 0.

  • A. Overview and Notation: Wiola embeds a length-T token sequence, processes it through L identical decoder layers, normalizes the result, and projects to vocabulary logits with tied embedding weights.The model uses hidden width d, H attention heads, and per-head width dh = d/H.
  • A. Overview and Notation: Each decoder layer contains pre-norm Gated Spiral Attention and Butterfly feed-forward sub-blocks, both wrapped in residual connections.The attention and feed-forward components are denoted GSA and FFN, respectively.
  • A. Overview and Notation: Root mean square normalization is used throughout, with a learned gain γ ∈ R^d and a small constant ε.Figure 1 presents the overall data flow, while Figure 3 details the attention block.
  • B. Spiral Rotary Positional Encoding: Standard rotary encoding partitions each head vector into dh/2 coordinate pairs and rotates pair i at position m by angle mθ_i under a geometric frequency schedule.This provides the baseline construction that Spiral Rotary Positional Encoding modifies.
  • B. Spiral Rotary Positional Encoding: Spiral encoding scales each frequency by a slowly growing, parameter-free factor controlled by one fixed hyperparameter α.The factor grows monotonically from 1 at the lowest pair toward 1+α at the highest pair.
  • B. Spiral Rotary Positional Encoding: Higher-frequency pairs receive a proportionally larger boost and separate phases more rapidly with position, while the lowest pairs remain almost unchanged and preserve local continuity.Setting α = 0 recovers standard rotary encoding exactly, making Spiral encoding a strict generalization.
  • B. Spiral Rotary Positional Encoding: Applying the same per-position rotation to queries and keys preserves rotary encoding’s relative-offset property: their dot product depends on t − s.The transformation is applied identically to query and key vectors at each position.

C. Gated Spiral Attention

Gated Spiral Attention uses a strictly causal, content-adaptive scalar gate for each head, applied before softmax to enable differentiable soft head selection. Gates are initialized near one half, preserving behavior close to ungated attention at training start.

  • Causal gate: By default, the gate uses pre-rotation queries so it depends on content rather than absolute position, with this choice exposed as a configuration flag.Queries and keys receive Spiral encoding, while the default gate input is taken before that rotation.
  • Causal gate: A causal cumulative mean of the query stream feeds a two-layer network that produces one scalar gate per attention head.The gate summarizes positions up to the current timestep, so it is strictly causal.
  • Causal gate: Zero initialization of b2 makes every gate approximately one half at training start, keeping attention close to the ungated baseline.The gate network uses W1 ∈ R^(H×d_h) and W2 ∈ R^(H×H).
  • Gated scores: The gate multiplies each head’s scaled attention scores before the softmax, under a causal mask, before outputs are concatenated and projected.The causal mask assigns zero to permitted positions and −∞ to future positions.
  • Gated scores: Small gates flatten attention rows toward near-uniform distributions, reducing head contributions, while unhelpful heads are driven toward zero without an auxiliary sparsity loss.This implements implicit, fully differentiable soft head selection.

D. Butterfly Feed Forward Block

The Butterfly feed forward block combines multiplicative gating with an intra-block bypass path. With d_ff = 2d, it matches the parameter count of a conventional four-times expansion block while supporting gradient flow in shallow stacks.

  • D. Butterfly Feed Forward Block: The block expands the hidden state into two halves, applies a multiplicative gate, and adds an intra-block bypass before the down projection.This design uses a single up projection, a bypass projection, and a down projection.
  • D. Butterfly Feed Forward Block: The multiplicative term matches the expressiveness of gated linear unit feed forward blocks, while the additive bypass provides a linear input-to-output path that improves gradient flow when L is small.The bypass is intended to benefit shallow stacks.
  • D. Butterfly Feed Forward Block: With d_ff = 2d, the three matrices contain 4d d_ff parameters, identical to a conventional four-times expansion block with a single activation.The added multiplicative and bypass structure therefore comes at no parameter-count cost.

E. Training Objective

Wiola is trained with the standard next-token objective. Tied input and output embeddings produce the logits, and optimization uses mean token negative log-likelihood.

  • The model uses the standard next-token training objective.
  • Tied embedding matrix E produces logits from the final-layer hidden state via RMSNorm(h(L))E⊤.
  • The loss is the mean token negative log-likelihood.

IV. THEORETICAL ANALYSIS … C. Exact Equivalence of Training and Cached Decoding

The theoretical analysis establishes Wiola’s parameter and compute budgets, then proves and numerically verifies that gated attention produces exactly identical outputs during full-sequence training and cached autoregressive decoding. The design preserves the layer’s asymptotic complexity while requiring only a small additional gate state during inference.

  • A. Parameter Budget: 12,915,888 parameters yield the Nano configuration’s target of about 12.9 million parameters across six layers with tied embeddings.The attention block contributes four projection matrices, the gate contributes Hdh + H2 + H parameters, and the Butterfly block contributes 4ddff.
  • B. Computational Complexity: The attention projections and Butterfly feed-forward block are quadratic in width and linear in sequence length T.These components contribute to the per-layer computational cost summarized in Table III.
  • B. Computational Complexity: The score and value products remain quadratic in sequence length, as in dense attention.Their contribution is part of the O(Td2 + T 2d) layer complexity.
  • B. Computational Complexity: The gate adds a cumulative mean and small two-layer network with linear cost in T, leaving layer complexity at O(Td2 + T 2d).Its cost is independent of T 2.
  • C. Exact Equivalence of Training and Cached Decoding: For any input sequence, Proposition 1 states that the per-head gate from full-sequence computation equals the gate from autoregressive decoding with cached running sums.The cached path maintains the same prefix sum and position count as the parallel path.
  • C. Exact Equivalence of Training and Cached Decoding: Identical gates make attention scores and outputs identical at every position because subsequent attention computations are deterministic functions of the gate, queries, keys, and values.Key and value caching is exact by construction.
  • C. Exact Equivalence of Training and Cached Decoding: Floating-point-tolerance agreement at every position on random inputs, with and without gating, confirms that cached decoding introduces no approximation.The only additional decoding state is the running sum St, a vector of dh values per layer.
  • C. Exact Equivalence of Training and Cached Decoding: Cached decoding updates each layer’s running sum with the new head-mean query, divides by the token count, and applies the two-layer sigmoid gate.The procedure also carries key and value caches, a running sum initialized to zero, and a count initialized to zero.

D. Inference Memory

Inference memory is dominated by the key-value cache, while the gated attention mechanism adds only a negligible running state during decoding.

  • D. Inference Memory: 2LTd elements: the key-value cache is the dominant memory term for a context of length T.The cache stores keys and values across all L layers and d dimensions.
  • D. Inference Memory: 6.3 megabytes: the Nano configuration’s key-value cache at a 1024-token context in half precision.This quantifies the cache memory requirement for the stated Nano inference setting.
  • D. Inference Memory: 192 values: the Nano gate’s total running state, represented by Ldh elements and negligible relative to the cache.The gate state is maintained during decoding and adds only a small memory overhead.

V. REPRODUCIBLE EVALUATION PROTOCOL · VI. DISCUSSION AND LIMITATIONS · VII. CONCLUSION

The paper specifies a reproducible tiny-story evaluation protocol with one-flag ablations, while acknowledging that architectural analysis does not establish accuracy gains. It concludes that Wiola’s three components are parameter-efficient, cache-exact where applicable, and supported by exact budget derivations and verification.

  • V. REPRODUCIBLE EVALUATION PROTOCOL: The Nano evaluation recommends the tiny story corpus and a byte-level tokenizer with a 32,000-token vocabulary.The tokenizer may be trained on the corpus or reused from an existing model.
  • V. REPRODUCIBLE EVALUATION PROTOCOL: Optimization uses AdamW with momentum coefficients 0.9 and 0.95 under a cosine schedule.
  • V. REPRODUCIBLE EVALUATION PROTOCOL: The ablation plan separately tests Spiral encoding, Gated Spiral Attention, and the Butterfly bypass path against specified alternatives.The tests set α = 0, disable the gate, or replace Butterfly with an equal-parameter conventional gated linear unit block.
  • V. REPRODUCIBLE EVALUATION PROTOCOL: Each ablation is implemented through a single flag change because all three toggles are exposed in the released configuration.
  • VI. DISCUSSION AND LIMITATIONS: The analysis establishes inexpensive components and cache-exact gated attention, but does not establish an accuracy improvement without empirical measurement.The paper states that accuracy must be measured under the specified protocol.
  • VI. DISCUSSION AND LIMITATIONS: The spiral coefficient α is fixed rather than learned, and the gate uses only a first moment of the query stream.The passage identifies learned or per-layer spiral schedules and richer gate statistics as possible alternatives.
  • VII. CONCLUSION: Wiola concentrates novelty in three drop-in components: parameter-free phase fanning, causal cache-exact soft head selection, and multiplicative bypass structure without increasing the parameter budget.The conclusion also reports exact parameter and computation budgets and proof with numerical verification.
Loading 2608.14604v1…