Bigram Subnetworks in Language & Knowledge Models

• Bigram subnetworks are distinct subsets of model parameters that capture essential two-unit interactions for tasks like next-token prediction and embedding augmentation.
• They are anatomically concentrated in early Transformer layers and identified via sparsity-driven optimization techniques to faithfully model bigram distributions.
• Their applications in language modeling, knowledge base completion, and word embeddings lead to improved metrics such as MRR, Hits@k, and word similarity scores.

A bigram subnetwork refers to a distinct subset of parameters or architectural components in a model that captures or implements bigram dependencies—statistical or functional relationships between pairs of discrete units, such as sequential tokens in language modeling or entity–relation pairs in knowledge base completion. The formalization, instantiation, and role of bigram subnetworks vary by application domain but consistently center on isolating minimal circuits or embeddings that are both necessary and sufficient for modeling the simplest two-unit interactions. Bigram subnetworks underpin critical behavioral and representational phenomena in contemporary neural architectures, with implications for interpretability, disentanglement, and inductive bias in machine learning.

1. Bigram Subnetworks in Neural LLMs

In large Transformer LLMs, a bigram subnetwork is defined as a sparse set of parameters that, in isolation, enable the model to reproduce the bigram distribution: the empirical next-token probabilities conditioned solely on the immediately preceding token. Formally, given model parameters $\theta$ and an explicit binary mask $M \in \{0,1\}^P$ ($P$ total parameters), the subnetwork $f_M(x) = f(x; M \odot \theta)$ matches the empirical bigram distribution $P_{\mathrm{bigram}}(w_i \mid w_{i-1})$ with surprisal correlation $r > 0.95$ (Chang et al., 21 Apr 2025).

Only a minute fraction of a model's parameters are required for this functionality. For models with up to 1B parameters, bigram subnetworks with $\approx 10^7$ parameters—under 0.2%—are sufficient for high-fidelity bigram prediction. These subnetworks are anatomically concentrated in the first Transformer MLP layer, which is responsible for a sharp geometric transformation of residual stream activations from representing the current token to aligning with next-token predictive space.

Isolation and ablation experiments demonstrate that bigram subnetworks are both necessary and sufficient for naive next-token prediction. Masking out the subnetwork causes a near-complete collapse of coherent language generation, closely mirroring the behavioral deficit from ablating an optimally pruned (cross-entropy optimized) subnetwork of the same size. Conversely, retaining only the bigram subnetwork is enough to achieve high bigram-surprisal correlation, establishing these subnetworks as minimal circuit components for next-token prediction in deep LLMs (Chang et al., 21 Apr 2025).

2. Bigram Embeddings in Knowledge Base Completion

Bigram subnetworks have been operationalized in knowledge base completion through pairwise ("bigram") embeddings learned for pairs of entities and relations, augmenting classical models reliant on singleton (unigram) representations. In the Factorization Machine (FM) framework (Welbl et al., 2016), each fact $f = (s, r, o)$ is encoded not only via embeddings for $s$ (subject), $r$ (relation), $o$ (object), but also for the bigrams $(s, r)$, $(r, o)$, and $(o, s)$. The scoring function for a fact sums all pairwise interactions among singletons and bigrams, allowing direct parameterization of joint statistics between pairs.

Empirical ablations reveal that bigram embeddings, particularly $(r, o)$ pairs, deliver substantial gains in mean reciprocal rank (MRR) and Hits@$k$ on the fb15k-237 dataset. For instance, the $(r, o)$ bigram subnetwork alone boosts MRR from 24.8 (DistMult baseline) to 32.0, and combining all bigram-only subnetworks yields MRR = 33.2 with a Hits@10 of 47.4—an absolute gain of 8.4 MRR points and 9.5 in Hits@10 over the baseline (Welbl et al., 2016). This suggests that directly learning representations for frequent or semantically critical entity–relation pairs is an effective means of capturing data sparsity and enhancing knowledge base completion performance.

3. Bigram Subnetworks in Word and Sentence Embedding Models

Bigram-augmented subnetworks have been incorporated into distributed word embedding schemes, such as CBOW-char and Sent2Vec, by maintaining explicit embedding tables for word bigrams and, optionally, trigrams (Gupta et al., 2019). Enriching the compositional context vectors with word-n-gram embeddings supplements the standard unigram and character-n-gram representations used for context prediction.

During training, all embedding tables (unigram, character n-gram, word n-gram) are optimized jointly via a negative sampling loss. By architectural separation—parametric isolation of word-n-gram tables—and strategic dropout over word-n-grams, the approach enforces a weak disentanglement: unigram embeddings specialize in standalone word semantics, while n-gram embeddings absorb contextual collocational information. At inference, only the unigram table is retained, yielding purer word vectors with improved utility (Gupta et al., 2019).

Empirically, this disentangled subnetwork structure results in consistent improvements across word similarity and analogy tasks. For example, CBOW-char + bigram augmentation improves WordSim353 relatedness scores from 0.626 to 0.652 and MEN from 0.769 to 0.786. The Sent2Vec unigram + bigram model achieves 0.683 on WS353-R, a substantial gain over the pure unigram model’s 0.593 (Gupta et al., 2019). Ablation studies confirm that omitting bigram subnetworks reduces Spearman’s ρ on similarity benchmarks by 0.02–0.05, demonstrating the essential role of bigram representations.

4. Methodologies for Identifying and Training Bigram Subnetworks

Bigram subnetworks in large neural architectures are identified using sparsity-constrained optimization. The continuous sparsification framework introduces a real-valued auxiliary mask $m_j$ for each parameter, transforms it using a sigmoid with annealed temperature $T$, and encourages sparsity via an $\ell_1$ regularization on the mask entries. The loss minimized during optimization comprises the cross-entropy between the masked model’s output and the target bigram distribution, with sparsity controlled by hyperparameter $\lambda$ (Chang et al., 21 Apr 2025). The sparsest subnetwork whose bigram-surprisal correlation remains high ($r > 0.95$) is selected.

In Factorization Machine-based knowledge base models, all pairwise (bigram) and singleton embeddings are learned together, with explicit ablation settings isolating subsets of the possible pairwise interactions to probe the contributions of different bigram types (Welbl et al., 2016).

For word and sentence embedding contexts, three lookup tables are co-trained. Strategic dropout over word-n-grams during training prevents the pollution of unigram vectors with context, compensating for the absence of explicit regularization terms (Gupta et al., 2019).

5. Anatomical Localization and Overlap with Pruned Subnetworks

In Transformer LLMs, the bigram subnetwork is anatomically concentrated in the earliest (first) MLP layer, which is responsible for a geometric "rotation" of the residual stream from current-token to next-token-predictive space (Chang et al., 21 Apr 2025). This indicates that the mapping from input embedding to immediate successor token prediction is largely instantiated early in the network, with later layers responsible for more complex (higher-order, long-range) dependencies.

Bigram subnetworks display significant overlap with masks found by optimizing for optimal pruning against the full next-token cross-entropy loss. For example, in Pythia 1B, 38% of the bigram subnetwork is included within the optimal pruning mask, a 15× enrichment over random overlap (Chang et al., 21 Apr 2025). Despite the optimal mask being five times larger, this overlap points to a shared core set of parameters essential for both naive and high-fidelity next-token prediction.

6. Broader Implications and Applications

Bigram subnetworks serve as minimal interpretable circuits underlying key predictive behaviors in language and knowledge modeling. Their sparsity, geometric transparency, and functional significance render them ideal starting points for bottom-up, compositional interpretability analysis of large neural networks (Chang et al., 21 Apr 2025). For instance, one can envision incrementally assembling higher-order circuits—such as those for n-grams, syntax, or semantics—by grafting additional subnetworks onto the foundational bigram subnetwork.

In the context of knowledge base completion, the explicit inclusion of bigram embeddings directly targets sparsity and memorization issues, outperforming more compositional models by allocating dedicated representational capacity to statistically or semantically salient pairs (Welbl et al., 2016).

In distributed embedding learning, the architectural and training choices supporting bigram subnetworks empirically yield more robust, context-independent word embeddings, augmenting downstream performance on lexical semantic benchmarks (Gupta et al., 2019).

7. Summary Table: Bigram Subnetworks Across Domains

Application Domain	Bigram Subnetwork Instantiation	Principal Role
Language Modeling (Transformer)	Sparse parameter mask over Transformer layers	Drives next-token (bigram) prediction; necessary and sufficient for bigram distribution
Knowledge Base Completion	Pairwise (entity–relation) embeddings in FM	Captures joint entity–relation statistics; increases MRR and Hits@k
Word/Sentence Embeddings	Explicit embedding tables for word bigrams	Disentangles context from unigrams; boosts similarity/analogy performance

Bigram subnetworks are thus a unifying architectural principle for implementing and analyzing pairwise dependencies in diverse neural systems, enabling both performance gains and new forms of mechanistic transparency.