Language Models Can Benefit from Domain-Specific Inductive Biases

AI summary

Overview

Research area: Machine learning / large language model architecture — specifically, how the choice of activation functions (nonlinearities) shapes what a transformer can learn from different kinds of pretraining data.

Technical level: Intermediate. The paper assumes familiarity with transformers, activation functions (GeLU, ReLU, softmax), and standard training recipes, but its central idea — swapping activations for learnable splines — is explained clearly enough for readers with a working ML background.

Scope: A large-scale empirical study (over 6,000 small models of 40–120M parameters) testing whether transformers whose nonlinearities are tuned for one data domain transfer well to other domains.

What This Paper Is About

Language models are usually built from one architecture trained on a huge, diverse mix of web text, code, and mathematics. This paper asks whether that one-size-fits-all design leaves performance on the table: could different subsets of pretraining data benefit from different built-in assumptions (inductive biases)? The authors test this by giving transformers nonlinearities specialized for specific datasets and then measuring how well those specialized models perform both on their own domain and on others.

Key Contributions

1. A method for domain-specialized nonlinearities. Standard GeLUs in MLP layers and the nonlinearity in attention layers (inserted before the QKV projections) are replaced with learnable linear splines. These are optimized during a warm-up phase on a source dataset, then frozen and transferred to a target dataset — only 2 × 64 = 128 parameters move between phases.

2. An improved training procedure over prior work. Compared to Teney et al. (2026), the authors remove optimization on held-out data, add staggered starts across multiple models sharing the same splines (M = 20 parallel models), and change where nonlinearities are placed in attention layers — the last of which they report provides strong benefits.

3. A systematic cross-domain transfer study. Using a mix-and-match design over seven datasets (FineWeb, Enwik8clean, TinyStoriesV2, Shakespeare, CodeSearchNet-Java, CodeSearchNet-Python, and a DeepMind Math subset with 56 question types), they measure whether architectures tuned on one domain help or hurt on another.

4. An analysis of why the gains occur. The paper rules out surface-statistics explanations (via scrambled-token controls) and identifies two mechanisms: wider loss basins and a greater abundance of "lottery tickets" — useful circuits already present at initialization.

Main Findings

• Standard transformers are rarely optimal for any single domain. On all seven tested datasets, the specialized variants improved training efficiency or final performance relative to baselines.

• Even a trivial activation swap already shows domain effects. GeLU MLPs performed better on natural language and code, while ReLU MLPs performed better on math.

• Transfer is highly asymmetric. Nonlinearities optimized for language were generally useful — as effective as baselines on Math, and clearly better on language and code. Nonlinearities optimized for Math performed worse than baselines on language and code, and most did not even transfer across math question types.

• Language tuned on FineWeb transfers to TinyStories. Architectures tuned on a natural language dataset such as FineWeb perform well on TinyStories, supporting the abstract's claim of within-domain compatibility.

• Clean, structured datasets transfer most reliably. Nonlinearities optimized for Enwik8 and Csn-Java transferred best. Transfer was also more successful in attention layers than in MLP layers, consistent with the view that attention handles structural patterns while MLPs handle semantic content, which is intrinsically more dataset-specific.

• Shakespeare produced outliers. The authors attribute poor transfer from Shakespeare to its peculiar dialogue formatting, speaker labels, stylistic regularities, and its much smaller size.

• Specialists benefit more than generalists on math. On the Math subset, learned nonlinearities improved a generalist model by +1.3% over a 28.7% baseline, versus +3.5% for specialist models over a 33.7% baseline. With fewer training tokens (24.1% baseline), specialists improved by +7.5%. Improvements persisted with up to ~20 tokens/parameter and with increased depth (12 vs. 4 layers).

• The gains are not surface statistics. Repeating the warm-up on scrambled token sequences — which preserves distributional properties but removes structure — produced nonlinearities that "work very poorly," indicating the splines capture non-trivial dataset properties.

• More winning lottery tickets at initialization. Freezing everything except token embeddings, a 4-layer model went from 22.3% to 39.5% accuracy with tuned nonlinearities; a 12-layer model went from 23.0% to 35.2%. Fully trained models improved more modestly (86.2 → 87.0 and 88.4 → 89.7).

• Wider loss basins. Baseline models often improved sharply at the end of training, coinciding with the learning rate decaying to tiny values, whereas specialized models converged more smoothly. When checkpoint averaging replaced LR decay, the advantage of the specialized models grew much larger.

• Character-level tokenization appears to need very different biases. A character-level language dataset discussed in Appendix C showed no transfer at all to other language datasets.

Methodology in Plain English

The authors start from a strong baseline: a record holder from the NanoGPT speedrun, a GPT-2-style decoder-only transformer with rotary embeddings, U-Net-style skip connections, LayerNorm, QK normalization, and tuned initializations. They strip out components they consider too exotic or likely to interfere with modified activations (the Muon optimizer, mixed learning rates across layers, mixed precision, and an increasing attention-context schedule), and run each experiment on a single NVIDIA RTX 3090 with AdamW, a batch size of 1, and 4,096 tokens per step.

Then comes the key intervention. Every standard nonlinearity is replaced by a linear spline: a piecewise-linear function defined by 64 keypoints spread uniformly over the interval [−15, 15], whose heights are learnable. Splines are used rather than small neural networks because they impose no arbitrary prior — they can represent a sine wave as easily as a step function.

Training happens in two phases. In the warm-up, the spline values are initialized as constant functions and trained jointly with model weights on a source dataset, while 20 models run in a staggered schedule sharing the same splines (this prevents the weights and splines of a single model from overfitting to each other, and makes the result more consistent across seeds). In the second phase, the splines are frozen and transferred, and a fresh model is trained normally on a target dataset — making the setup directly comparable to a standard transformer except for those 128 transferred parameters. By crossing every source dataset with every target dataset, the authors build a matrix of which inductive biases suit which data.

For the mechanism analysis, they use two probes: scrambled-token warm-ups to test whether splines merely encode surface statistics, and heavily frozen models (only token embeddings trainable) to measure how many useful circuits exist at initialization.

Why This Matters

Impact on research. The paper pushes back on the prevailing convergence toward a small set of general-purpose architectures and optimizers. It shows that even fine-grained subsets of pretraining data — not just broad domains like vision or molecules — can benefit from different learning mechanisms, and that minor architectural changes can substantially affect learnability. It also frames this as a modern, practical instance of the no-free-lunch theorem, while acknowledging recent arguments that real-world data regularities limit that theorem's relevance.

Real-world applications (implied by the findings):

• Small and on-device language models, where data efficiency matters most and where the paper's 40–120M parameter regime is directly relevant.
• Domain-specific assistants in code, mathematics tutoring, or structured technical text, where a specialized activation choice could reduce the data or compute needed.
• Compute-constrained pretraining pipelines, where improving training efficiency without more data or parameters is a direct cost saving.
• Architecture search and hyperparameter tooling — the warm-up procedure is a cheap, continuous way to search a small but impactful slice of architecture space.

Industry relevance. The intervention is tiny and cheap: 64 parameters per spline, a short warm-up, and no change to the overall training recipe. For organizations training many models on distinct data mixtures, this suggests a low-cost lever beyond scaling. The paper is explicit, however, that current mixture-of-experts designs do not capture this benefit, since their experts share the same architecture and learning algorithm and MoEs primarily target inference cost rather than training data efficiency.

Future Directions

• Does the benefit survive scale? The authors state this as the key open question: their study covers only 40–120M parameter models, and they note prior work (Teney et al., 2026) showing learned activations can improve systematic generalization in ways that do not emerge with more data or parameters — but whether language-task gains hold for larger models is unresolved.

• Characterizing inductive biases. The field lacks descriptors to compare architectures beyond empirical evaluation — what, precisely, distinguishes a GeLU model from a ReLU one? Most complexity measures are described as too coarse to distinguish architectures.

• Matching inductive biases to data without trial and error. Prior work has done this in toy settings (Boolean functions, NLP classification tasks, world models); extending it to finer-grained dataset distinctions and learning methods remains open.

• Constructing learning methods with desired biases by design. Translating desired inductive biases into architectures, optimizers, and initializations — for example, reverse-engineering desirable neural tangent kernels into custom activation functions.

• Mixture-of-experts with heterogeneous experts. The paper suggests MoEs whose experts use different architectures or optimizers as one concrete route to combining specialized mechanisms inside a general-purpose model.

Target Audience

Researchers and engineers working on language model architecture, pretraining efficiency, or data-efficient training who already understand transformers and want evidence that architectural inductive biases interact meaningfully with data domain. It is also useful for practitioners building small or domain-focused models, and for anyone interested in the debate over monolithic versus specialized learning systems. Readers looking for downstream benchmark results or large-scale validation will not find them here — the paper evaluates perplexity and question-answering accuracy on small models and explicitly flags this as a limitation.

Authors’ abstract

Language models are typically trained with a single choice of architecture on diverse domains and data sources. Modern transformers have been extensively refined for such general-purpose learning, but are they equally effective across different subsets of the pretraining data? We investigate this question by training small language models from scratch on representative subsets of pretraining data such as web text, code, and mathematics. We introduce variants of transformers in which standard activation functions are replaced with learned splines tuned on specific datasets, providing a simple way to modify their inductive biases. Using a mix-and-match of architectures and datasets, we evaluate cross-domain transfer, i.e. how well does an architecture tuned for one domain perform on data from another? We find that the tuned transformer variants improve training efficiency or generalization on all seven of our tested datasets, but the gains are often domain-specific. For example, architectures tuned on a natural language dataset such as FineWeb perform well on TinyStories, whereas those tuned on mathematics perform poorly on natural language. We identify several mechanisms underlying these improvements, including wider loss basins and more abundant "lottery tickets". Our results show that different subsets of pretraining data can benefit from different learning mechanisms, highlighting a limitation of the prevailing trend towards a small set of general-purpose architectures and optimizers. This points to an opportunity to improve data efficiency by developing domain-specific learning strategies that retain the ability to train general-purpose models.