Token Mixing for Industrial Ranking: RankMixer and TokenMixer-Large

Jie Zhu et al. (ByteDance) — RankMixer, CIKM 2025, arXiv:2507.15551 Yuchen Jiang et al. (ByteDance) — TokenMixer-Large, arXiv:2602.06563 (2026)

RankMixer TL;DR

TokenMixer-Large TL;DR

Relations

Builds on: MLP-Mixer (no note yet), DLRM (no note yet), DCN V2 (no note yet), DHEN, Wukong (no note yet) Concepts used: Mixture of Experts, Neural Scaling Laws, Memory-Bound Inference, Standard Attention

Table of Contents

Part I: RankMixer (CIKM 2025)

1. Background and Motivation
   • The DLRM Baseline and Its Limitations
   • Model FLOPs Utilization: Formal Definition
   • Memory-Bound vs Compute-Bound: The Roofline Dichotomy
   • Why Self-Attention Fails in Recommendation
2. Architecture
   • Feature Tokenization
   • Multi-Head Token Mixing
   • Per-Token Feed-Forward Network (PFFN)
   • Full RankMixer Block
   • Mermaid Diagram
3. Hardware Efficiency Analysis
   • Arithmetic Intensity of Self-Attention
   • Arithmetic Intensity of Token Mixing and PFFN
   • MFU Measurement and Serving Cost Decomposition
   • Engineering Optimizations
4. Scaling Experiments
   • Offline Baselines at 100M Parameters
   • Scaling Law Curves
   • Optimal Scaling Directions
5. Sparse MoE Variant
   • Motivation: Two Failure Modes of Vanilla MoE in RankMixer
   • ReLU Routing
   • Dense-Training Sparse-Inference (DTSI-MoE)
   • Scalability Results
6. Online A/B Results — RankMixer
   • Feed Recommendation
   • Advertising and Search
7. Ablation Studies — RankMixer
8. Discussion and Limitations — RankMixer

Part II: TokenMixer-Large (2026)

9. Why Scale Beyond 1B: The Three Failure Modes
10. Architecture Innovations
    • Tokenization
    • Mixing-and-Reverting Operation
    • Per-Token SwiGLU
    • Residuals and Normalization
    • Inter-Residual and Auxiliary Loss
    • Block Architecture Diagram
11. Sparse Per-Token MoE
    • Formulation
    • First Enlarge, Then Sparsify
    • Shared Expert
    • Gate Value Scaling
    • Down-Matrix Small Initialization
12. Scaling to 7B–15B
    • Offline Scaling Curves
    • Data Hunger at Scale
    • DCN Diminishing Returns
13. Training and Serving Optimizations
    • Custom MoE Operators
    • FP8 Quantization
    • Token Parallel Distributed Training
14. Online Experiments — TokenMixer-Large
    • Business Metrics
    • Feed Recommendation Breakdown
15. Ablation Studies — TokenMixer-Large
16. Discussion and Limitations — TokenMixer-Large
17. References

Part I: RankMixer (CIKM 2025)

1. Background and Motivation

🏛️ Industrial recommendation ranking systems must evaluate hundreds of millions of candidate items per second under strict latency budgets (typically <20 ms end-to-end). The dominant architecture family — Deep Learning Recommendation Models (DLRMs) — pairs sparse embedding lookup with a dense neural stack that computes feature interactions. Despite years of accuracy-focused engineering, this lineage suffers from a structural hardware mismatch: the interaction modules were designed for CPUs and are catastrophically inefficient on modern GPUs.

1.1 The DLRM Baseline and Its Limitations

Definition (DLRM). A DLRM decomposes input into sparse categorical features \(\{f_i\}_{i=1}^{N}\) and dense numerical features \(\mathbf{x}_{\text{dense}} \in \mathbb{R}^{d_{\text{dense}}}\). Each categorical feature is looked up in an embedding table:

\[e_i = \text{EmbeddingLookup}(f_i) \in \mathbb{R}^{d_e}\]

The resulting embeddings are passed through an interaction layer (e.g., element-wise products, DCN, self-attention), concatenated with dense features, and fed to an output MLP that produces a click/conversion probability.

The interaction layer is where the architecture family fragments. DCN V2, AutoInt, DHEN, and related models all propose different interaction operators layered on top of this core. Each operator is typically a small, irregular computation — pairwise inner products, attention score matrices over a handful of tokens — that generates very little arithmetic relative to the bytes it touches. On a GPU, this makes the layer memory-bandwidth-bound, not compute-bound.

1.2 Model FLOPs Utilization: Formal Definition

Definition (MFU). Let \(\Pi_{\text{HW}}\) denote the GPU’s peak theoretical FLOPs per second and \(T_{\text{wall}}\) the wall-clock time for a forward pass consuming \(C_{\text{model}}\) floating-point operations. Then:

\[\text{MFU} = \frac{C_{\text{model}}}{T_{\text{wall}} \cdot \Pi_{\text{HW}}}\]

MFU lies in \((0, 1]\); it equals 1 only if the GPU is running at full arithmetic throughput with no stalls from memory latency, kernel launch overhead, or IO-compute imbalance. State-of-the-art LLM training achieves 40–60% MFU on A100/H100 hardware. The DLRM baseline at ByteDance achieves 4.51% — less than one-tenth of what modern hardware should deliver.

1.3 Memory-Bound vs Compute-Bound: The Roofline Dichotomy

Definition (Arithmetic Intensity). For a computation requiring \(F\) floating-point operations and \(B\) bytes of memory traffic (weights + activations read/written), the arithmetic intensity is:

\[I = \frac{F}{B} \quad \text{[FLOPs/byte]}\]

Definition (Ridge Point). For a GPU with peak compute throughput \(\Pi\) [FLOPs/s] and peak memory bandwidth \(\beta\) [bytes/s], the ridge point is:

\[I^* = \frac{\Pi}{\beta}\]

A kernel with \(I < I^*\) is memory-bandwidth-bound: performance is limited by how fast bytes can be transferred, not how fast arithmetic can proceed. A kernel with \(I > I^*\) is compute-bound.

Modern GPUs (e.g., A100 SXM4) have \(\Pi \approx 312\) TFLOP/s (BF16 tensor cores) and \(\beta \approx 2\) TB/s, giving \(I^* \approx 156\) FLOPs/byte. To be compute-bound, a layer must do at least 156 arithmetic operations per byte loaded — a threshold only large matrix multiplications (GEMMs) consistently reach.

The design insight of RankMixer is to restructure every computation into large GEMMs so that \(I \gg I^*\) throughout, pushing MFU from 4.5% to near 45%.

1.4 Why Self-Attention Fails in Recommendation

Self-attention computes pairwise similarities between token pairs. For \(T\) tokens of dimension \(D\), the attention matrix \(A = \text{softmax}(QK^\top / \sqrt{D/H}) \in \mathbb{R}^{T \times T}\) requires \(O(T^2 D)\) FLOPs. In NLP, this is justified because all tokens share a unified vocabulary embedding space — inner products between token embeddings are semantically meaningful.

In recommendation, feature tokens are heterogeneous: a user-ID embedding and an item-category embedding live in different, unrelated spaces. The inner product between them has no semantic grounding and must be learned from scratch via attention weight matrices. The paper’s ablation confirms this: replacing multi-head token mixing with self-attention costs only −0.03% AUC but uses +16% parameters and +71.8% FLOPs. Self-attention is not catastrophically wrong — it is simply compute-inefficient for the heterogeneous-feature regime.

2. Architecture

🏗️ A RankMixer model processes an input of \(T\) feature tokens through \(L\) successive blocks, then applies mean pooling to produce a final representation for task-specific scoring.

2.1 Feature Tokenization

Raw inputs include hundreds of heterogeneous fields: user IDs, video IDs, author metadata, sequence features, and numerical signals. These are first converted to dense embeddings:

\[e_i = \text{EmbeddingLookup}(f_i) \in \mathbb{R}^{d_i}\]

Definition (Feature Tokenization). Features are grouped into \(T\) semantically coherent clusters via domain knowledge. The \(i\)-th token is projected to the common hidden dimension \(D\):

\[x_i = \text{Proj}\!\left(e_{\text{input}}\left[d \cdot (i-1) : d \cdot i\right]\right) \in \mathbb{R}^D, \quad i = 1, \ldots, T\]

The full token matrix fed to the backbone is:

\[\mathbf{X}_0 = \text{stack}[x_1, \ldots, x_T] \in \mathbb{R}^{T \times D}\]

2.2 Multi-Head Token Mixing

Given \(\mathbf{X} \in \mathbb{R}^{T \times D}\) at the input to a block, each token \(x_t \in \mathbb{R}^D\) is split into \(H\) heads of dimension \(D/H\) (the paper sets \(H = T\)):

Definition (Head Splitting). For each token \(x_t\), define the \(h\)-th head as:

\[x_t^{(h)} = x_t\!\left[(h-1) \cdot \frac{D}{H} : h \cdot \frac{D}{H}\right] \in \mathbb{R}^{D/H}\]

Definition (Token Mixing). The \(h\)-th mixed token \(s^{(h)}\) is assembled by concatenating the \(h\)-th head slice from every input token:

\[s^{(h)} = \text{concat}\!\left[x_1^{(h)},\, x_2^{(h)},\, \ldots,\, x_T^{(h)}\right] \in \mathbb{R}^{T \cdot D/H}\]

The full mixed output is:

\[\mathbf{S} = \text{stack}\!\left[s^{(1)}, \ldots, s^{(H)}\right] \in \mathbb{R}^{H \times (T \cdot D/H)}\]

Because \(H = T\), this is a matrix in \(\mathbb{R}^{T \times D}\) — the same shape as the input \(\mathbf{X}\).

Proposition (Token Mixing is a Parameter-Free Permutation). The token mixing operation is an index permutation on the entries of \(\mathbf{X}\). Entry \((t, d)\) in \(\mathbf{X}\) maps to:

\[\text{head index:}\; h = \left\lfloor \frac{d \cdot H}{D} \right\rfloor, \qquad \text{position in }s^{(h)}\text{:}\; (t-1) \cdot \frac{D}{H} + \left(d \bmod \frac{D}{H}\right)\]

This is a bijection on \(\{1, \ldots, TD\}\), so no information is lost and no parameters are consumed.

With residual and layer normalization:

\[\mathbf{S} = \text{LN}\!\left(\text{TokenMixing}(\mathbf{X}) + \mathbf{X}\right) \in \mathbb{R}^{T \times D}\]

2.3 Per-Token Feed-Forward Network (PFFN)

After token mixing, each mixed token \(s_t \in \mathbb{R}^D\) is processed by a dedicated two-layer MLP — one per token position, not shared across positions.

Definition (PFFN). For the \(t\)-th token:

\[v_t = f_{\text{pffn}}^{t,2}\!\left(\text{GELU}\!\left(f_{\text{pffn}}^{t,1}(s_t)\right)\right) \in \mathbb{R}^D\]

with \(W_{\text{pffn}}^{t,1} \in \mathbb{R}^{D \times kD}\) and \(W_{\text{pffn}}^{t,2} \in \mathbb{R}^{kD \times D}\), where \(k\) is the expansion factor.

Why per-token weights? After token mixing, token \(t\) contains a concatenation of one head from each original input token. Because the semantic content of each mixed token position is structurally distinct, a shared FFN would apply the same weights to representations with different semantic structure. The per-token weights allow each position to learn a transformation appropriate for its specific head mixture.

The ablation confirms this: replacing PFFN with a shared FFN costs −0.31% AUC.

Parameter and FLOPs count. For \(T\) tokens, \(D\) hidden dimension, \(L\) layers, expansion factor \(k\):

\[\#\text{Param} \approx 2kLTD^2, \qquad \text{FLOPs} \approx 4kLTD^2\]

Parameter count scales linearly with \(T\), enabling fine-grained capacity control. The FLOPs/Param ratio is 2, independent of architecture details. The baseline DLRM achieves 5.9 GFLOPs/param(M), Wukong 3.6, and RankMixer-1B only 2.1 — more parameters per unit compute.

2.4 Full RankMixer Block

Definition (RankMixer Block). Let \(\mathbf{X}_{n-1} \in \mathbb{R}^{T \times D}\) be the input to block \(n\):

\[\mathbf{S}_{n-1} = \text{LN}\!\left(\text{TokenMixing}(\mathbf{X}_{n-1}) + \mathbf{X}_{n-1}\right)\]

\[\mathbf{X}_n = \text{LN}\!\left(\text{PFFN}(\mathbf{S}_{n-1}) + \mathbf{S}_{n-1}\right)\]

After \(L\) blocks, the final representation is obtained by mean pooling across tokens:

\[\hat{y} = \text{MLP}_{\text{out}}\!\left(\text{MeanPool}(\mathbf{X}_L)\right)\]

2.5 Mermaid Diagram

flowchart TD
    RAW["Raw Features
user/item/sequence/cross fields"]
    EMB["Embedding Lookup
e_i in R^d_i per field"]
    TOK["Feature Tokenization
x_i = Proj(e_input slice)
X in R^{TxD}"]
    RAW --> EMB --> TOK

    subgraph BLOCK["RankMixer Block (repeated L times)"]
        LN1["LayerNorm"]
        MIX["Multi-Head Token Mixing
parameter-free permutation
s^h = concat(x_1^h,...,x_T^h)"]
        ADD1["+ residual"]
        LN2["LayerNorm"]
        PFFN["Per-Token FFN
token-specific W_pffn^t
GELU activation"]
        ADD2["+ residual"]
        LN1 --> MIX --> ADD1 --> LN2 --> PFFN --> ADD2
    end

    TOK --> BLOCK
    BLOCK --> POOL["MeanPool across tokens"]
    POOL --> OUT["Output MLP -> CTR / engagement score"]

    subgraph SMOE["Optional Sparse MoE Extension"]
        direction LR
        RELU_R["ReLU Router
G_ij = ReLU(h(s_i))"]
        EXPERTS["N_e Per-Token Experts
(scaled-down PFFN blocks)"]
        RELU_R --> EXPERTS
    end

    PFFN -.->|"1B+ scale"| SMOE

Figure 1 (Zhu et al., 2025): The RankMixer block architecture. Multi-head token mixing (parameter-free permutation) interleaves head slices across tokens; the per-token FFN (PFFN) then applies token-specific MLP transformations. The optional SMoE extension replaces the dense PFFN at 1B+ scale.

3. Hardware Efficiency Analysis

⚡ The central claim of RankMixer is a 10× MFU improvement over the DLRM baseline.

3.1 Arithmetic Intensity of Self-Attention

It is tempting to assume that attention is inefficient simply because it has \(O(T^2)\) complexity, or because of the heterogeneous feature problem discussed in §1.4. The real explanation is more precise and more instructive.

A full attention block contains several distinct operations. They do not all have the same arithmetic intensity:

The projection GEMMs are efficient for exactly the same reason as RankMixer’s PFFN: the weight matrices are parameters shared across all \(B\) samples in the batch, so memory traffic for weights is \(O(D^2)\) independent of \(B\), giving \(I = B\). These steps are not the problem.

The problem is the attention score computation \(QK^\top\) and the weighted sum \(AV\). These are different in kind: both operands are activations, not parameters. They are distinct for every sample in the batch.

Definition (Activation-to-activation GEMM). A GEMM of the form \(C = AB\) where \(A\) and \(B\) are both activation tensors (i.e., outputs of prior layers, different per sample) has arithmetic intensity independent of batch size.

To see why, account for batch size \(B\) explicitly:

\[I_{QK^\top} = \frac{B \cdot 2T^2 D}{B \cdot 2 \times (2TD + HT^2)} = \frac{2T^2 D}{2 \times (2TD + HT^2)}\]

The \(B\) cancels. Both FLOPs and memory traffic scale linearly with \(B\), so the ratio is constant. For \(T = 32\), \(D = 1536\), \(H = 32\) (FP16, 2 bytes/element):

\[I_{QK^\top} = \frac{2 \times 1024 \times 1536}{2 \times (2 \times 32 \times 1536 + 32 \times 1024)} \approx 12 \text{ FLOPs/byte}\]

Surprisingly, this is far below the A100 ridge point of \(I^* \approx 156\) FLOPs/byte, placing the attention score kernel firmly in the memory-bandwidth-bound regime — regardless of batch size.

3.2 Arithmetic Intensity of Token Mixing and PFFN

Token mixing is a parameter-free data permutation — it touches each byte exactly once and performs no arithmetic. Its arithmetic intensity is 0 FLOPs/byte. The entire computation in a RankMixer block therefore comes from the PFFN.

For the PFFN upward projection applied to a batch of \(B\) samples:

• FLOPs: \(B \times T \times 2D \times kD = 2BkTD^2\)
• Memory traffic (weights, assuming activations in SRAM): \(2kTD^2\) bytes

\[I_{\text{PFFN}} = \frac{2BkTD^2}{2kTD^2} = B\]

For batch size \(B = 1024\), arithmetic intensity \(= 1024\) FLOPs/byte, far above the ridge point of \(\approx 156\). The PFFN is deeply compute-bound.

The general principle, unifying §3.1 and §3.2:

3.3 MFU Measurement and Serving Cost Decomposition

The 100× parameter increase from DLRM to RankMixer translates to only a 3% latency decrease. Two factors multiply: (1) FLOPs/Param ratio is 2.8× lower (fewer FLOPs per parameter), and (2) MFU is 9.9× higher (each FLOP is more productive). Together they nearly cancel the raw FLOPs increase, leaving latency nearly unchanged despite 100× more parameters.

3.4 Engineering Optimizations

Three system-level techniques further reduce inference latency for the deployed 1B model:

1. Per-token FFN operator fusion — multiple PFFN computations merged into a single 3D tensor operation: +30% throughput.
2. Mixed-precision inference (FP16) — matrix multiplications use FP16; LayerNorm uses FP32: +45% throughput, −31.5% latency.
3. Sparse-GEMM acceleration — custom sparse-GEMM for the SMoE variant (§5): −40% latency.

4. Scaling Experiments

📈

4.1 Offline Baselines at 100M Parameters

Experiments use Douyin’s production training data over a two-week window, with 300+ input features. An improvement of 0.01% AUC is considered confidently significant at production scale.

RankMixer-100M outperforms every baseline including Wukong while using 47% fewer FLOPs per batch. Scaling to 1B adds a further +0.31% Finish AUC with a 9× FLOPs increase.

4.2 Scaling Law Curves

The paper plots Finish AUC gain as a function of both parameter count and FLOPs across five architectures. Key observations:

• RankMixer exhibits the steepest slope on both the AUC vs parameters and AUC vs FLOPs curves.
• Wukong’s parameter-curve slope is steep but its FLOPs-curve slope is gentler.
• DHEN shows non-ideal scaling, reflecting limited scalability of cross-structure stacking.

The steepness of RankMixer’s scaling curve is the primary architectural claim: for a given parameter or FLOPs budget, RankMixer extracts more AUC gain than any alternative tested.

Figure 2 (Zhu et al., 2025): Scaling laws comparing Finish AUC gain vs parameter count (left) and FLOPs (right) across DLRM-MLP, DCN V2, DHEN, Wukong, and RankMixer. RankMixer exhibits the steepest slope on both axes. The x-axis is logarithmic.

4.3 Optimal Scaling Directions

RankMixer scales along four orthogonal axes: token count \(T\), hidden dimension \(D\), number of layers \(L\), and number of MoE experts \(E\). The paper finds:

• Model quality correlates primarily with total parameter count; different \((T, D, L)\) combinations achieving the same total reach nearly identical AUC.
• Increasing \(D\) (wider) is preferable to increasing \(L\) (deeper): wider \(D\) generates larger GEMM shapes, achieving higher MFU.

Final configurations chosen: - RankMixer-100M: \(D = 768\), \(T = 16\), \(L = 2\) - RankMixer-1B: \(D = 1536\), \(T = 32\), \(L = 2\)

5. Sparse MoE Variant

⚙️ Scaling RankMixer beyond 1B parameters while maintaining fixed inference latency requires a Sparse Mixture-of-Experts (SMoE) extension that decouples parameter count from active FLOPs.

5.1 Motivation: Two Failure Modes of Vanilla MoE in RankMixer

Standard sparse MoE (Switch Transformer style with top-\(k\) + softmax gating) degrades markedly when naively applied to RankMixer’s PFFN:

1. Uniform routing ignores token information content. Different feature tokens carry different amounts of information — a rich user behavior sequence token conveys far more signal than a sparse cross-feature token. Top-\(k\) routing allocates the same number of expert activations to every token regardless, wasting capacity on low-information tokens.

2. Expert under-training from token-count explosion. PFFN already multiplies parameter count by \(T\). Adding \(N_e\) non-shared experts multiplies further: total experts becomes \(T \times N_e\). With a fixed routing budget of \(k\) experts per token, most experts receive very few gradient updates, leading to expert starvation.

5.2 ReLU Routing

Definition (ReLU Routing). Let \(h(\cdot) : \mathbb{R}^D \to \mathbb{R}^{N_e}\) be a learned router (linear layer). For token \(s_i\) and expert \(j\), the gate value is:

\[G_{i,j} = \text{ReLU}\!\left(h(s_i)_j\right) \geq 0\]

The aggregated output for token \(s_i\) is:

\[v_i = \sum_{j=1}^{N_e} G_{i,j} \cdot e_{i,j}(s_i)\]

Unlike softmax + top-\(k\), ReLU routing allows a variable number of experts to activate per token. Sparsity level is steered via a regularization loss:

\[\mathcal{L} = \mathcal{L}_{\text{task}} + \lambda \mathcal{L}_{\text{reg}}, \qquad \mathcal{L}_{\text{reg}} = \sum_{i=1}^{N_t} \sum_{j=1}^{N_e} G_{i,j}\]

The \(\ell_1\) penalty on gate values directly penalizes the total number of active expert slots, adaptively controlling the average activation ratio without fixing it per-token.

5.3 Dense-Training Sparse-Inference (DTSI-MoE)

Definition (DTSI-MoE). Two router functions \(h_{\text{train}}\) and \(h_{\text{infer}}\) are both updated during training. The regularization loss \(\mathcal{L}_{\text{reg}}\) is applied only to \(h_{\text{infer}}\):

• \(h_{\text{train}}\): unregularized, tends toward dense activation, providing broad gradient coverage to all experts.
• \(h_{\text{infer}}\): regularized via \(\mathcal{L}_{\text{reg}}\), tends toward sparse activation, enabling fast inference.

At inference, only \(h_{\text{infer}}\) is used. The experts are trained on dense gradient signals but evaluated under sparse routing, preventing the “dying expert” failure where sparsity at training time starves most experts.

5.4 Scalability Results

• Vanilla SMoE (top-k): AUC degrades monotonically as sparsity increases.
• Vanilla SMoE + load-balancing loss: Some recovery, but still substantially below dense.
• DTSI + ReLU routing: Near-flat AUC curve from full activation down to 1/8 sparsity.

RankMixer with DTSI + ReLU routing scales to 8× sparsity with nearly no AUC loss and a +50% throughput improvement, validating the approach as a practical path to 10B+ parameters without proportional cost increase.

Figure 3 (Zhu et al., 2025): AUC performance of RankMixer variants at decreasing expert activation ratios (1, 1/2, 1/4, 1/8). Dense-training + ReLU-routed SMoE (DTSI-MoE) preserves near-full accuracy of the 1B dense model across all sparsity levels, while vanilla top-k MoE degrades monotonically.

Figure 4 (Zhu et al., 2025): Activated expert ratio for different token positions in RankMixer. ReLU routing produces adaptive, token-dependent activation patterns — information-rich tokens activate more experts than sparse or low-signal tokens, naturally allocating capacity where it is most needed.

6. Online A/B Results — RankMixer

🚀 RankMixer-1B was deployed for full production traffic across Feed Recommendation, Advertising, and Search on Douyin. Experiments ran for five months.

6.1 Feed Recommendation

Low-active users benefit disproportionately — their active day lift (0.46%) is 3.7× that of high-active users (0.12%). Larger model capacity helps most when personal history is sparse: the model draws on richer cross-feature and cross-user statistical patterns.

6.2 Advertising and Search

The search −1.0% query change rate means users find what they want with fewer reformulations, indicating genuine relevance improvement.

7. Ablation Studies — RankMixer

🔬

Block component ablations (100M scale)

Token routing strategy comparison (100M scale)

The comparison to self-attention is particularly revealing: self-attention costs only −0.03% AUC but requires +71.8% more FLOPs. Token mixing trades a negligible 0.03% AUC for a 42% FLOPs reduction relative to self-attention.

8. Discussion and Limitations — RankMixer

💬

Architectural unification. The core thesis is that a single well-designed block — parameter-free token mixing followed by per-token FFN — subsumes the functionality of an entire zoo of handcrafted interaction modules (DCN, AutoInt, DHEN, FM-based approaches).

Hardware-aware design philosophy. The design choices are explicitly reverse-engineered from GPU arithmetic intensity requirements: large GEMMs, parameter-free permutations for cross-token interaction, and batch-size-amplified FLOPs-per-byte.

Limitations:

• Residual dimension mismatch. The token mixing operation changes the layout of the token matrix (from \(\mathbb{R}^{T \times D}\) to \(\mathbb{R}^{H \times (T \cdot D/H)}\)), creating an impedance mismatch for inter-block residuals at large depths. The paper uses only \(L = 2\) blocks. Part II (§9) addresses this.

• No head count (\(H\)) ablation. The paper sets \(H = T\) throughout but provides no sensitivity analysis.

• Retrieval stage unexplored. All results pertain to re-ranking (scoring shortlisted candidates). Whether the architecture transfers to embedding-based retrieval is not addressed.

• Single-task framing. Multi-task optimization challenges that arise at scale are not discussed.

Part II: TokenMixer-Large (2026)

9. Why Scale Beyond 1B: The Three Failure Modes

🏛️ TokenMixer-Large begins from a precise diagnosis of why naively scaling RankMixer past ~1B parameters fails. Three architectural failure modes are identified:

1. Dimension mismatch. The mixing output \(\mathbf{H} \in \mathbb{R}^{H \times (T \cdot D/H)}\) has a different layout than the input \(\mathbf{X} \in \mathbb{R}^{T \times D}\) even though they contain the same number of scalars. A reshape is required before any residual connection back to \(\mathbf{X}\), which breaks pre-norm symmetry and degrades gradient magnitude at depth.

2. Gradient vanishing at depth. Without skip connections spanning multiple blocks, gradients to early layers become vanishingly small as depth grows beyond ~20 blocks.

3. Uniform dense FFN. The per-token SwiGLU treats all tokens identically. At 7B+ parameters, this wastes capacity by forcing every expert computation to fire for every token.

TokenMixer-Large addresses each failure mode systematically: Mixing-and-Reverting for (1), inter-residual connections for (2), and Sparse Per-token MoE for (3).

10. Architecture Innovations

🏗️

10.1 Tokenization

Each raw categorical feature \(F_i\) is embedded:

\[e_i = \text{Embedding}(F_i, d_i) \in \mathbb{R}^{d_i}\]

Features are organized into \(T-1\) semantic groups \(G_0, \ldots, G_{T-2}\). Each group is projected to dimension \(D\) by a group-specific MLP:

\[X_i = \text{MLP}_i\!\bigl(\text{concat}[e_l, \ldots, e_m]\bigr) \in \mathbb{R}^D\]

A global token \(X_G\) aggregates cross-group information:

\[X_G = \text{MLP}_g\!\bigl(\text{concat}[G_1, \ldots, G_{T-1}]\bigr) \in \mathbb{R}^D\]

The full token matrix is \(\mathbf{X} = \text{concat}[X_G, X_0, \ldots, X_{T-1}] \in \mathbb{R}^{T \times D}\).

10.2 Mixing-and-Reverting Operation

The central innovation is a two-phase symmetric transform that resolves the dimension-mismatch problem.

Definition (Mixing Phase). Given \(\mathbf{X} \in \mathbb{R}^{T \times D}\), for each head \(h\), concatenate the \(h\)-th slice from every token:

\[\text{Mix}: \quad H_h = \text{concat}[x_1^{(h)}, x_2^{(h)}, \ldots, x_T^{(h)}] \in \mathbb{R}^{T \cdot D/H}\]

Stacking gives \(\mathbf{H} = \text{stack}[H_1, \ldots, H_H] \in \mathbb{R}^{H \times (T \cdot D/H)}\). A pSwiGLU is applied in this mixed layout to produce \(\mathbf{H}' \in \mathbb{R}^{H \times (T \cdot D/H)}\).

Definition (Reverting Phase). The reverting operation is the inverse permutation of mixing: for each position \(t\), gather slice \(h\) from \(H_h'\) to reconstruct the token:

\[\text{Revert}: \quad X_t^{\text{rev}} = \text{concat}[x'^{(1)}_t, x'^{(2)}_t, \ldots, x'^{(H)}_t] \in \mathbb{R}^D\]

This yields \(\mathbf{X}^{\text{rev}} \in \mathbb{R}^{T \times D}\) — exactly the same shape as the input \(\mathbf{X}\).

Definition (TokenMixer-Large Block Output).

\[\mathbf{X}^{\text{next}} = \text{Norm}\!\bigl(\text{pSwiGLU}(\mathbf{X}^{\text{rev}}) + \mathbf{X}\bigr) \in \mathbb{R}^{T \times D}\]

The residual \(+ \mathbf{X}\) is now dimensionally consistent. Reverting is not merely a reshape — it explicitly recombines mixed-head representations back into per-token vectors, allowing the subsequent pSwiGLU to operate in the original token-feature space.

10.3 Per-Token SwiGLU

TokenMixer-Large uses per-token SwiGLU (pSwiGLU), where each token position \(t\) has its own projection matrices:

\[\text{pSwiGLU}(\cdot) = FC_{\text{down}}\!\bigl(\text{Swish}(FC_{\text{gate}}(\cdot)) \odot FC_{\text{up}}(\cdot)\bigr)\]

with token-specific projections \(FC_i(\mathbf{x}) = W_i^t x_t + b_i^t\), where \(\{W_{\text{up}}^t, W_{\text{gate}}^t\} \in \mathbb{R}^{D \times nD}\) and \(W_{\text{down}}^t \in \mathbb{R}^{nD \times D}\).

10.4 Residuals and Normalization

Pre-Norm with RMSNorm throughout, consistent with modern LLM practice:

\[\text{Output} = \text{SubLayer}(\text{RMSNorm}(\mathbf{X})) + \mathbf{X}\]

10.5 Inter-Residual and Auxiliary Loss

For networks beyond ~20 blocks, standard residuals are insufficient to propagate gradients to early layers. TokenMixer-Large introduces inter-residual connections: skip connections that bypass 2–3 consecutive blocks.

Definition (Inter-Residual). Let \(\mathbf{X}^{(\ell)}\) denote the output of block \(\ell\). An inter-residual with stride \(s\) adds:

\[\mathbf{X}^{(\ell+s)} \leftarrow \mathbf{X}^{(\ell+s)} + \mathbf{X}^{(\ell)}\]

at regular intervals \(s \in \{2, 3\}\) throughout the network.

An auxiliary loss is applied at intermediate block outputs:

\[\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{CE}}(y, \hat{y}^{(L)}) + \lambda \sum_{\ell \in S} \mathcal{L}_{\text{CE}}(y, \hat{y}^{(\ell)})\]

Ablation: removing both inter-residuals and auxiliary loss costs −0.04% AUC at the 4B scale; removing the standard within-block residual costs −0.15% AUC.

Figure 5 (Jiang et al., 2026): Inter-residual connections and auxiliary loss in TokenMixer-Large. Skip connections with stride \(s \in \{2, 3\}\) bypass groups of consecutive blocks, and auxiliary cross-entropy losses at intermediate layers provide direct gradient signals to early blocks — essential for stable training at 50+ block depth.

X^(0) ──► Block 1 ──► X^(1) ──► Block 2 ──► X^(2)
   └──────────────── (+) ──────────────────►  X^(2) += X^(0)

X^(2) ──► Block 3 ──► X^(3) ──► Block 4 ──► X^(4)
   └──────────────── (+) ──────────────────►  X^(4) += X^(2)

10.6 Block Architecture Diagram

flowchart TD
    X["X in R^{TxD}
input tokens"] --> RN1["RMSNorm"]
    RN1 --> MIX["Mixing Phase
split heads, concat across tokens
H in R^{Hx(T D/H)}"]
    MIX --> SWIGLU1["pSwiGLU
(in mixed layout)"]
    SWIGLU1 --> REV["Reverting Phase
gather heads to per-token
X^rev in R^{TxD}"]
    REV --> ADD1["+ residual X"]
    ADD1 --> RN2["RMSNorm"]
    RN2 --> SWIGLU2["Per-token SwiGLU
token-specific W^t"]
    SWIGLU2 --> ADD2["+ residual"]
    ADD2 --> OUT["X^next in R^{TxD}"]
    X -.->|"inter-residual
every s=2-3 blocks"| SKIP["downstream block"]
    OUT --> SKIP

Figure 6 (Jiang et al., 2026): The TokenMixer-Large block architecture. Each block consists of (RMSNorm → Mixing → SP-MoE → Reverting → RMSNorm → SP-MoE) with inner and inter-block residuals. The Reverting phase restores the token-major layout after mixing, enabling dimensionally consistent residuals at any depth. SP-MoE replaces the dense pSwiGLU at 7B–15B scale.

11. Sparse Per-Token MoE

⚡ At 7B–15B parameters, it becomes infeasible to activate all parameters for every token at every layer. TokenMixer-Large introduces Sparse Per-token MoE (SP-MoE).

11.1 Formulation

Let there be \(E\) experts, each a scaled-down pSwiGLU. For token \(x_t\), a gating network selects the top-\(k\) experts:

\[\text{SP-MoE}(x_t) = \sum_{j=1}^{k} g_j(x_t) \cdot \text{Expert}_j(x_t)\]

where \(g_j(x_t) = \text{softmax}(\text{top-}k(W_g x_t))_j\). Each expert has width \(nD/E\), so that each expert is \(1/E\) the width of the corresponding dense pSwiGLU.

Per-token routing means each token independently selects its \(k\) experts. This contrasts with sequence-level MoE (Switch Transformer), which routes entire sequence positions identically — suboptimal for ranking where different semantic token groups (user vs. item vs. context) should leverage different experts.

11.2 First Enlarge, Then Sparsify

The strategy is sparse train, sparse infer — sparsity is fixed at training time so inference requires no special conversion. For a \(1:2\) sparsity model, active FLOPs drop by approximately \(2\times\) while parameter count doubles relative to a single-expert baseline.

Figure 7 (Jiang et al., 2026): “First Enlarge, Then Sparsify” illustration. Starting from a dense baseline, the model is first enlarged by adding experts (increasing total parameters), then sparsified so only a fraction of experts activate per token at inference. This sequence — train sparse from the start — avoids the dense-to-sparse conversion gap that degrades accuracy in post-hoc pruning.

11.3 Shared Expert

One expert is always active, regardless of gating:

\[\text{SP-MoE}(x_t) = \sum_{i=1}^{k-1} g_i(x_t) \cdot \text{Expert}_i(x_t) + \text{SharedExpert}(x_t)\]

The shared expert acts as a “default path” ensuring all tokens receive at least one full transformation, preventing catastrophic forgetting of common patterns when the router is uncertain. Removing it costs −0.02% AUC.

Figure 9 (Jiang et al., 2026): SP-MoE expert activation load balance at 1:2 sparsity. Each bar shows the fraction of tokens routed to each expert across the sequence. The distribution is approximately uniform, confirming that the auxiliary load-balancing loss prevents expert collapse even under per-token routing.

11.4 Gate Value Scaling

A scalar \(\alpha\) is applied to the gated sum, set inversely proportional to the sparsity ratio: \(\alpha = 2\) for \(1:2\) sparsity, \(\alpha = 4\) for \(1:4\) sparsity. Without this correction, the softmax gate values sum to 1 over \(k-1\) selected experts, so the magnitude of the routed contribution decreases as \(k\) shrinks. Removing \(\alpha\) costs −0.03% AUC.

11.5 Down-Matrix Small Initialization

The down-projection matrix \(W_{\text{down}}^{t,j}\) of each expert pSwiGLU is initialized with standard deviation \(0.01\) (vs. default \(1.0\)). This forces expert outputs near-zero at initialization, so the model starts close to a passthrough and learns expert specialization gradually. This is analogous to small-init residual branches in NTK theory. Removing it costs −0.03% AUC.

12. Scaling to 7B–15B

📈

12.1 Offline Scaling Curves

The paper fits a log-linear relationship between parameter count \(N\) and AUC gain:

\[\Delta\text{AUC}(N) \approx a \cdot \log N + b\]

Key offline results (vs. DLRM-MLP-500M baseline):

A key finding: beyond 1B parameters, scaling requires balanced expansion across width \(D\), depth \(L\), and expansion factor \(n\) simultaneously — scaling any single dimension yields diminishing returns.

Figure 8 (Jiang et al., 2026): Scaling laws comparing AUC gain vs parameter count (left) and FLOPs (right) for Wukong, RankMixer (TokenMixer), and TokenMixer-Large across multiple Douyin scenarios (e-commerce 15B, 7B, 4B). TokenMixer-Large achieves consistent log-linear AUC improvements up to 15B parameters, with the steepest scaling slope among all tested models. The x-axis is logarithmic.

The paper further notes that DCN-style cross-network components become less valuable at larger scales:

This suggests the token-mixing backbone subsumes the cross-feature interaction function that DCN was designed to provide, making DCN redundant at scale.

12.2 Data Hunger at Scale

The 2.3B model trained for 30 days underperforms the 500M model — it simply has not seen enough data to fill its capacity. At 60 days, it recovers and exceeds the 500M baseline. Larger models require proportionally more training data, consistent with Chinchilla-style scaling laws.

12.3 DCN Diminishing Returns

13. Training and Serving Optimizations

🔧

13.1 Custom MoE Operators

The permute and unpermute operations reorder token activations so that all tokens routed to expert \(j\) are contiguous in memory before the GroupedFFN kernel executes — essential for batched matrix multiplication efficiency. The GroupedFFN dominates at both training (89%) and serving (98%), making it the target for FP8 quantization.

13.2 FP8 Quantization

The MoEGroupedFFN is quantized to FP8 (E4M3 format) for serving, providing 2× memory bandwidth reduction vs FP16 and hardware-accelerated matrix multiplication on H100 GPUs. Result: 1.7× serving speedup with negligible AUC degradation.

13.3 Token Parallel Distributed Training

Token parallelism partitions the \(T\) tokens across \(P\) GPUs, keeping model parameters on each device while splitting the sequence. Each GPU processes \(T/P\) tokens; an all-reduce aggregates after the per-token FFN.

A 4-way token parallel configuration yields: - 29.2% throughput improvement (raw, without communication overlap) - 96.6% throughput improvement with communication-computation overlap - MFU improved to 60% in the advertising backbone

14. Online Experiments — TokenMixer-Large

🚀

14.1 Business Metrics

The +2.98% GMV gain on e-commerce is the headline result, representing one of the largest single-model improvements reported in recent industrial recommender papers.

14.2 Feed Recommendation Breakdown

Low-active users benefit most — ~12× vs high-active users on like rate — consistent with larger model capacity helping most when user histories are sparse.

15. Ablation Studies — TokenMixer-Large

🔬

TokenMixer-Large Block ablations (4B scale)

SP-MoE ablations (4B scale)

The last row is the most informative: replacing per-token routing with global (sequence-level) MoE routing costs −0.10% AUC at zero parameter or FLOPs overhead. This confirms that per-token routing is qualitatively more expressive than global routing for the heterogeneous-feature setting of recommendation ranking.

16. Discussion and Limitations — TokenMixer-Large

💬

Pure model design. A recurring theme is the elimination of fragmented operators — ad-hoc task-specific components accumulated over years of system iteration. TokenMixer-Large argues that a single well-designed block (Mixing-and-Reverting + pSwiGLU + SP-MoE) subsumes the functionality of these operators while being easier to scale and profile.

Scaling to 15B is not free. Three practical constraints:

1. Data constraint. The 2.3B model requires ≥60 days of Douyin training data to converge; scaling to 15B implies even longer horizons or faster data pipelines.
2. Latency constraint. Serving a 15B model within a 10 ms SLA requires careful FP8 quantization, model parallelism, and hardware-specific kernel tuning.
3. Scenario-specific saturation. The e-commerce scenario saturates at 7B; live streaming at 4B. Continued scaling should be driven by scenario-specific offline scaling laws, not a global parameter target.

Limitations:

• No retrieval-stage results. TokenMixer-Large is a ranking model; transfer to embedding-based retrieval (ANN search) is unexplored.
• Cursory load balancing analysis. Expert activation distributions are shown but routing collapse or expert specialization are not quantified.
• No ablation on head count \(H\). A critical hyperparameter for Mixing-and-Reverting, but no sensitivity analysis is provided.

References