Multi-Head Neural N-gram Layer

• The paper introduces a layer that biases heads to attend only to localized n-gram spans, improving computational efficiency and imposing useful inductive biases.
• It details variants such as sliding-window masking, linear-MLP local layers, and induction heads that balance expressivity with speed and domain-specific performance.
• Empirical evaluations show minimal BLEU penalties with significant speedups, enhanced training stability, and robust performance across translation, language modeling, and reinforcement learning tasks.

A Multi-Head Neural N-gram Layer generalizes multi-head attention by constraining or biasing individual heads to operate on restricted, locally-constrained or linguistically-motivated n-gram spans, rather than allowing each head to attend globally. This design provides computational efficiency, strong locality biases, and novel inductive structures with minimal or sometimes positive impact on performance in canonical generation, translation, and in-context learning tasks. There exist several concrete formalisms for multi-head neural n-gram layers, including: sliding-window n-gram masking in self-attention, local head MLPs over n-gram windows, n-gram–based attention over phrase memories, and explicit n-gram pattern induction masks. Each variant yields a different trade-off between expressivity, speed, and domain specificity.

1. Architectural Variants and Formal Definitions

The principal architectures for Multi-Head Neural N-gram layers are:

• N-gram Masked Self-Attention: Self-attention weights at target position $i$ are computed solely over the local window $j \in [\max(1, i-N+1), i]$; a per-position, per-head binary mask $M$ restricts the receptive field. For each attention head $h$ in layer $\ell$, the query, key, and value projections

$$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$

are aggregated over the mask window. The masked logits are

$$\widetilde{S}_{i,j} = S_{i,j} + M_{i,j},\qquad \alpha_{i, :} = \mathrm{softmax}(\widetilde{S}_{i, :}),$$

with $M_{i,j}=0$ for $j$ in the window, $-\infty$ otherwise (Chelba et al., 2020).

• Multi-Head Neural N-gram (Linear-MLP Local Layer): For each sequence position $j \in [\max(1, i-N+1), i]$0, a local window of $j \in [\max(1, i-N+1), i]$1 neighboring tokens is concatenated, yielding $j \in [\max(1, i-N+1), i]$2 (where $j \in [\max(1, i-N+1), i]$3 for unidirectional, $j \in [\max(1, i-N+1), i]$4 for bidirectional context). Each head $j \in [\max(1, i-N+1), i]$5 computes

$j \in [\max(1, i-N+1), i]$6

The head outputs are concatenated, projected, then passed through residual and normalization layers as in standard Transformers (Loem et al., 2022).

• N-gram Induction Heads (Pattern Masked Head): Individual heads are hard-wired with fixed binary affinity masks $j \in [\max(1, i-N+1), i]$7, activating only if the preceding $j \in [\max(1, i-N+1), i]$8 inputs at positions $j \in [\max(1, i-N+1), i]$9 and $M$0 match. The n-gram head computes

$M$1

using learned projections, possibly gated and softly scaled (Zisman et al., 2024).

• Multi-Granularity Heads (Phrase-level Attention): Transformers reserve a subset of heads to attend not to tokens, but to phrase memory banks composed from n-gram segmentation and composition (often via intra-phrase self-attention). The dot-product attention then operates over phrase-level keys/values rather than tokens (Hao et al., 2019).

2. Computational Complexity and Efficiency

Restricting attention or aggregation to local n-gram windows makes the n-gram layers substantially more efficient than standard multi-head self-attention:

Layer Type	Time per Layer	Memory per Layer
Multi-Head Self-Attn	$M$2	$M$3
N-gram Masked Attn	$M$4	$M$5
MH N-gram (MLP window)	$M$6	$M$7

For a sequence of length $M$8, the attention workload is reduced by a factor of $M$9, often yielding 2–3x speedups in self-attention computation and 1.5–2x end-to-end throughput improvements for $h$0 (Chelba et al., 2020, Loem et al., 2022). Memory requirements scale linearly with $h$1 and $h$2 rather than quadratically in $h$3.

3. Implementation Methodology

N-gram Masked Self-Attention (Sliding Window)

• Replace full-causal attention by windowed attention: at each position $h$4, restrict summation to $h$5.
• Share the same binary $h$6-gram mask across all heads.
• For incremental decoding, maintain a circular buffer of the last $h$7 key/value pairs per head.

Multi-Head Neural N-gram Layer (Linear-MLP Window)

• For each position $h$8, concatenate $h$9 nearest neighbor vectors.
• Apply $\ell$0 parallel linear-ReLU projections (one per head), concatenate, and project through a learned output matrix.
• Compose with residual and LayerNorm.
• For the encoder, use a bidirectional window; for decoder, last $\ell$1 tokens.

N-gram Induction Heads

• For each designated head and n-gram order $\ell$2, precompute binary affinity masks $\ell$3 encoding n-gram token matches.
• At each layer, compute output by combining a standard projection path and an n-gram masked path, modulated by a learned scalar gate and (optional) mask scaling parameter.

Multi-Granularity Attention

• In the encoder's first layer, partition the token sequence into non-overlapping n-grams, compose phrase representations, and reserve groups of heads for n-gram attention (alongside token heads).
• Aggregate head outputs via concatenation, followed by linear projection as in standard multi-head attention.

4. Empirical Performance and Trade-offs

Language Modeling and Machine Translation

• On WMT’18 En→De with N-gram Masked Self-Attention: for $\ell$4 from 4 to 8, BLEU reduction is ≤0.4 (e.g., baseline 28.6, 8-gram 28.2), with negligible perplexity increases (Chelba et al., 2020).
• On WMT’14 En→Fr, a 6–10-gram window provides near-identical BLEU as baseline (e.g., baseline 41.1, 8-gram 40.7), with 1.7x total decoding speedup for $\ell$5.
• Multi-Head Neural n-gram (linear-MLP): on IWSLT De→En, achieves BLEU 35.49 (vs. Transformer-base 35.34); on WMT En→De, BLEU 27.15 (vs. 27.20). Increasing head count is crucial (single-head n-gram → 25.33 BLEU) (Loem et al., 2022). Hybrid stacks (mixing n-gram and MHSA layers) sometimes further improve BLEU by +0.5.
• Multi-Granularity Self-Attention: Adding n-gram heads yields +0.52 BLEU (27.83 vs. 27.31) on WMT’14 En→De, and further gains with explicit syntactic heads and auxiliary losses. Probing confirms explicit n-gram heads encode broader syntactic phenomena (Hao et al., 2019).

In-Context Reinforcement Learning

• N-gram Induction Heads reduce required data in Algorithm Distillation for RL by up to 27× and sharply improve training stability (from 400 to 20 hyperparameter trials needed to reach a solution in grid-world) (Zisman et al., 2024).
• In low-data regimes, models with n-gram heads succeed in generalizing with much less supervision or tuning.

5. Analysis of Inductive Bias and Ablations

• Across masking and explicit n-gram head designs, a uniform shared n-gram window or mask across all heads is typically sufficient; there is no observed performance benefit from head-specific $\ell$6 or adaptive segmentation, but implementation complexity increases (Chelba et al., 2020).
• Local n-gram and global MHSA heads are complementary: hybrid and layered approaches yield marginal BLEU gains and improved robustness (Loem et al., 2022).
• In multi-granularity models, syntactic or surface n-gram heads in lower layers improve phrase representation, while at higher layers, switching to broader-context (or full MHSA) enhances model expressivity (Hao et al., 2019).
• Adding global context vectors to the encoder (e.g., max-pooled span) has mild impact, but does not substitute for local n-gram inductive bias (Loem et al., 2022).
• For n-gram induction heads, performance is highly sensitive to the number of n-gram heads, order $\ell$7, gate initialization, and the scale of fixed mask projections (Zisman et al., 2024). Initializing n-gram heads with identity matrices and small mask scaling improves early-stage training stability.

6. Applications and Limitations

• Multi-Head Neural N-gram layers enable deployment of efficient sequence models in resource-constrained or low-latency settings.
• Particularly in incremental/auto-regressive decoding (e.g., machine translation, speech), n-gram masking reduces computation and memory, with minimal BLEU/pplx trade-off.
• In in-context RL, n-gram induction heads directly encode induction/copy constructs, resulting in faster generalization and greater robustness to hyperparameter variation.
• A plausible implication is that hard n-gram biases allow faster discovery of compositional structure where such motifs are prevalent in the data, but may degrade if global dependencies dominate.
• No evidence has been reported of benefit in dynamically learning the window size per head, or combining adaptive n-gram spans in a data-driven fashion, though heterogeneity in n-gram order across heads does not materially harm performance.

7. Integration and Practical Guidelines

• For n-gram masked self-attention: set $\ell$8 for En-De, $\ell$9 for En-Fr; end-to-end speedup $$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$01.7x with BLEU penalty $$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$1; keep the same $$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$2 (number of heads) as baseline (Chelba et al., 2020).
• For MLP-style neural n-gram: use $$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$3 throughout, or schedule $$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$4 to be smaller in lower layers, increasing towards the top; $$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$5 heads and depth matched to Base/Big Transformer (Loem et al., 2022).
• For n-gram induction heads in RL, allocate a small but nontrivial number of heads (e.g., 3/8) to n-gram induction, orders $$q_i^{(\ell, h)} = W_q^{(\ell, h)} x_i^{(\ell-1)},\quad k_j^{(\ell, h)} = W_k^{(\ell, h)} x_j^{(\ell-1)},\quad v_j^{(\ell, h)} = W_v^{(\ell, h)} x_j^{(\ell-1)}$$6; initialize mask pathway with small weights and turn on gradually via scalar gates (Zisman et al., 2024).
• For multi-granularity models, phrase-level heads are best placed in lower encoder layers; auxiliary syntactic losses provide measurable gain if gold phrase boundaries are available (Hao et al., 2019).
• In all cases, multi-head organization is essential—single-head n-gram variants underperform substantially relative to multi-head counterparts.

• "Faster Transformer Decoding: N-gram Masked Self-Attention" (Chelba et al., 2020)
• "Are Neighbors Enough? Multi-Head Neural n-gram can be Alternative to Self-attention" (Loem et al., 2022)
• "Multi-Granularity Self-Attention for Neural Machine Translation" (Hao et al., 2019)
• "N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs" (Zisman et al., 2024)