--- title: Multi-Head Neural N-gram Layer url: https://www.emergentmind.com/topics/multi-head-neural-n-gram-layer type: topic --- # Multi-Head Neural N-gram Layer 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 [2001.04589]. - **Multi-Head Neural N-gram (Linear-MLP Local Layer):** For each sequence position $t$, a local window of $n$ neighboring tokens is concatenated, yielding $c_t \in \mathbb{R}^{m d}$ (where $m = n$ for unidirectional, $2n-1$ for bidirectional context). Each head $k$ computes \[ h_t^{(k)} = \mathrm{ReLU}(c_t W_k + b_k). \] The head outputs are concatenated, projected, then passed through residual and normalization layers as in standard Transformers [2207.13354]. - **N-gram Induction Heads (Pattern Masked Head):** Individual heads are hard-wired with fixed binary affinity masks $A^{(n)} \in \{0, 1\}^{T \times T}$, activating only if the preceding $n$ inputs at positions $i$ and $j$ match. The n-gram head computes \[ H^{(n)} = W_1^{(n)} H^l + W_2^{(n)} [A^{(n)}]^T H^l \] using learned projections, possibly gated and softly scaled [2411.01958]. - **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 [1909.02222]. ## 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 | $O(T^2 d)$ | $O(T^2)$ | | N-gram Masked Attn | $O(N T d)$ | $O(N T)$ | | MH N-gram (MLP window) | $O(L n d)$ | $O(L n)$ | For a sequence of length $T \gg N$, the attention workload is reduced by a factor of $T/N$, often yielding 2–3x speedups in self-attention computation and 1.5–2x end-to-end throughput improvements for $N=8$ [2001.04589, 2207.13354]. Memory requirements scale linearly with $T$ and $N$ rather than quadratically in $T$. ## 3. Implementation Methodology ### N-gram Masked Self-Attention (Sliding Window) - Replace full-causal attention by windowed attention: at each position $i$, restrict summation to $j \in [\max(1,i-N+1), i]$. - Share the same binary $N$-gram mask across all heads. - For incremental decoding, maintain a circular buffer of the last $N-1$ key/value pairs per head. ### Multi-Head Neural N-gram Layer (Linear-MLP Window) - For each position $t$, concatenate $n$ nearest neighbor vectors. - Apply $K$ 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 $n$ tokens. ### N-gram Induction Heads - For each designated head and n-gram order $n$, precompute binary affinity masks $A^{(n)}$ 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 $N$ from 4 to 8, BLEU reduction is ≤0.4 (e.g., baseline 28.6, 8-gram 28.2), with negligible perplexity increases [2001.04589]. - 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 $N=8$. - 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) [2207.13354]. 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 [1909.02222]. ### 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) [2411.01958]. - 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 $N$ or adaptive segmentation, but implementation complexity increases [2001.04589]. - Local n-gram and global MHSA heads are complementary: hybrid and layered approaches yield marginal BLEU gains and improved robustness [2207.13354]. - 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 [1909.02222]. - 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 [2207.13354]. - For n-gram induction heads, performance is highly sensitive to the number of n-gram heads, order $n$, gate initialization, and the scale of fixed mask projections [2411.01958]. 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 $N=8$ for En-De, $N=8\text{--}10$ for En-Fr; end-to-end speedup $\sim$1.7x with BLEU penalty $\leq 0.4$; keep the same $H$ (number of heads) as baseline [2001.04589]. - For MLP-style neural n-gram: use $n=5$ throughout, or schedule $n$ to be smaller in lower layers, increasing towards the top; $K$ heads and depth matched to Base/Big Transformer [2207.13354]. - For n-gram induction heads in RL, allocate a small but nontrivial number of heads (e.g., 3/8) to n-gram induction, orders $n=1\text{--}3$; initialize mask pathway with small weights and turn on gradually via scalar gates [2411.01958]. - 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 [1909.02222]. - In all cases, multi-head organization is essential—single-head n-gram variants underperform substantially relative to multi-head counterparts. ## References - "Faster Transformer Decoding: N-gram Masked Self-Attention" [2001.04589] - "Are Neighbors Enough? Multi-Head Neural n-gram can be Alternative to Self-attention" [2207.13354] - "Multi-Granularity Self-Attention for Neural Machine Translation" [1909.02222] - "N-Gram Induction Heads for In-Context RL: Improving Stability and Reducing Data Needs" [2411.01958] Source: https://www.emergentmind.com/topics/multi-head-neural-n-gram-layer