Tern-1: Ternary Hierarchical Sliding-Linear Transformers for Efficient Long-Context Reasoning
Gugu8 / Tern-1
~105M-parameter ternary looped language model with native 1M-context architecture
JAX / Flax · Kaggle TPU · Apache 2.0
| Property | Value |
|---|
| Parameters | 104.97 M |
| Hidden size (d) | 880 |
| Physical layers (L) | 6 |
| Loops per layer (R) | 8 |
| Effective depth | (L \times R = 48) |
| Attention | THSL (window + memory) |
| Native max sequence | (1,048,576) |
| Train sequence length | 128 |
| Vocabulary | 50 257 (GPT-2) |
Abstract
We introduce Tern-1, a compact language model that combines three design principles rarely applied jointly near the 100 M scale: (i) ternary weights via BitLinear-style quantisation with a straight-through estimator, (ii) weight-tied depth looping that expands six physical layers into forty-eight effective residual blocks without increasing parameter count, and (iii) a novel Ternary Hierarchical Sliding-Linear (THSL) attention mechanism that supports a native context length of one million tokens while remaining memory-efficient during training.
THSL realises long-range modelling through a hierarchical decomposition: a local causal sliding-window softmax of width (W) together with a residual memory bank that is read and written across loops. The architecture is further augmented with a chain-of-thought (CoT) latent token injected at sequence position zero and residual loop gates that allow the model to adaptively emphasise or suppress successive loop iterations.
Tern-1 was trained end-to-end on TPU with JAX/Flax under a curriculum of next-token prediction followed by heavy synthetic multi-step arithmetic CoT and real web/math corpora. Held-out evaluation on GSM8K and HellaSwag (never used in training) is reported. All weights, the tokenizer, and a complete training recipe are released.
1. Introduction
Large language models achieve strong reasoning performance through scale, yet many research and deployment settings require models that remain tractable on a single TPU or modest GPU while still supporting multi-step reasoning and long contexts. Two orthogonal research directions address complementary aspects of this constraint:
- Low-bit / ternary quantisation reduces the cost of matrix multiplications and the memory footprint of weights.
- Looped (recurrent-depth) Transformers reuse the same physical layers multiple times, increasing effective depth at constant parameter count.
Tern-1 unifies both ideas and adds a purpose-built attention mechanism (THSL) that is linear in sequence length for the global component and therefore compatible with a 1 M token native context length. A lightweight residual memory bank and a CoT latent prior further encourage iterative reasoning inside the fixed training context of 128 tokens.
Contributions
- A ~105 M ternary Transformer with (8\times) weight-tied loops (6 physical layers (\to) 48 effective passes).
- THSL attention: local sliding-window causal softmax of width (W=64) combined with a cross-loop residual memory bank, configured for
max_seq_native = 1 048 576.
- Residual loop gates and a 32-slot memory bank that accumulate state across loops.
- A CoT latent token injected at the first position of every sequence.
- Full open weights (Flax msgpack), GPT-2 tokenizer, and a reproducible TPU training recipe.
2. Model Architecture
2.1 Ternary Linear Layers
All linear projections follow a BitLinear-style ternary quantisation. Given a full-precision weight matrix (W \in \mathbb{R}^{d_{\mathrm{in}} \times d_{\mathrm{out}}}),
$$
s = \frac{1}{d_{\mathrm{in}} d_{\mathrm{out}}} \sum_{i,j} |W_{ij}| + \varepsilon, \qquad
W_q = \mathrm{clip}\bigl(\mathrm{round}(W / s), -1, 1\bigr).
$$
The forward pass uses the straight-through estimator (STE):
$$
y = x \bigl( W + \mathrm{stopgrad}(s W_q - W) \bigr).
$$
This yields exact ternary multiplies at inference while permitting gradient flow through the continuous weights during training.
2.2 Weight-Tied Looping
Each of the (L=6) physical layers is applied (R=8) times. Inside a layer the residual update is modulated by a learned per-loop gate:
$$
g_r = \sigma(\gamma_r), \qquad
x \leftarrow x + g_r \cdot f_\theta(x; r),
$$
where (\gamma \in \mathbb{R}^R) is a vector of logits and (f_\theta) denotes the layer transformation (attention + MLP). A per-loop scale vector (s_r \in \mathbb{R}^d) further multiplies the incoming residual stream, allowing the model to re-weight features at successive depths.
2.3 Residual Memory Bank
A bank (M \in \mathbb{R}^{B \times S \times d}) with (S=32) slots is maintained across loops. At each loop iteration a query is formed from the current residual stream,
$$
a = \mathrm{softmax}(W_q , \mathrm{RMSNorm}(x)), \qquad
x \leftarrow x + g_r \cdot a M.
$$
A write is performed by adding a projected mean-pooled state into a cycling slot:
$$
M_{:,, r \bmod S,, :} \leftarrow M_{:,, r \bmod S,, :} + 0.05 \cdot W_w , \mathrm{mean}(x).
$$
The memory therefore functions as a hierarchical long-range pathway that is independent of the attention window.
2.4 THSL Attention
Ternary Hierarchical Sliding-Linear attention is designed so that the dominant memory term scales as (O(T W)) rather than (O(T^2)), while still supporting a native context of (T = 2^{20}).
Given input (x \in \mathbb{R}^{B \times T \times d}), ternary projections produce queries, keys and values which are reshaped into (H=10) heads of dimension (h = d / H). For a local window of width (W = \min(64, T)):
- Keys and values are left-padded by (W-1) positions.
- A sliding gather yields tensors of shape ((B, H, T, W, h)).
- Scaled dot-product scores are computed only inside the window:
$$
\mathrm{scores}{b,h,t,w} = \frac{1}{\sqrt{h}} , q{b,h,t}^\top k_{b,h,t,w}.
$$
- Softmax is taken over the window dimension and the weighted values are contracted back to ((B, T, d)).
The global / hierarchical component is supplied by the residual memory bank described above; together they form the hierarchical design. Because the windowed path never materialises a full (T \times T) matrix, the same kernel can be used for sequences up to the configured native maximum of (1,048,576) tokens.
2.5 CoT Latent Token
A learned vector (c \in \mathbb{R}^d) is added to the first position of the residual stream after the token + position embeddings:
$$
x_{:,0,:} \leftarrow x_{:,0,:} + c.
$$
This provides a persistent “scratch-pad” prior that is present for every forward pass and is intended to bias the model toward explicit intermediate reasoning.
2.6 Full Forward Pass
$$
\begin{align*}
x &\leftarrow \mathrm{Embed}(ids) + \mathrm{Pos}(0:T-1) \
x_{:,0} &\leftarrow x_{:,0} + c \
M &\leftarrow 0 \
\text{for } \ell &= 1 \dots L: \quad
\text{for } r = 1 \dots R: \
&\quad x, M \leftarrow \mathrm{SmartBlock}\ell(x, M, r) \
\mathrm{logits} &\leftarrow \mathrm{RMSNorm}(x) , W{\mathrm{te}}^\top
\end{align*}
$$
Weight tying is used between the input embedding and the final projection.
3. Training
3.1 Data Curriculum
GSM8K and HellaSwag are held out completely.
- Pre-training on FineWeb-Edu, OpenWebMath and Cosmopedia-style mixtures.
- Chain-of-thought fine-tuning with multi-step arithmetic and word problems that force intermediate reasoning.
Tokenisation uses the GPT-2 BPE tokenizer (50 257 tokens). Sequences are packed to a fixed length of 128 with right padding.
3.2 Optimisation
- Optimizer: AdamW ((\beta_1=0.9), (\beta_2=0.95), weight decay (0.01))
- Learning rate schedule from (1.0 \times 10^{-3}) (pre-train) to (2.5 \times 10^{-4}) (CoT)
- Gradient clipping: global norm 1.0
- Precision: bfloat16 throughout
- Parallelism:
jax.pmap over available TPU cores, fixed-shape JIT compilation
3.3 Implementation Notes
- All linear layers are ternary (STE).
- Loop indices are taken modulo the number of loops / memory slots so that the same physical parameters are reused.
- The model is compiled once with a fixed ((B, T)) shape; subsequent training steps never trigger recompilation.
4. Evaluation
4.1 Protocol
Neither GSM8K nor HellaSwag appears in any training stage. Evaluation uses greedy decoding (temperature 0) and a maximum of 32 newly generated tokens for GSM8K. HellaSwag uses standard multiple-choice log-probability ranking of the four endings given the context.
| Benchmark | Split | (n) evaluated | Metric |
|---|
| GSM8K | test | 120 | exact numeric match after “Answer:” |
| HellaSwag | validation | 200 | accuracy of highest-log-prob ending |
4.2 Results
| Benchmark | Accuracy |
|---|
| GSM8K | 0.83 % |
| HellaSwag | 18.5 % |
These are the honest measured results of the released ~105 M checkpoint. The synthetic + web/math curriculum drives training loss down but has limited transfer to the linguistic form of GSM8K and to HellaSwag ranking when those sets are fully held out.
4.3 Throughput
Under the fixed-shape training regime the model sustained several thousand tokens per second on TPU hardware after compilation.
5. Files and Usage
5.1 Repository Contents
| File | Description |
|---|
flax_model.msgpack | Main model parameters (Flax, ~201 MB) |
config.json | Architecture hyper-parameters and evaluation metrics |
tokenizer.json / tokenizer_config.json | GPT-2 tokenizer |
README.md | This document |
5.2 Loading the Model (Flax)
1import jax.numpy as jnp
2from flax.serialization import from_bytes
3# Instantiate the Tern1THSL module graph matching the config
4
5with open("flax_model.msgpack", "rb") as f:
6 params = from_bytes(None, f.read())
7
8ids = jnp.array([[...]]) # shape (B, T) int32, T <= 128 for the released checkpoint
9logits = model.apply({"params": params}, ids)
5.3 Native 1 M Context
The released checkpoint was trained at (T=128). The THSL attention implementation itself never materialises a full attention matrix and is therefore usable, with appropriate memory-bank and positional handling, up to the configured max_seq_native = 1_048_576. Extending the positional embeddings and re-tuning the memory write schedule for ultra-long sequences is left as future work.
5.4 Inference Constraints
- Pad or truncate to a fixed length (\le 128) for the current weights.
- Causal masking is implicit in the sliding-window construction.
- Prefer fixed shapes + JIT for decode steps to avoid repeated compilation.
5.5 License
Apache 2.0. The tokenizer follows the GPT-2 / Hugging Face conventions.
6. Design Rationale
Why ternary + loops?
Ternary weights reduce arithmetic cost; loops buy serial depth. The combination yields a deeper compute graph per parameter, which is attractive when activation memory is the limiting resource.
Why THSL?
A pure local window cannot capture long-range dependencies; a pure linear global attention loses the sharp local inductive bias that helps language modelling. THSL keeps both, routes the long-range signal through a cheap residual memory, and stays (O(T W)) in the dominant term.
Why a CoT latent token?
Providing an explicit, always-present scratch-pad prior encourages the model to utilise the looped depth for intermediate computation rather than collapsing to a single-pass prediction.
7. Limitations
- Current released weights were trained only at sequence length 128.
- Transfer to open-domain mathematical language remains limited (as reflected by the GSM8K score).
- The STE introduces a train–inference discrepancy if the quantisation scales are not carefully maintained.
- No official Hugging Face
PreTrainedModel wrapper is provided yet; users must instantiate the Flax module graph.
8. Related Work
- BitNet / BitLinear ternary and low-bit Transformers
- Universal Transformers and other looped / recurrent-depth architectures
- Sliding-window and linear-attention variants for long context
- Chain-of-thought prompting and process supervision
9. Citation
1@misc{tern1-2026,
2 title = {Tern-1: Ternary Hierarchical Sliding-Linear Transformers for Efficient Long-Context Reasoning},
3 author = {Gugu8},
4 year = {2026},
5 howpublished = {\url{https://huggingface.co/Gugu8/Tern-1}},
6 note = {~105M ternary BitLinear, weight-tied loops, THSL attention, residual memory bank, CoT latent token; trained on TPU with JAX/Flax}
7}
10. Changelog
- v1 – Base ternary looped model.
- v2 – Memory bank + loop gates + CoT token.
- v3 – THSL attention (window + hierarchical memory), native 1 M configuration.
- v4 – Scaled to ~105 M parameters (dim 880, 6 layers × 8 loops), full curriculum, honest evaluation.
- v4.1 – README rewritten using official Hugging Face KaTeX delimiters (
$$...$$ display, \(...\) inline).
Built for curiosity: small models, deep loops, ternary weights, and hierarchical long-context attention.
