Can Transformers Really Do It All? On the Compatibility of Inductive Biases Across Tasks
Authors: Damien Teney
Liangze Jiang
Hemanth Saratchandran
Simon Lucey
Venue: ICLR 2026 Poster (2026)
Keywords: Transformers
language models
inductive biases
length generalization
activation functions
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: We optimize transformers by replacing the GeLUs/softmax with parametrized splines optimized on held-out data. This is used as a tool to reveals when task-specific architectures can be better than standard transformers.

Abstract: Transformers are remarkably versatile and their design is largely consistent across a variety of applications. But are they optimal for any given task or dataset? The answer may be key for pushing AI beyond merely scaling current designs.

**Method.** We present a method to optimize a transformer architecture for a given dataset, which we use as a tool to study optimal task-specific inductive biases. This method replaces the most important non-linearities (GeLUs,;softmax) with functions learned on held-out data. We then train the resulting architectures on other datasets, as a way to evaluate the compatibility between pairs of tasks.

**Findings.** On algorithmic toy tasks, we identify new architectures with dramatic improvements in learning speed, in- and out-of-distribution generalization, and stability across seeds. The new designs prove very task-specific however, and indicate that these tasks require inductive biases very different from those of standard transformers. On code and language modeling datasets, we also find architectures with consistent, yet smaller improvements. These designs transfer much better across datasets and domains (English & computer code).

**Implications.** Our results show that standard transformers are rarely a local optimum in the space of architectures. Simple alternatives can perform much better but sacrifice universality. This suggests that there may be room for improved architectures that better support multiple capabilities simultaneously, such as fluency and robust reasoning.

Published as a conference paper at ICLR 2026
CAN TRANSFORMERS REALLY Do IT ALL? ON THE COMPATIBILITY OF INDUCTIVE BIASES ACROSS TASKS
Damien Teney1.3 Liangze Jiang Hemanth Saratchandran Simon Lucey* Idiap Research Institute 2EPFL. Adelaide University
ABSTRACT
Transformers remarkably versatile and their design is largely consistent across a variety of applications. But are they optimal for any given task or dataset? The answer may be key for pushing AI beyond merely scaling current designs.
Method. We present method to optimize a transformer architecture for given dataset, which we use as a tool to study optimal task-specific inductive biases. This method replaces the most important non-linearities (GeLUs, softmax) with functions learned on held-out data. We then train the resulting architectures on other datasets, way evaluate the compatibility between pairs of tasks.
Findings. On algorithmic toy tasks, we identify new architectures with dramatic improvements in learning speed, in- and out-of- disttirbibttto generalization, and stability across seeds. The new designs prove very task- sasc-specffic however, and indi- cate that these tasks require inductive biases very different from those of standard transformers. On code and language modeling datasets, we also find architectures with consistent, yet smaller improvements. These designs transfer much better across datasets and domains (English & computer code).
Implications. Our results show that standard transformers are rarely local op- timum the space of architectures. Simple alternatives can perform much better but sacrifice universality. This suggests that there may be room for improved architectures that better support multiple capabilities simultaneously, such as flu- ency and robust reasoning.
INTRODUCTION
Inductive biases of transformers. The recent history of machine learning has seen a convergence of architectures across modalities. Most otate-of-the-art models for vision, language, and speech are based on transformers, barring only minor differences (Vaswani et al.) 2017). This success contrasts with earlier task-specific models, and it has prompted the hypothesis that transformers implement very generic inductive biases" well suited to various types of real-world data (Goldblum al., 2023] Teney al., 2024 2025). Considering the vast space of possible architectures, following question remains (QI).
Are transformers uniquely endowed with such generic inductive biases?
Uneven performance across domains. Transformers perform remarkably, for example as language models. But they also fail on elementary tasks such as learning arithmetic (Nikankin et al. 2024). These failures have motivated a plethora of alternative components such as positional encodings (Cai ct al., 2025 Jelassi et al. 2024) and attention mechanisms (Katharopoulos et al. 2020; Saratchan dran et al. 2024b Schlag al. 202 I). But these designs are rarely adopted beyond toy tasks. This suggests that the inductive biases of standard transformers may not be suited to domains as different natural language and arithmetic. This raises the following question (Q2).
Should we even seek to address such different domains with the same learning method?
The inductive biases of learning algorithm can be seen as prior over the space of functions MMitchell 1980 Mingard et al, 2021 such that particular (types of) functions are favored among the many that fit the data. We focus on biases encoded architectures, not in choices of optimizer, objective function, initialization,

---

Published as a conference paper at ICLR 2026
Transformer block (nx)
Decimal addition
MANO (Allen-Zhu et al)
Standard transformer
Input-
Multi-head
MLP
(..)
Output
Test accuracy
attention
Softaa
Gelle
Optimized architecture
Learned
non-linearities
Training steps (one curve per seed)
Figure 1: Our approach to discover better task-specific inductive biases. (Left) We replace the non-linearities in a transformer (softmax, GeLUs) with splines optimized from scratch for one cho- sen dataset. (Right) The new resulting architectures (with frozen splines as non-linearities) allow us to train models with dramatically better convergence, generalization, and stability across seeds, on algorithmic tasks and code/language modeling datasets. We also mix-and-match the architec- tures across tasks (not pictured) to evaluate the compatibility of inductive biases across tasks.
The above questions are relevant for developing better Al systems. Current designs have largely relied on scaling up models and data (Mayilvahanaan et al., 2025). This growth is not infinitely sustainable. and their deficiencies (e.g. in systematic generalization and data efficiency) indicate need for better inductive biases. Understanding the suitability of the inductive biases of transformers to specific tasks (Q1) and the compatibility between tasks (Q2) are steps in this direction.
Our approach. We address the above questions with a method that seeks better inductive biases for a chosen dataset by tweaking the transformer architecture. We replace non-linearities (GeLUS, softmax) with parametrized splines, optimized on held-out data. This yields new architectures that match or surpass standard transformers. The improvement in learning speed and/or generalization indicates how far the standard transformer is from local optimum in the space of architectures for specific task (Q1). We also mix mix-and-mater these new architectures across tasks to assess how the inductive biases tuned for one task perform for another, thus assessing their compatibility (Q2).
Findings. We study two domains: algorithmic skills and language modeling. For algorithmic skills, use toy tasks commonly used to evaluate architectures (see c.g. AAlennZZhu 2005). For all consid- ered tasks, our approach finds architectures that dramatically improve learning speed, generalization, and stability across random seeds (Section 3.1). Our task-specific variants of transformers are vastly superior to standard designs, using only minor modifications such as replacing the GeLUs. The cross-task evaluation also reveals that the new architectures are quite task-specific. This explains why many components proposed in the literature (e.g. attention mechanisms, positional encodings) are rarely adopted beyond toy tasks. These results also challenge the view that a single architecture can be optimal for vast set of tasks (Goldblum et al.) 2023).
For language modeling, we evaluate multiple datasets of natural language and computer code. In most cases, we find optimized architectures that slightly improve over a baseline transformer. We stress that these improvements are practically not directly useful, because standard components are more computationally efficient. But they matter indirectly, because they are evidence that standard transformers are neither a unique nor a local optimum in the space of architectures. In contrast to algorithmic tasks, the cross-task evaluation shows here that the improvements can transfer across natural language datasets, and across tokenization levels (character vs. subword). Overall, our results indicate that standard transformers are intrinsically better suited to modeling natural language than code, and clearly ill-equipped to learn algorithmic skills.
Our contributions are summarized as follows.
A method to optimize transformer architecture for any given dataset (Section 2). We replace GeLUs and softmaxes with parametrized components optimized on held-out data. The optimized architecture then be used with standard training evaluate its suitability to any other dataset. An application to algorithmic tasks (Section 3). We find that optimized architectures dramat- ically improve learning speed, generalization, and stability across seeds. They also prove very task-specific, showing the utility of inductive biases very different from standard transformers'.
An application to language modeling (Section 4). We obtain small, albeit consistent improve- ments, showing that standard transformers are neither unique nor optimal designs, even for com- mon code and natural language modeling tasks.
We discuss the implications for the development of future AI models in Section
2

---

solely for Published as a conference paper at ICLR 2026
2 PROPOSED METHOD TO OPTIMIZE ARCHITECTURES
Goal, As a baseline architecture, we consider a standard decoder-only transformer (GPT-2-style, see details in Appendix B). Our goal is evaluate whether this choice is optimal for specific datasets. We also seek better variants of transformers, i.c. identifying inductive biases better suited to each task. Evaluating the new architectures across tasks can then measure the compatibility of pairs of tasks. All the tasks we consider are formulated as sequence completion of natural language, computer code, or abstract tokens.
Replacing non-linearities with learnable parametrized functions. We replace the main non- linearities in a transformer with components that can be optimized (see[Figure 1). Indeed, the main difference between transformer and simple linear model hinges on few non-linear operations in the attention and MLP layers, which we will alter to obtain different inductive biases.
An MLP layer is defined as: x f W' 6(Wx w + b) + b' where æ is a vector of activations, W, W', b, b' learned weights and biases, and o: - R an element-wise non-linearity. In the baseline architecture, 6 is a GeLU. In our model, ΦDXLP is a ID linear spline parametrized by learnable keypoints MLP., capable of approximating a variety of functions (details in Appendix B).
An attention layer in the baseline transformer is defined as: x softmax QK)) V, where x is the output vector of activations and Q K, V are linear projections of the input. This is a special case of the kernel version of attention: < Σ}= KQQ, K,) V; -j=1 K K(Q,, K,) where the similarity between Q and K is measured with a kernel function K(Q, K). In the baseline transformer, Kamax (n,,axQQQ K) = @ K K/ In our model, we introduce learnable non- linearity o' : R + R giving K(Q, = =Q (((Q) o'(K). We implement ø' as a linear spline % with keypoints A that can be optimized.
Two-stage setting. Our experiments proceed in two stages. In stage I, we optimize the architecture for a chosen dataset by training both the model's weights and its parametrized non-linearities (0 MMPP) on D. In stage II, the non-linearities are frozen, and we retrain the model in standard manner from scratch on any dataset D'. The models obtained from stage are thus fairly comparable with the baseline architecture When ID' + ID, i.e. mmix-and-match setting, stage II serves to evaluate whether the inductive biases optimized for suit the learning of DD'.
Optimizing architectures. Our method may seem similar to prior work about learning activation functions (e.g. (Alexandridis et al.] (2025) but their goals are very different. These works seek to improve performance by continuously updating the activation during training. Whereas we seek to identify inductive biases that can remain hard-encoded in the architecture and further reused to train new models with other seeds and datasets (stage II). We make this possible with a two-loss training. During stage we hold out fraction of the training data (e.g. 20%) that optimizing the non-linearities, while we optimize the weights in a standard manner on the training set. This prevents ao-aaapatioon, that could cause the non-linearities to overfit particular weights or seed. This is particularly important for our experiments on algorithmic toy tasks, and even more so for improving length generalization (Section distribution (OOD) split of data (see Seccionn 3.1 such that the weights are optimized for one range of sequence lengths, and the architecture for a different wider range. This forces the architecture to capture an inductive bias for length generalization. In stage II, the non-linearities are frozen, and the model weights are trained in a standard manner on the whole training split of the target dataset.
((Secto 3.1) In this latter case, we hold out an out-of-
A second innovation to prevent the co-adaptation of weights and non-linearsties in stage I is multi- model training. We train M models in parallel (e.g. M = 4) that use different seeds but share the non-linearities being optimized. The resulting optimized architecture is naturally more likely to generalize in stage to other weights and datasets (see Appenddi Dj. This also proves particularly helpful for algorithmic tasks because the variance across seeds of the baseline architecture is often high. We provide a complete description of our method as Alloriitmm in the appendix.
In stage II, A0 OLLF) are frozen and hetter viewed as pre-tuned hyperparameters than extra model capacity. 'The henefit of the two-loss training is smaller for language modeling hecause the models are heavily over- parametrized and never risk of overfitting the training data.
3

---

Published as a conference paper at ICLR 2026
Rationale for splines. We parametrize our non-linearities as linear splines because they offer the most unbiased, tractable parametrization for an R- R function. For example, a spline can represent the identity function as easily as step function or a sine wave. Prior work on trainable activation functions uses e.g. small MIPs that enforce priors like smoothness or monotonicity (Apicella et al.) [222 Greydanus & Kobak 2020). Such parametrizations would struggle capture sharp transitions like those in Figure 6 We also favor linear splines over higher-order (c.g. cubic) ones because they
behave similarly while being much cheaper, as evaluated by Teney et al. 2025 Appendix D).
3 EXPERIMENTS ON ALGORITHMIC REASONING TASKS
In this section, we apply the proposed method to set of tasks commonly used to evaluate the al- gorithmic skills of transformers (Table I). These tasks seem elementary but they are remarkably challenging for transformers, and often used to highlight their limitations. The tasks all formu- lated as sequence completion: each example comprises an "input" part, followed by a separator, then an "output" part. The models are trained for next-token prediction the output part of training sequences. Unless otherwise noted we use i.id. sets of training, validation, and test data.
Experimental setup. For each task ID we first train the baseline architecture and tune its hyper- parameters (width, depth, learning rate, batch size, etc.) for high accuracy and fast convergence on the validation set. We then run the proposed method (stage I, M = 8) to optimize the architecture for D. We then re-train model from scratch with the optimized architecture (stage II), keeping the same hyperparameters (we get no further improvements by re-tuning them). In (Scction 3.2] we also re-train models on other tasks D' as way to evaluate the compatibility between ID and IV" and the generality of the optimized architectures. All results are averages over random seeds.
Table I: Algorithmic tasks used in our experiments. They are sized such that they require similar model capacity, except for MANO (Allen-Zhu 2025) which is intrinsically more complex.
Task
Examples
MEMORIZE, Simple memorization of mapping between two-integer key and integer value. with all integers in [1.32]. Each sequence consists of the key. a separator, and the value. This lusk has no set: performunce is simply training accuracy Zhong Andreas (2024]
10
11 32 27
18
PARENTHESES. Recognition of Dyck language. Fach sequence contains paremteses followed by separator and marker indicating whether they are balanced or not. Sequences leneths are [1,20] in the training set, and [21,40] in the validation and test sets [hoog & Amdreas 2024..
( [ Kunbalanced
(C) ] [ ) } <balanced>
[ ) ( <unbolancad>
AnDMon. Modular addition mod N, with 95% of the 2 examples used for training Zhong & Andreas 2024). We use N=97.
12 | 1.5
96 2
HAYSTACK. Needle-in-a-haye ha-lll-t--- recall. The model gets a sequence [ma, Ci ma, Ck2 mu] of markers my and values cq. It search for the first occurrence m. and return successor Cu (7hong & Andreas 2024. We use & E [1, 10] and my Cg E [1,64].
2 9 k 3 1
2
2 p 9 k b
ADD, Docimal addition of 4-digit mumbers with digit-wise tokens. Zhong & Andreas
2024.
1009 1 10
ADDREVERSED. ADD with reversed numbers, known to be easier to learn (Lee et al
2023
9.0 1 + 2
CopY. Repeating the input. Elementary but unsolved for length generalization CCai et al. 2025
2 2
Tokens in (1.8). Seg. longths in [2.10] for training, [2,16] validation. 16, for testing.
9 3
9
MANO. Synthetic task proposed by Allen-Zhu 2025 to evaluate large pretrained models. Each
(1+3)44 1:0
sequence specifies nested arithmetic operations mod N with number-leve tokens. Our scaled- down version uses N=7 and a number of operations per sequences in [1.3].
(2-(6-11)+3
(3x 61) 1
Buseline
Ours
Baseline
Ours
HAYSTACK
CoPY
ADDREVERSED
MANO
Figure 2: Training curves (test accuracy vs. training step, one curve per random seed) of models trained on algorithmic tasks with a baseline transformer or our optimized architectures. The latter converge much faster and show less variance across seeds. See See Appendixx Clfor other tasks.
4

---

Published as a conference paper at ICLR 2026
3.1 IMPROVEMENTS ON INDIVIDUAL TASKS
Faster convergence. The most striking improvement with optimized architectures is the learning speed (Figure 2). For the ADD and MANO tasks for example, convergence occurs 2 -33x faster. The learning rate of the baseline was tuned to maximum stable value for every task.
Reduced variance. On some tasks, baseline transformers show huge variance in accuracy and training speed across random seeds. This suggests tasks that are underspecified (Teney et al. [ 2022) and misaligned with the model's inductive biases (Zhou et al. 2024). In these cases, the optimized architectures eliminate the problem and make the training much more reliable (Figure 2).
Better generalization. For some tasks, baseline transformers do not reach perfect test accuracy though they perfectly fit the training data. This shows again a misalignment between the target func- tion and the inductive biases. Optimized architectures solve this problem (see e.g. MANO, Figure 2).
Improved length generalization. An outstanding challenge for transformers is the generaliza- tion to sequences longer than seen during training. Even the COPY task is unsolved and baseline transformer completely fails on unseen lengths (Figure 3). Among the plethora of existing partial solutions, the Alibi positional encodings (Press al., (002T) bring non-trivial accuracy slightly longer sequences. We use our method to optimize the Alibi architecture. We use the two-loss mech- anism of ofgoogrit Ioo optimize the transformer weights lengths 2-10 and the non-linearities on 2-15. This forces the optimized architecture to capture an inductive bias for length generalization. As a result, model trained with the optimized architecture reaches higher accuracies on longer sequences. While this is not complete solution to length generalization, shows that inappropriate inductive biases in the base architecture are one of the obstacles length generalization.
Baseline
accuracy
Press
2021
(sequence in shuding shows +/-1 std. dev.)
Alibi Ours losses
Alihi - Ours (full method)
Traning
lengtha
Sequence length
Figure 3: Length generalization on the COPY task. The baseline completely fails on unseen lengths (> 10). Alibi positional encodings (Press et al., (021) help. Optimizing the Alibi archi- tecture with our method further improves the accuracy and extends the benefits to longer sequences.
Performance with smaller models. We train models of different widths for each task. Results in Figure 4 show that the accuracy drops more sharply on some tasks with the baseline architecture than optimized ones. Intuitively, when the architecture is already aligned with the task, less capacity needed its weights. Equivalently, fixed number parameters offers more capacity.
Better accuracy on small models
No clear difference
Test accurney (%)
50
077
52
23
12
32
32
84
120
PARENTHESES
MEMORIZE
ADDREVERSED
HAYSTACK
ADD
ADDMOD
Figure 4: Test accuracy of models of different widths (X axi).. On some tasks, optimized architec-
tures maintain higher accuracy than the baseline when reducing the width of the model.
5

---

Published as a conference paper at ICLR 2026
3.2 COMPATIBILITY OF OPTIMIZED ARCHITECTURES ACROSS ALGORITHMIC TASKS
We now train models on each task using architectures optimized for any other task D' to evaluate the pairwise compatibility of their inductive biases. The results in Figure 5 show that the optimized architectures are very task-specific. Few of the benefits transfer across tasks, mostly across closely related tasks like ADD and ADDREVERSED. Many perform worse than a standard transformer. This shows that the specialization to our algorithmic tasks comes at the cost of universality, These tasks are very narrow however and remains an open question whether the negative impact is inevitable. A future step to study this question could be multi-task optimization in Algorithm
Parstthozd
Memorite
AddiMoo
Haysaak
Add
Copy
Memorize
-43
40
+33
-10
AddMod
-28
428
41A
Add
Architectures optimized for spocific tasks
Addfieverzed
Mano
19
+3
+5
+22
+22
+13
+17
Haystack
+1
+3
+12
+15
Copy
43
429
Parentheses
Worse
Target tasks
Figure 5: Compatibility of architectures across algorithmic tasks. We plot the absolute difference in test accuracy (%) with the baseline after a fixed number of steps (details in in/Appendix B). The best option per task (column) is usually the diagonal, meaning that the optimized architectures quite task-specific, while still yielding some positive transfer.
M ADD
MEMORIZE
PARENTHESES
ADDMoD
HAYSTACK
ADDREVERSED
COPY
MANO
Figure 6: MLP non non-linearities optimized for each algorithmic task.
Take-away. On algorithmic tasks, optimized architectures can dramatically outperform standard transformers, but the benefits are quite task-specific. This means that these tasks require inductive biases very different from those of standard transformers.
4 EXPERIMENTS ON LANGUAGE MODELING
We now apply the same experimental setup Section for computer code (English, Java) and natural language of various complexity levels (Table 2). Our goal is to understand whether different types of data benefit from different inductive biases. Current practices for building LLMs show that data diversity is beneficial (Longpre et al.) 2024) and that code complementary to natural language (Aryabumi et al., 2024 Petty et al., (0224). But because all kinds of data mixed during training, is unknown whether they could each exploit elicit different mechanisms in model. We also consider versions the datasets tokenized at the character subword level (BPE; details in Appendix B). These choices are motivated by Mayilvahanan al. who showed that L.LM performance is mostly determined by data diversity and tokenization.
language modeling. We use datasets
2025)
4.1
IMPROVEMENTS ON INDIVIDUAL DATASETS
TINYSTORIES. We compare in Figure models trained with baseline or optimized architectures. The latter do slightly better. The improvement is small but consistent at different model sizes. Train- ing curves (Figure 15) show that the improvement larger early during training then diminishes. We

---

Published as a conference paper at ICLR 2026
Table Datasets used in our experiments for language modeling (see Appendix B for details).
Dataset
Excerpt
TINYSTORIES,
Children stories generaled with GPT.3.5. was designed to capture core aspects natural language (syntax, coberence, compositionalit) with limited vocabulary. This allows smaller scale experiments than web scale open domain corpora (Eldan & La 2023
Once upor a time, there was clever little dog named Max loved to run (..)
Max.
Plays and sonnets by William Shakespeare, often used early research language modelling It includes recognizable patterns of grammnrr thythm, and vocahulary, as well unique structure because of the speaker labels and dialogue formatting (Karathy 2015).
SHAKESPEARE.
BENVOLIO: Good-morrow, cousin. ROMED; Is the day young? RENVOLAT:: But (.)
English Wikipedia (Mahoney 2006 We use clean version from Yong 2025) with only text visihle to human readers, without links and meta data. This data provides dense, real-world text with mix of vocabulary, syntax, and formatting.
EN RIK.,8 First 100 bytes
anarchiam originated as a term of abuse firat used againat eazly working (..
COnE SOACCEFRR NET--AAA--AVA & PYYTON.. Dataset of computer onde originally created to support research code search and code-text understanding (Husel 2019). We discard comments and descriptions in natural language following! La et nl. [022 to Tocus exclusively on code.
batch, "imit 100, self. it itarliti (..)
Transformer width
Optimized from scratch
Optimized from a GeLU
Retter
+0.3
Num. layers 40.5
40.6
Baseline
-22
+0.3 +0.2
+0.4
Figure 7: (Left) Absolute improvements in token prediction accuracy (%) of the best optimized architectures on TIORSIO compared to our baseline transformer. The accuracy is consistently slightly better at different model sizes. (Right) Visualization of MLP non-lineasssies optimized from scratch (results on the left) or from GeLU initialization (GeLU + Ours in Figure 8). Although they resemble generic wavelets, we show in Appendix D that fine details in these functions matter.
find best to optimize non-linearities only in MLPs (i.e. replacing GeLUs; see Figure 8). Replac- ing softmaxes with learned components barely matches or underperforms the baseline, indicating difficult optimization. We experimented with alternative parametrizations that exactly mimic a soft- max at initialization. This solution would barely move away from this initialization (not reported in tables), suggesting that a softmax is close to a local optimum.
We visualize in Figure 7 (right) the optimized MLP non-linearities, which are remarkably similar to sine wavelets. We evaluate a non-exhaustive selection of activation functions and attention variants from the literature in Table 3 None of them works better than ours. The gated linear units (GLUs) are a popular design that adds multiplicative interactions to the MLPs. We show that we can also improve them by introducing our learned spline in GLUS in lieu of their internal Swish activations. This provides similar improvements as over standard MLPs, cf. GLU/Swish and GL GLU/Ours in Table 3 We also evaluate in Alppeeddix D the importance of fine details in the learned non-linearities. We try to make them more periodic or symmetric, but they then always perform worse.
Table 3: Performance of models trained on TINYSTORIES with existing alternative attention and MLP designs (2 layers, width 256). None works better than ours. See SeelAppendix D for references.
Atlention
smax
smax
smax
smaa
smax
smax
smIa
smaa
smux
P1
F3 GeLU
Adaptive Gel.U
NornSmax GeLUU
MLP
Linear
GeLUU
Ours
GL.U/SwishhL.U/OO
ReLU
ReLU"
Tanll
Sinc
Gaussian
GeLL
Tr. perplexity
1.78
1.58
1.57
1.59
1.58
1.60
1.60
1.71
2.50
1.64
1.62
1.60
1.58
1.58
Val. acc. (%)
59.9
63.7
64.0
63.6
61.2
63.0
63.7
63.7
63.7
SHAKESPEARE & ENWIK8. These datasets differ from TinyStories in their richer vocabulary and sentence structure. SHAKESPEARF also follows a particular formatting presenting dialogues with speaker labels (see Table 2p. The results in Figures & 13 show that some optimized architectures slightly improve over the bascline. Optimizing non- loneaiicasssies in the MLPs is again more useful than in the attention. However, differences with the baseline are small, which suggests that standard transformers are inherently well-suited to language modeling.
The improvement is slightly clearer on character-level datasets than tokenized ones (marked -CHAR in (Fguree 8). We hypothesize that the target function to be learned by the transformer layers for leeva language modeling is more complex, because of the lesser capacity available
7

---

Published as a conference paper at ICLR 2026
TINYSTORIEE SHAKESPEARE:
SHAKESPEARE
SEWIK&-CHAR
ENWIKS
CSN-JAVA
CsN-PYTHON
It: perpleskyy
Attention:
Sunax
Ours
Ours
smax
MLP:
Linear
GeLU
GeLU Ours
GeLU
Ours
Ours
Figure 8: Perplexity on code and natural language (lower is better, numbers on bars correspond to the difference with the baseline architecture). Some optimized architectures perform slightly better than the baseline, often simply with optimized MLP non-linearities (.). Datasets of code (CSN- JAVA, CSN-PYTHON also benefit relatively more than datasets of natural language.
in the model's token embeddings (embeddings can otherwise make up significant fraction of the model parameters for tokenized datasets). This could be the reason why learned non-linearities are particularly helpful, since they can help learn and represent complex functions (Teney et al. 2025
We also evaluate a version of our optimized MLP non-linearities initialized as a GeLU rather than constant zero (GeLU Ours in Figure 8). With this, the model starts stage with a non-linearity known to perform well. And because the optimization is non-convex, the optimized solution remains in the local search space near GeLUs (see Figurre 7] right). The models trained with these non- linearities perform in-between GeLUs and those optimized from scratch. This means that GeLUs usually not an optimal solution, not even local one. But note also that best solutions not guaranteed to be globally optimal and better ones may exist.
CODESEARCHNET (CsN-JAVA, CSN-PYTHON), The results in Figure 8 show that our optimized non-lincarities in MLPs improve again over the baseline. The gains are larger for code than for natural language, relative to the gap between the baselines with linear and GeLU MLPs. These larger gains may reflect the larger importance of systematic structure and compositionality in code than natural language. The task of modeling code may thus resemble some the algorithmic tasks Seccinn 3 which benefited greatly from optimized architectures. Therefore, the architectures best suited to natural language may not be simultaneously optimal for code.
4.2 COMPATIBILITY OF OPTIMIZED ARCHITECTURES ACROSS LANGUAGE DATASETS
Our final results examine the compatibility the optimized architectures across language modeling datasets. We consider seven datasets plus MANO, the most complex of our algorithmic tasks. We train models for every task D using architectures optimized for any other task ID' The results in Figure 9 show that the variations across architectures very small. This contrasts with the results on algorithmic tasks (Figure 5). These optimized architectures thus encode much less taek-spec specialization. This suggests that the skills required across code and language modeling datasets much more uniform. We discuss the implications of these results in Section 6
Shakes Sharee char
TinyStories
Shakesseeae
emvild char
CauPython
chwold
Co-Jass
Mano
+23
+37
+350
+23
+13
+13
+15
Wona
Compatibility archi- tectures across code and language datasets (relative difference in per- plexity with the baseline in %, lower is better). The differences are much less dramatic than with algorithmic tasks (Figure 5), indicating smaller benefit dataset-specific specialization.
Figure 9:
TinyStories
-3
-13
-7
9
Shakespesre-char
Shakespeare
+1
122
Architectures optimized for specific tasks
enwikö char
-2
111
-8
9
12
enwikS
+4
+99
-28
+3
3
6
Con- Python
5
Better
Targot Lakks
8

---

Published as a conference paper at ICLR 2026
Take-away. For code and natural language modeling, the optimized architectures improve much less than for algorithmic tasks. This means that standard transformers are intrinsically closer to a local optimum in the space of architecture for these tasks, than for algorithmic skills.
RELATED WORK
Understanding inductive biases in NNs. Much of the prior on understanding neural net- works (NNs) has focused their simplicity bias, i.e. their preference for representing functions of low Kolmogorov (Zhou et al.] 2023) or spectral complexity (Bhattamishra et al. (0022). The simplicity bias depends primarily on the choice of activation function (Mingard al. 2019 Teney
et al. 2024), and its suitability was questioned (Domingos) 1999) by evaluating alternative activation functions MLPs (Teney al.j 2024). We extend this inquiry transformers and larger settings. In particular, we introduce a method to optimize non non-lineatt non-linearities in both attention and MLP layers, and apply to tasks relevant to the state of the art (code, natural language, algorithmic reasoning).
Improving transformers. Current LLMs all use very similar architectures, and Mayilvahanan et al. (2025) show that small design differences play little role in their performance. Prior work has however studied at length the impact of various components of transformers including their nonlin- earities (Jha & Reagen! 2025 Newhouse et al. 2025). Proposed improvements include alternative attention mechanisms (Katharopoulos al. 2020 Saratchandran et al. 2024b Schlag et al. 1
Tamayo-Rousseal et al. [2025] Veličković et al. (2024) and activation functions for MLPs (Hu et al,
2025 Teney al. 20255) and transformers (Mirzadeh al. 2023 So et al. 2021a). This motivates our work by suggesting that standard transformers are not a uniquely optimal choice of architecture.
Architecture search. Our method to optimize architectures is reminiscent of neural architecture search (NAS) (Goyal et al., 2019; Hong 2025] Liu et al. 2018 Manessi & Rozza 20018: Ramachan- dran al., 2018 Zoph & Le 2017). 7) The goals and approach are different though. NAS uses RL or evolutionary algorithms to search through pre-defined design choices. We directly use gradient descent to optimize relatively unrestricted parametrization of the non-linearities of transformers. Our goal is not to find better models (our designs are often computationally expensive). Instead, our method is tool to understand the compatibility of the inductive biases required for various tasks.
DISCUSSION
We have presented a method optimize a transformer architecture for specific datasets and used it to study the compatibility of inductive biases across tasks. We found that standard transformers are often suboptimal, but minor tweaks (replacing GeLUs and softmax operations) can substantially improve training speed, generalization, capacity, and stability across random seeds.
Our results show that different tasks benefit from different inductive biases, aligning with the no- free-lunch theorem (Wolpert 1996). Yet, transformers seem uniquely suitable to vast range of applications (Goldblum et al. 2023 our results can be seen as probing the limits of this hypothesis.
Architecture vs. scale. Prior work showed that the choice of architecture can become less important with scale (Bachmann et al. 2023 Tay et al. 2022). But this also means that the current need to build ever-larger models may be due to suboptimal inductive biases. In this work, tweaked the non-linearisies of transformers, but different inductive biases could be obtained by completely different means such initializations (Shinnick al. 2025) optimizers (Pascanu et al., 2025).
Do we need domain-specific models? Our results show a higher compatibility across language and code than algorithmic tasks. The latter are often used to highlight limitations of transformers. If they really represent desirable capabilities in LLMs, perhaps new architectures are required to combine capabilities for both language and algorithmic skills. A future step could be to apply our method to optimize architectures for multiple tasks simultaneously. It would also be interesting to evaluate whether the improved architectures transfer other domains such as vision and speech.

---

Published as a conference paper at ICLR 2026
Limitations. First, our search space of architectures is limited. Complex forms of attention (Hashemi et al. 2025) or interactions like gated linear units (GLU, Shazeer 2020) cannot be repre- sented in our formulation. Further gains possible with larger search space, but the optimization will also be more challenging. Second, the scale of our experiments is tiny relative to state-of-the- art LLMs.. The effects of different architectures may vanish with more data, but improving data effi ciency is a key objective of this line of work. So, effects at a small scale are intrinsically meaningful. Third, our architectures with optimized non-linearities can seem computationally costly. However, we describe efficient implementations in Appendix B and find that polynomial approximations of the learned splines make this a non-issue
Code is available at
REFERENCES
Konstantinos Panagiotis Alexandridis, Jiankang Deng, Anh Nguyen, and Shan Luo. Adaptive para- metric activation. In ECCV, 2025.
Zeyuan Allen-Zhu. Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers. SSRN Electronic Journal, May 2025.
Cem Anil, Yuhuai Wu, Anders Andreassen, Aitor Lewkowycz, Vedant Misra, Vinay Ramasesh, Am- brose Slone, Guy Gur-Ari, Ethan Dyer, and Behnam Neyshabur. Exploring length generalization in large language models. NeurIPS, 2022.
Andrea Apicella, Francesco Isgro. and Roberto Prevete. simple and efficient architecture for trainable activation functions. Neurocomputing, 2019.
Andrea Apicella, Francesco Donnarumma, Francesco Isgrò, and Roberto Prevete. A survey on modern trainable activation functions. Neural Networks, 2021.
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. NeurIPS, 2019.
Viraat Aryabumi, Yixuan Su, Raymond Ma, Adrien Morisot, Ivan Zhang, Acyr Locatelli, Marzieh Fadaee, Ahmet Üstün, and Sara Hooker. To code, or not to code? exploring impact of code in pre-training. arXiv:2408.10914, 2024.
Gregor Bachmann, Sotiris Anagnostidis, and Thomas Hofmann. Scaling MLPs: A tale of inductive bias. arXiv:2306. 13575, 2023.
Samuel James Bell and Levent Sagun. Simplicity bias leads to amplified performance disparities. In Proceedings of the 2023 ACM Conference on Fairness, Accountability, and Transparency, pp- 355-369, 2023.
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. arXiv:2211.123: 2022.
Garrett Bingham. William Macke, and Risto Mükkulainen. Evolutionary optimization of deep learn- ing activation functions. In Genetic and Evolutionary Computation Conference, 2020.
Ziyang Cai, Nayoung Lee, Avi Schwarzschild, Samet Oymak, and Dimitris Papailiopoulos. Ex- trapolation by association: Length generalization transfer in transformers. arXiv:2506.09551, 2025.
Irit Chelly, Shahaf E Finder, Shira Ifergane, and Oren Freifeld, Trainable highly-expressive activa- tion functions. In ECCV. Springer, 2024.
CodeGPT. Microsoft CodeGPT (available on HuggingFace). 2024. htttps:/huuginngface
Giacomo De Palma, Bobak Kiani, and Seth Lloyd. Random deep neural networks are biased towards simple functions. NeurlPS, 2019.
10

---

Published as a conference paper at ICLR 2026
Kamaludin Dingle, Chico Camargo, and Ard A Louis. outptt maps strongly biased towards simple outputs. Nature communications, 2018.
Pedro Domingos. The role of occam's s razor in knowledge discovery, Data mining and knowledge discovery, 1999.
Shiv Ram Dubey, Satish Kumar Singh, and Bidyut Baran Chaudhuri. Activation functions in deep learning: A comprehensive survey and benchmark. Neurocomputing, 2022.
Stanislas Ducotterd, Alexis Goujon, Pakshal Bohra, Dimitris Perdios, Sebastian Neumayer, and
Michael Unser. Improving lipschitz-constrained neural networks by learning activation functions. Journal of Machine Learning Research, 2024.
Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers on compositionality. NeurIPS, 36, 2023.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent english? arXiv:2305.07759, 2023.
Richard Futrell and Kyle Mahowald. How linguistics learned to stop worrying and love the language models. Annual Review of Linguistics, 2023.
Philip Gage, A new algorithm for data compression. C Users Journal, 12(2):23-38, 1994.
Gallant. There exists neural network that does not make avoidable mistakes. In IEEE International Conference on Neural Networks. IEEE, 1988.
Robert Geirhos, Jörn-Henrik Jacobsen, Claudio Michaelis, Richard Zemel, Wieland Brendel, Matthias Bethge, and Felix A Wichmann. Shortcut learning in deep neural networks. Nature Machine Intelligence, 2020.
Micah Goldblum, Mare Finzi, Keefer Rowan, and Andrew Gordon Wilson. The no free lunch theorem, kolmogorov complexity, and the role of inductive biases in machine learning. arXiv:2304.0533 05366, 2023.
Anirudh Goyal and Yoshua Bengio. Inductive biases for deep learning of higher-level cognition. Proceedings of the Royal Society A, 2022.
Mohit Goyal, Rajan Goyal, and Brejesh Lall. Learning activation functions: A paradigm of understanding neural networks. arXiv:/906.09529, 2019.
Sam Greydanus and Dmitry Kobak. Scaling down deep learning with MNIST-ID, arXiv preprint arXiv:2011.14439, 2020.
Michael Hahn and Mark Rofin.
Why are sensitive functions hard for transformers?
arXiv:2402.09963. 2024.
Michael Hahn, Dan Jurafsky, and Richard Futrell. Sensitivity as a complexity measure for sequence classification tasks. Transactions of the ACL, 2021.
Baran Hashemi, Kurt Pasque, Chris Teska, and Ruriko Yoshida. Tropical attention: Neural algorith- mie reasoning for combinatorial algorithms. arXiv:2505.17/90. 2025.
Dan Hendrycks and Kevin Gimpel. Gaussian enror linear units (GeLUs), arXiv:1606.08415. 2016.
Katherine L Hermann and Andrew K Lampinen. What shapes feature representations? exploring datasets, architectures, and training. arXIv:2006.12433. 2020.
Jinwook Hong. Asnn: Learning suggest neural architectures from performance distributions.
arXiv:2507.20164, 2025.
Sara Hooker. The hardware lottery. CACM, 64(12):58-65 2021. 2021.

---

Published as a conference paper at ICLR 2026
Leyang Hu, Matteo Gamba, and Randall Balestriero. Curvature tuning: Provable training- ffrenngg-rreeeee model steering from a single parameter. arXiv:2502.07783 2025.
Hamel Husel, Ho Ho-Hsiang Wu, Tiferet Gazit, Miltiadis Allamanis, and Mare Brockschmidt. Code- searchnet challenge: Evaluating the state of semantic code search. arXiv:1909.09436, 2019.
Ameya D Jagtap and George Em Karniadakis. How important are activation functions in regression and classification? a survey, performance comparison, and future directions. Journal of Machine Learning for Modeling and Computing, 4(1), 2023.
Ameya D Jagtap, Kenji Kawaguchi, and George Em Karniadakis. Adaptive activation functions accelerate convergence in deep and physics-infommed neural networks. Journal of Compulational Physics, 2020.
Samy Jelassi, David Brandfonbrener,, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models at copying. arXiv:2402.01032. 2024.
Nandan Kumar Jha and Brandon Reagen. arXiv:2501.03499 2025.
Entropy-guided attention for private llms.
Zixuan Jiang, Jiaqi Gu, and David Z Pan. Normsoftmax: Normalizing the input of softmax to accelerate and stabilize training. In IEEE International Conference on Omni-layer Intelligent Systems (COINS). IEEE, 2023.
Keller Jordan, Jeremy Bernstein, Brendan Rappazzo, @fernbearbssky.social,, Boza Vlado, You Jiacheng, Franz Cesista, Braden Koszarsky, and @Grad62304977. modded-nanogpt: Speedrunning the nanogpt baseline, 2024. URL httph::tttsssie https://gttubb.cmm//eeleeooorrt
paadedd-annopt
Dimitris Kalimeris, Gal Kaplun, Preetum Nakkiran, Benjamin Edelman, Tristan Yang, Boaz Barak, and Haofeng Zhang. SGD on neural networks learns functions of increasing complexity. NeurIPS, 2019.
Andrej Karpathy. The unreasonable effectiveness of recurrent neural networks.
http://
2015.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In ICML. PMLR, 2020.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. NeurIPS, 2023.
Jan Kukačka, Vladimir Golkov, and Daniel Cremers. Regularization for deep learning: A taxonomy. 1xXiv:110.11066 2017.
Nayoung Lee. Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv:2307.03381, 2023.
Chenxi Liu, Barret Zoph, Maxim Neumann, Jonathon Shlens, Wei Hua, Li-Jia Li, Li Fei-Fei, Alan L. Yuille, Jonathan Huang, and Kevin Murphy. Progressive neural architecture search. In 15th European Conference on Computer Vision (ECCV), 2018.
Ziming Liu, Yixuan Wang, Sachin Vaidya, Fabian Ruchle, James Halverson, Marin Soljačie, Thomas Y Hou, and Max Tegmark. Kan: Kolmogorov-aroold networks. arXiv:2404.19756 2024.
Shayne Longpre, Gregory Yauney, Emily Reif, Katherine Lee, Adam Roberts, Barret Zoph, Denny Zhou, Jason Wei, Kevin Robinson, David Mimno, et al. A pretrainer's 's guide to training data: Measuring the effects of data age, domain coverage, quality, & toxicity. In Proceedings of the Conference of the North American Chapter of the Association for Computational Linguistics, 2024.
12

---

Published as a conference paper at ICLR 2026
Shuai Lu, Daya Guo, Shuo Ren, Junjie Huang, Alexey Svyatkovskiy, Ambrosio Blanco, Colin Clement, Dawn Drain, Daxin Jiang, Duyu Tang, et al. CodeXGlue: A machine learning bench- mark dataset for code understanding and generation. arXiv:2102. 04664, 2021.
Kaifeng Lyu, Zhiyuan Li, Runzhe Wang, and Sanjeev Arora. Gradient descent on two-layer nets: Margin maximization and simplicity bias. NeurlPS, 2021.
Andrew L Maas, Awni Y Hannun, Andrew Ng, et al. Rectifier nonlinearities improve neural network acoustic models. In ICML, 2013.
Matt Mahoney, enwik8: First 100m bytes of english wikipedia. Available at https://
mat mattmahoney net dtte/////eettaaaaa httm 2006.
Franco Manessi and Alessandro Rozza. Learning combinations of activation functions. In Interna- tional Conference on Pattern Recognition (ICPR), 2018.
Prasanna Mayilvahanan, Thaddäus Wiedemer, Sayak Mallick, Matthias Bethge, and Wieland Bren- del. Lims on the line: Data determines loss-to-loss scaling laws, arXiv:2502. 12120, 2025.
Raphaël Millière. Language models as models of language. Mind & Language, 2024.
Chris Mingard, Joar Skalse, Guillermo Valle-Pérez, David Martinez-Rubio, Vladimir Mikulik, and Ard A Louis. Neural networks are priori biased towards boolean functions with low entropy. arXiv: 11522, 2019.
Chris Mingard, Guillermo Valle-Pérez, Joar Skalse, and Ard A Louis. Is SGD a bayesian sampler? well, almost. Journal of Machine Learning Research, 2021.
Chris Mingard, Henry Rees, Guillermo Valle-Pérez, and Ard A Louis. Do deep neural networks have an inbuilt occam's razor? arXiv:2304.0660 066700, 2023.
Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo C Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mchrdad Farajtabar. Relu strikes back: Exploiting activation sparsity in large language models. 1X.1:23100005644 2023.
Tom M Mitchell. The need for biases in learning generalizations. Rutgers University CS tech report CBM-TR-117, 1980.
Aaron Mueller and Tal Linzen. How to plant trees in language models: on the emergence of syntactic inductive biases. Transactions of the ACL, 2022.
Data and architectural effects
Laker Newhouse, R Preston Hess, Franz Cesista, Andrii Zahorodnii, Jeremy Bernstein, and Phillip Isola. Training transformers with enforced lipschitz constants. arXivv2507.13338, 2025.
Yaniv Nikankin, Anja Reusch, Aaron Mueller, and Yonatan Belinkov. Arithmetic without algo- rithms: Language models solve math with bag of hcuristics. arXiv:2410.21272, 2024.
Isabel Papadimitriou and Dan Jurafsky. Injecting structural hints: Using language models to study inductive biases in language learning. Proceedings ACI., 2022.
Razvan Pascanu, Clare Lyle, Ionut- Vlad Modoranu, Naima Elosegui Borras, Dan Alistarh, Petar Velickovic, Sarath Chandar, Soham De, and James Martens. Optimizers qualitatively alter solu- tions and we should leverage this. arXiv preprint arXiv:2507. /2224, 2025.
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra. Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. NewrIPS, 2024.
Jackson Petty, Sjoerd van Steenkiste, and Tal Linzen. How does code pretraining affect language model task performance? arXiv:2409.04556, 2024.
Mohammad Pezeshki, Oumar Kaba, Yoshua Bengio, Aaron C Courville, Doina Precup, and Guil- laume Lajoie. Gradient starvation: A learning proclivity in neural networks. NewrIPS, 2021.
13

---

Published as a conference paper at ICLR 2026
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. 12409, 2021.
Aahlad Manas Puli, Lily Zhang, Yoav Wald, and Rajesh Ranganath. Don't blame dataset shift! shortcut learning due to gradients and cross entropy. NeurIPS, 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Nasim Rahaman, Aristide Baratin, Devansh Arpit, Felix Draxler, Min Lin, Fred Hamprecht, Yoshua Bengio, and Aaron Courville. On the spectral bias of neural networks. JCML. PMLR, 2019.
Prajit Ramachandran, Barret Zoph, and Quoc V Le. arXiv:1710.05541, 2017.
Swish:
self-gated activation function.
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions. In ICLR, 2018.
Riccardo Rende, Federica Gerace, Alessandro Laio, and Sebastian Goldt. A distributional simplicity bias in the learning dynamics of transformers. NeurlPS, 37, 2024.
Vishwanath Saragadam, Daniel LeJeune, Jasper Tan, Guha Balakrishnan, Ashok Veeraraghavan, and Richard G Baraniuk. Wire: Wavelet implicit neural representations. In CVPR, 2023.
Hemanth Saratchandran, Sameera Ramasinghe, Violetta Shevchenko, Alexander Long, and Si- mon Lucey. A sampling theory perspective on activations for implicit neural representations. arXiv:2402.05427, 2024a.
Hemanth Saratchandran, Jianqiao Zheng, Yiping Ji. Wenbo Zhang, and Simon Lucey. Rethinking softmax: Self-attention with polynomial activations. arXiv:2410./880 2024b.
Simone Scardapane, Steven Van Vaerenbergh, Simone Totaro, and Aurelio Uncini. Kafnets: Kernel- based con-parametric activation functions for neural networks. Neural Networks, 2019.
Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers. In ICML. PMLR, 2021.
Noam Shazeer. Glu variants improve transformer. arXiv:2002.05200 2020.
Zachary Shinnick, Liangze Jiang, Hemanth Saratchandran, Anton van den Hengel, and Damien Teney. Transformers pretrained on procedural data contain modular structures for algorithmic reasoning. arXiv preprint arXiv:2505.223085 2025.
David So, Wojeiech Manke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc Le. Searching for efficient transformers for language modeling. NeurlPS, 34, 2021a.
David So, Wojeiech Manke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc V Le. Primer: Searching for efficient transformers for language modeling, 2022. arXiv:210900666, 2021b.
Leon René Sütfeld, Flemming Brieger, Holger Finger, Sonja Füllhase, and Gordon Pipa. Adaptive blending units: Trainable activation functions for deep neural networks. In Intelligent Computing: Proceedings of the Computing Conference. Springer, 2020.
Remi Tachet, Mohammad Pezeshki, Samira Shabanian, Aaron Courville, and Yoshua Bengio. On the learning dynamics of deep neural networks. arXiv: / 1809.088489 2018
Camilo -ouaye-Ruusseau, Yunjia Zhao, Yiqun Zhang, and Randall Balestriero. Your attention matters: to improve model robustness to noise and spurious correlations. arXiv:2507.20453. 2025.
Yi Tay, Mostafa Dehghani, Samira Abnar, Hyung Won Chung. William Fedus, Jinfeng Rao, Sharan Narang, Vinh Q Tran, Dani Yogatama, and Donald Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling? arXiv:2207. 10551, 2022.
14

---

Published as a conference paper at ICLR 2026
Damien Teney, Ehsan Abbasnejad, Simon Lucey, and Anton den Hengel. Evading the simplic- ity bias: Training diverse set models discovers solutions with superior ood generalization. arXiv:2105.05612, 2021.
Damien Teney, Maxime Peyrard, and Ehsan Abbasnejad. Predicting is not understanding: Recog- nizing and addressing underspecifieation in machine learning. In ECCV. Springer, 2022.
Damien Teney, Armand Mihai Nicolicioiu, Valentin Hartmann, and Ehsan Abbasnejad. Neural redshift: Random networks are not random functions. In CVPR, 2024.
Damien Teney, Liangze Jiang, Florin Gogianu, and Ehsan Abbasnejad. Do we always need the simplicity bias? looking for optimal inductive biases in the wild. In CVPR, 2025.
Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the parameter-function map is biased towards simple functions. vrXiv:8005.0830 2018.
Bhavya Vasudeva, Deqing Fu, Tianyi Zhou, Elliott Kau, Yougi Huang, and Vatsal Sharan. Simplicity bias of transformers learn low sensitivity functions. arXiv:2403. 06925, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurIPS, 2017.
Petar Veličković, Christos Perivolaropoulos, Federico Barbero, and Razvan Pascanu. softmax is enough (for sharp out-of out-of-distribution)) arXiv:2410.01100 2024.
Alex Warstadt and Samuel R Bowman. What artificial neural networks can tell us about human language acquisition. Trends in Cognitive Sciences, 2020.
Colin White, Mahmoud Safari, Rhea Sukthanker, Binxin Ru, Thomas Elsken, Arber Zela, De- badeepta Dey, and Frank Hutter. Neural architecture search: Insights from 1000 papers. arXiv:2301.08727, 2023.
David H Wolpert. The lack of a priori distinctions between learning algorithms. Neural computation, 1996.
David H Wolpert. The supervised learning no-free-lunch theorems. Soft computing and industry: Recent applications, 2002.
Zhi-Qin John Xu, Yaoyu Zhang, Tao Luo, Yanyang Xiao, and Zheng Ma. Frequency principle:
Fourier analysis sheds light on deep neural networks. arXiv:1901.06532 2019.
Xiulin Yang et al. (anything) goes? crosslinguistic study of (im)possible language learning in Ims. In Proceedings of ACL, 2024.
York Yong. Clean version of the wikipedia. Available at /
vl pt//aataaett//oommaa///yrryrooot
zip 2025.
Ailing Zeng, Muxi Chen, Lei Zhang, and Qiang Xu. Are transformers effective for time series forecasting? In AAAI, 2023.
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin John Xu. Initial- ization is critical to whether transformer fits composite function by inference or memorizing. arXiv:2405.00 05409, 2024.
Ziqian Zhong and Jacob Andreas. arXiv:2410.04363 2024.
Algorithmic
capabilities of random transformers.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? a study in length generalization. 12Xi::2310.16223 2023.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371, 2024.
Barret Zoph and Quoc V. Le. Neural architecture search with reinforcement learning. In ICLR, 2017.
15

---

Published as a conference paper at ICLR 2026
A ADDITIONAL RELATED WORK
Inductive biases in deep learning are due to choices of architecture (Goyal & Bengio 2022 and
of the learning algorithm (optimizer, objective, regularizers Kukačka et al. 2017). We focus on the former. The simplicity bias has been studied from both aspects. Most explanations attribute to loss functions (Pezeshki et al. 202I) and gradient descent (Arora al. 2019; Hermann & Lampinen 2020 Lyu et al. 2021] Tachet et al. 2018). But work on untrained networks shows that it can be explained with architectures alone (De Palma al., 2019; Goldblum et al. 2023] Mingard et al., 2019 Teney et al. 2024] Perez et al. 2018). Teney et al. al.] (2024) showed that the choice of al., 2019 Kalimeris et al., (2119) or frequency principle (Xu al. (019) is a related but different effect related to training dynamics: NNs approximate low-frequency components of the target function earlier during training with SGD.
activation function can modulate the simplicity bias. The spectral bias (Rahaman
Simplicity bias in transformers. The hypothesis of a simplicity bias in NNs has also been studied specifically in transformers. |Hahn et al.] (2021) shows that common models in NLP are biased to learn low-sensitivity functions. Bhattamishra et al. (2022) shows that transformers are more biased for simplicity than LSTMs. DDiri et al. (l.(()))) examine large pretrained models and determine that that tend rely on shortcut learning on simple reasoming tasks. Zhou et al. (2023) focus on length generalization and show that transformers learn the shortest program in the RASP language that fits the training data -a specific form of the simplicity bias. Rende et al. (2024) study BERT-like models and find that they learn simple functions first during the course of training. Zhang et al. (2024) find
that the scale of initialization can influence a transformer's learning of a generalizing or memorizing solution. Vasudeva et al. (2024) further study the bias of transformers for learning low-sensitivity functions using the NTK theory. HHahn & Rofin (2024) show that sensitive functions are hard to learn for transformers because they correspond to sharp solutions in their optimization landscape as side-effect of the simplicity bias.
Activation functions are key for introducing son-linearities in NNs. Many options were considered early on, e.g. sine activations in the Fourier Neural Networks from 1988 (Gallant 1988). ReLUs are often credited for enabling the rise of deep learning by avoiding vanishing gradients (Maas et al.) 2013). However they are also essential in inducing the simplicity bias (Teney et al. 2024) which may be just important. The research community has slowly converged towards smooth handcrafted variants of ReLUs such as GeLUs (Dubey et al., [2022] Hendrycks & Gimpel 2016: Ramachandran let al. 2017). Some works proposed learning activation functions using extra parameters optimized
alongside the weights of the network (Alexandridis et al. [025 (Apicella
| et al. 2019; 2021 ( [Bingham
let al. 2020; Chelly et al.] 2024 Ducotterd et al.] 2024 Jagtap al. 2020 [Sarrdapane et al.) [019]]
Sütfeld et al. 2020). See Jagtap & Karniadakis 2023) for a comprehensive review. The goal is to
better fit the training data with an activation function that can evolve during training. In contrast, we use meta learning to find an activation function that induces better inductive biases, such that training with this fixed activation provides better generalization. This requires bi-level optimization, episodic training. and unbiased parametrization that allows to learn activations very different from existing ones. Kolmogorov-Arnold Networks (Liu et af., 2024) parametrize the connections in a NN, which is equivalent to learning different activation functions across channels and layers. They use parametrization as splines similar to ours. Their benefits in physics-related problems likely result from the alterations to the inductive biases studied in this paper. Our method differs from neural architecture search (White et al. 2023) in its ability to discover novel activation functions from scratch, rather than selecting from predefined candidates (Sütfeld al. 2020) or from narrow set of parametric functions (Alexandridis et al. 2025).
Length generalization refers to the ability of a model to generalize to sequences longer than seen during training, especially for algorithmic tasks (e.g. arithmetic operations numbers with more digits). This remains a challenge despite extensive work on positional encodings, which only par- tially address the problem (Anil et al. 2022 Kazemnejad et al. 2023 Zhou et al. 2024). This paper shows that other aspects of the architecture can be important. We use the COPY task as proof of concept and show that different MLP activation functions can bring a significant improvement the existing Alibi encodings (Press et al. 2021).
16

---

Published as a conference paper at ICLR 2026
Learnability and inductive biases. The learnability of any given task is fundamental question in machine learning. It is well know that inductive biases are indispensable for generalization to unseen data (Mitchell 1980) and that learning algorithm is universally useful, as per one of the no-free lunch theorems [Wolpert ).02)). Meanwhile, neural networks have nevertheless proved widely successful. The broad applicability of transformers, in particular, suggests that their inductive bias has a broad relevance to real-world data (Goldblum et al.) (023). The simplicity bias is broad and vague characterization of these properties. Various studies have established however that
the simplicity bias is not universally beneficial (Domingos) 1999. [Teny et al. 2025 Zeng et al.
2023) and even responsible for failure cases such as shortcut learning (Geirhos et al. 2020 Puli
et al. 2023 Teney et al. 2021) the amplification of biases and performance disparities (Bell & Sagun 2023 Even the underlying principle supporting the simplicity bias, known as Occam's razor, has long been debated in the philosophical literature because it lacks a justification from first principles (Mingard et al., 2023 Appendix A). A prominent argument for simplicity is rooted in algorithmic information theory (Dingle et al. 2018) with results stating essentially that "a bias in the distribution of target functions must be towards low complexiry". However, this only means that simplicity is good prior average, but not necessarily the best choice for any task dataset.
Studies in linguistics and cognitive science have also examined the question of learnability. This includes studies the influence of architectures and data generalization during language ac- quisition, both for humans and machines (Futrell & Mahowald 20233 Millière 2024 Warstadt & Bowman 2020). This explains how syntactic and structural biases arise and how they can be con- trolled (Mueller & Linzen 2022, Papadimitriou & Jurafsky 2022, Yang et al. 2024). Our paper complements this line work since it helps clarify the impact of architectures on generalization. Our approach is quite different though. Our method allows searching through the space of archi- tectures via the optimization of lineaiiesss This matters because current popular designs (e.g. transformers) are contingent on external factors (see the Hardware Lottery, Hooker 2021).
Connection with prior work. This paper is a follow-up the study by Teney et al. (2025) that uses trainable non-linearities to study whether the simplicity bias of standard neural architectures always desirable. was however limited to MLPs and toy data, and relied on an expensive optimization method unsuitable to modern architectures. In comparison,
main innovations are: a formulation of trainable non-linearities that applies to transformers' MLPs and attention layers;
a tractable optimization method replacing the expensive bi-level approach from prior work;
the study of mainstream domains (language modeling, algorithmic reasoning);
the study of cross-task compatibility, whereas prior work focuses individual datasets; a demonstration of massive improvements on algorithmic tasks;
17

---

Published as a conference paper at ICLR 2026
B IMPLEMENTATION DETAILS
Proposed method. We provide formal deseription of our method in Algorithm
Algorithm Proposed method (stage 1) to optimize a transformer architecture for specific task. Input:
Training data D = {s as token sequences $ E S.
Baseline architecture T instantiable as next-token prediction model Te: S S of weights 0.
-): Loss function.
a: Fraction of held-out data. M: Number of parallel models.
Method:
Define new architecture Tea ONnL by replacing
softmaxes with Σ, K( K2) V, m, ((Q, Ký), where K(Q, K) co, (Q) do, (K).
GeLUs with linear spline 00NLPT
where architecture hyperparameters A and MLLP specify the value of splines @ at their keypoints. Instantiate M untrained models of architecture 7 as Fy, TM
Split into Darh and Dwaa of sizes a n and (1-a))n.
while not converged SGD training loop
Sample mini-batch D from Darch and D DM from Dwts
Eval. loss of each individual model its own data D L +- Ee EDM L(Tmm (s), Eval. combined loss of all models together on DO: Larch &- Em Ea EDU L L(fom, @m (8), Update weights of each model: Vm, 6. SGD(O Vol.
Update architecture: (0x, O M + SGD((OA, OMLP), V (arn
(0 M)P) + x0 0MLP).
Output: optimized architecture Te, 0
Now T can be used like any other architecture, treating 0% and 0% MIL as fixed hyperparameters.
Baseline transformer architecture. Our baseline is GPT-2-style architecture (Radford et al, 2019). uses standard multi-head attention, GeLU activation functions in the MLPs, post-norm layers, learned absolute positional embeddings, and width multiplier of 4 in the MLP hidden layers. All weights are initialized from Gaussians of standard deviation 0.02 truncated 2 standard deviations.
Parametrization of non-linearities linear splines. We want search space free of priors such as the smoothness and monotonicity enforced in similar work on the learning of activation functions (e.g. /Apicella et al. 22019;; [Chelly et al. (2024).. We therefore choose to learn non-linearity as a linear spline dgg + R with control points defined by O. We define nee points spread regularly in interval (a, bj. typically ne 122 points in [-20, +20] for spacing of 1/3 between points (see hyperparameters in Table Then o represents piecewise linear segments interpolating values specified in the learned parameters ;= (a). e (b))] R". The function can represent simple and complex functions, including smooth curves, periodic functions, sharp transitions, etc.
Datasets for algorithmic tasks. For most tasks, we generated data with code adapted from Zhong com j2222022220/22z2 ranoomm trans fo rrerrs While this prior work generates some of the data on-the-fly, pre-generate all the data ensure that the training/validationntsss splits are strictly disjoint.
& Andreas (2024):
For MANO, re-implemented the data generation based on the description by Allen -hu 2025
Compared to this prior work, we scaled down the task to allow using smaller models. We generated 1c5 training examples, with number of operations in each sequence in [1,3], modulus of and without tokens signaling the number of operations.
For all algorithmic tasks,
of 1e3 examples, strictly disjoint from the training
18

---

Published as a conference paper at ICLR 2026
Datasets for language modeling. For datasets tokenized the character level, every character or symbol in the data simply corresponds to one token. For the TINYSTORIES, SHAKESPEARE, and ENWIK8 datasets tokenized at the subword level, use the byte-pair encoding (BPE, Gage (994)) tokenizer from GPT-2 Radford et al. (2019). We consider it consistent choice suitable to our different datasets since it originally trained on very diverse data. For the CORESEARCHNET datasets, we use the tokenizer of the CodeGPT model (CodeGPT) 2024). For each dataset, we discard tokens with fewer than 200 training occurrences. This significantly reduces the vocabulary size and training costs. This should not undermine the results of our experiments: if anything, including more rare tokens could reveal larger differences across datasets.
Metries. For the algorithmic tasks, we measure performance as the token- wise accuracy of the "out- put" part of the generated sequences (the same part of the sequences as used to compute the training loss). This allows finer-grained evaluation of partial success than the sequence-wise wise accuracy.
For the CoPy task, we use the sequence-wir wise accuracy because the token- coken-wise accuracy can remain falsely high when model fails at length generalization.
For the language modeling tasks, we measure performance using the training perplexity (exponential of cross-entropy loss) well as token-wise accuracy on validation data as more intuitive measure of performance. For the accuracy, we measure on the latter half of the context window to ensure that we evaluate predictions with enough conditioning.
For the compatibility across algorithmic tasks (Figure 5), we plot the difference in test accuracy with the baseline after fixed number of steps. We adapt the number of steps to each task to capture improvements in generalization and/or training speed depending on the task. This is because both the baseline and optimized architectures saturate perfect accuracy for multiple tasks, hence the final accuracy alone is not informative.
MEMORIZE: 150 steps.
PARENTIESESS: 300 steps.
ADDMOD: 300 steps.
HAYSTACK: 400 steps.
ADD: 700 steps.
ADDREVERSED: 350 steps. COPY: 2,000 steps.
MANO: 3,000 steps.
Hyperparameters. We tuned the hyperparameters in Table 4 for standard transformer on each task, to make sure that our optimized architectures are compared against strong baselines. For example, we use the Canon layers proposed by AlllenZZhu (2025) for many tasks (sequence-wise ID convolutions) because they clearly improve the performance of the baseline.
Table 4: Hyperparameters used for each task.
MEMORIZE
PARENTHESES ADDMOn HAYSTACK
ADD AnDREVERSED CoPY MANO| Language datnsets
Num. layors
2
Num. alt, heads
2
2
2
2
4
Width
32
32
32
1
128
128
128
128
128
|
512
Tied embeddings
No
Yos
Canon layers
No
Yos
Num. It. steps
500
500
1,000
1,000
11,000
1.000
2,000
5.000
3,000
Peak LR
.005
.001
.02
.001
.001
.001
.004
.001
.001
Optimizer
Adam
Adam (18,82))
(0.9. 0.999)
(0.92, 0.98)
(0.9, 0.999)
LR schedule
5% linear warm-up. 50% cosine cool-down (act pecessary for algorithmic tasks; used on all tasks for consistency)
Weight decay
(betler on all tasks than using any weight decay)
Drupout
Parallel models M
8
Spline range [e,b]
20, 20
Spline spacing n4
1/9
1/3
19

---

Published as a conference paper at ICLR 2026
C ADDITIONAL RESULTS ON ALGORITHMIC TASKS
Training curves. Figure 10 shows that the optimized architectures (2nd and 3rd columns) always converge significantly faster than baseline transformer (Ist column) and show less variance across seeds. There is little difference between the 2nd and 3rd columns, which means that most of the benefits come from optimizing the non-linearity within the MLP layers rather than the attention.
Attention:
Ours
GeLU
MEMORIZE
PARENTHESES
ADDMOD
1100
1100
100
ADD
ADDREVERSED
COPY
MANO
5200
Figure 10: Training curves (test accuracy vs. training steps, one curve per seed) of models trained on algorithmic tasks with bascline transformer (first column) or optimized architectures (second and third columns).
20

---

Published as a conference paper at ICLR 2026
Compatibility of architectures across algorithmic tasks. We present below the full results fol- lowing the format of Figure 5] We show the effect when optimizing the non- con-linearities in MLP or attention layers, or both. Optimizing the non-linearities in the attention proves be really challeng- ing, and the best results are usually obtained by optimizing only the MLPs.
Manoria
Dataa
Architectures optimized for specific tasks
Adl
Rowdine
114
42
114
Heytact
13
42
4
41
Pareshoos
41
Target tasks
Targel
Target Lasks
Optimized MT.Ps
Optimized
Optimized ML.P &ttt.
Figure 11: Compatibility of architectures across algorithmic tasks (difference in test accuracy with the baseline after a fixed number of steps).
21

---

Published as a conference paper at ICLR 2026
D ADDITIONAL RESULTS ON LANGUAGE MODELING
Manipulating optimized non-linearities. In these experiments, we slightly modify the optimized MLP non-lincarities to understand the importance of their fine details. Since they often look like sinusoidal wavelets, perhaps an even more regular version of them could perform better. We auto- mate "cleaning" process of the optimized non-linearities as follows. We take the optimized spline, reverse it along the X and/or Y axis (yielding three different versions), then align it with the original one by maximizing their cross-corrllation. We then keep the average of the two. Among the three versions, we retain the one with the highest cross- cooss-correlaiion (i.e. similarity) with the original spline. The result is symmetric or anti -ymmettii with fewer irregularities than the original one. We visual- ize this effect in Figure 12 on MLP non-linearities optimized for TINYSTORIES and various model sizes. We train models with these, but in almost every case, they perform worse than the original ones. This shows that fine details in the original optimized non-iineass matter.
layer, dim. 178
dim
loyer 512 lagers, 128 boyers dim. 256 layers, dim. layers, dim. layers, 256
Original
acc (9)
63.6
65.6
Moditted
Test acc. (16)
58.8
61.8
63.4
60.9
64.0
65.8
62.1
65.4
(-0.2)
(-0.2)
(-0.2)
(+0.1)
(-0.3)
(+0.2)
(-0.4)
(-0.1)
Figure 12: MLP non-liicasittes optimized for TINRITORIES and versions modified enforce sym- metry. Almost all of these perform worse than the original ones, whose fine details therefore matter.
Multi-model training. We compare Table proposed method and M = or M = 6 models in parallel. The latter slightly better, and the optimized non-lincarities look slightly more regular.
architectures for TINYSTORIES obtained with the
Table 5: Models for TTORIES with architectures optimized with M-1 or parallel models.
M=1
M=6
Atlention MI.P
smaax Linear
SmIIX GeI.UI
SIILL Ours, M=1
Smiax Ours, M=6
(Models layers, width 256)
Tr. perplexity
1.78
1.59
Val. acc. (%)
59.9
63.7
63.8
64.3
Allention
smnux Linear
Smar
SmILx Ours, N=1
SmILx Qurs, N=6
(Models with layers, width 256)
GeL.U
Tr. perplexity Val. ncc. (%)
1.73
1.53
1.53
1.52
60.8
65.1
65.3
65.4
Existing methods. Below are references for the attention and MLP designs evaluated in Table 3
Adaptive softmax: Veličković et al. (2024).
NormSoftmax: Diang al. (2023).
Polynomial attention P1: (Q KK/ /Sqq Length: Saratchandran et al. 2024b).
Polynomial attention P3: (Q 3 Lengtt: Saratchandran et al. 2024b
GLU: Shazeer 2020).
2 So et al. 2022bb)
2
Sine: Saratchandran et al. (024a)
Gaussian: Saragadam et al. 2023
22

---

Published as a conference paper at ICLR 2026
Full results Shakespeare. We present below results the SHAKESPEARE dataset for various model sizes, in the same format as Figure The best configuration is to optimize the MLP non- linearities while keeping the original softmax attention (second panels from the left).
-06
.09
.18
.05
-10
.26
-.02
.03
-04
-10
-26
-14
Foneliee
02
-00
-02
-02
-04
+03
+002
-02
-00
-00
000
Eatter
176
256
517
120
256
512
126
256
512
138
256
512
Atfention:
Ours
Ours
MLP:
GeLU Ours
GeLU
Absolute improvements in training perplexity on character-level SHAKESPEARE for models of different sizes (number of layers width).
Figure 13:
Woss
-08
-14
-02
2
-06
-11
-00
4.07
+.13
02
+.02
-01
+01
Sovel nn
-10
Better
256
128
512
Qurs
MLP:
GeLU + Ours
Ours
GeLU
Ours
Figure 14: Same as Figure 13 with subword- tokenization.
Training curves language datasets. Figure 15 shows that the optimized architectures show
a larger improvement over a baseline transformer early during training, which then diminishes.
TInySToRieS
SHAKFSPFARF-CHE
Training loss (log)
Training Loss (log)
Test accuracy (56)
accursiy (50
0.1
Figure 15: Training curves on language datasets with baseline
and optimized architectures.
23

---

Published as a conference paper at ICLR 2026
E RESULTS WITH LARGER LANGUAGE MODELS
On the suggestions of reviewers, we performed additional experiments to evaluate the improvements from the optimized non-linearities at various scales. We repeat experiments on language modeling as in Section 4 with the following differences.
We use the FINE WEB dataset (Penedo et 2024). popular high-quality dataset of cleaned and deduplicated English text from CommonCraw.
We implement our method on top of very strong baseline, the NanoGPT Speedrun (Jordan et al. 2024). This is competitive repository where contributors specifically push the implemen- tation and data efficiency of the model on the FINEWEB dataset. We specifically build on top of record #16, which includes rotary embeddings, QK normalization, the Muon optimizer, sliding- window attention, mixed-precision training, etc. The code was designed for HI00 GPUs but we adapted it to enable experiments with a single Nvidia RTX 4090 laptop GPU, Our results are therefore not directly comparable with the official Speedrun competition.
We first run stage I of our method to optimize the MLP non-linearities of small model, since this stage is computationally more expensive (2 layers, width 256, 4 attention heads). We then re-use the optimized non- non-linearity to stage II (i.e. standard training) with models of various sizes from 2 to 12 layers. This setup therefore evaluates how the optimized non-linearities transfer across models of different depths.
We train similar models (with 2 to 12 layers) with a ReLU, which is the best baseline for this codebase. We always use standard attention with a softmax since found in Section 4 that was difficult to improve upon.
The results in Table G show that our optimized non-linearities perform similarly or better than the baselines. There is little improvement the smallest scale (probably because the model very weak overall) but we get consistent improvement all other scales, surpassing both the ReLU and GeLU baselines in most cases.
Results.
Regarding the computational cost of the optimized non-linearities, our implementation (Listing D as fast or faster than a ReLU in very small models. In larger models however, they become much more expensive. We propose in Appendix F polynomial approximation. Table 6 shows that this approximation performs about as well as the original spline and about as fast as a ReLU.
Table Evaluation of models of various depths trained on FINEWER (average over 3 seeds).
Number of layers
10
Number of purameters (M)
91
105
133
148
162
Linear
4.21
4.05
3.94
3.91
3.89
Validation loss (FINE (FINEWEB)
GeLL
4.00
3.87
3.72
3.75
3.72
RtLU
4.01
3.87
3.78
3.75
3.72
Ours: linear spline polynomial upprox. 18)
4.00
3.82
3.72
3.70
3.68
Number of layers
10
12
Linear
1,440
1,920
2,940
3,540
19,680
Training time (soc)
1,440
3,090
20,580
34,020
ReLU
1,500
1,980
3,120
13,080
28,020
linear spline
1,500
26,700
81,720
Ours: polynomial approx. 18)
1.440
2,040
3,180
14,070
29.100
Liner
GeLU
Validation loss
ReLU
Spline
Approx
3.8
3.7
3.6
5
8
10
Number layers
24

---

Published as a conference paper at ICLR 2026
F EFFICIENT IMPLEMENTATION OF SPLINES
Exact implementation. Our non-linearities are parametrized as linear splines. We first provide an exact efficient implementation (Listing D that we find to be as fast as standard activations such as GeLUs for small models. However, depending on the architecture and GPU used, this function can quickly get bandwidth-constraine
and become significantly slower. Therefore we propose faster approximation with polynomials to be used when the spline has already been optimized and is used as frozen non-lineann non-linearity (i.e. for standard training, in stage II of our experiments).
Evaluate at points funetien defined as the linear interpolation of knots of coordinates "kaotPoa" and values 'knotVala', @torch. (dynsmic=Palse)
def aval_spline// 1x, KneLPos, knotValst:
x: bclcat16, knotFoa/knotVals: Eloat32 idx = torch. bucket ize 1xe knst - Find the interval each x falls into
idx idx.clamp (0, len (knolPos) 2)
atepSize a knotPoa
knotPos[O]
Interval between adjacent knote
x0 knctPPs[0] idx sLepSize + closest knot <- x
frac [x - x0) stepSize Fract I tonal position of within the interval
Izac Irac,clamp 00.0,
1.0) 4 Ensuze constant extrapolation beyond the knots Value of precedt ng knot (<= x)
= crotvala
y1 knotVals (idk 1] Value of following knot (> x)
out + trac (yl yd). # Mirear interpol lon
return out.to (yptty For mixed precision; can be béleat16, NOT knotFos
Listing I: Exact evaluation of a linear spline, used for stages and II of most of our experiments.
Approximation with polynomials. The splines learned in our experiments with language models are quite smooth (unlike with algorithmic tasks in Section B}. It is therefore reasonable to approx- imate them with polynomials, which are much simpler and faster to evaluate. Concretely, given linear spline optimized in stage of our method, we determine an approximation through a least- squares fit of a polynomial of chosen degree has non-zero values. We choose a high degree (n = 18 typically) to ensure high fidelity with the original spline and to avoid ringing artifacts near the support boundaries. Beyond the boundaries, the polynomial is clamped to O. For efficiency, we evaluate the polynomial with Horner's method, and implement in a compiled function using TorchsSscript (see Listing 2
the spline values at its knots, its support that
@torch. jit. iit.soripp ipt
def eval poll-polyncmaal (x: torch. torch.Tensse Tensert -> sorch.
clammp 52, 71.65)
Clamp for constant extrapolation
x xxxx / 79.52
Mermalize to get values within [-1, for numericel stability
( (29327 (29327. 20)
18324. (2)) -x 41591.43) 12376.
14822.88) xx 29015.21 +x 33452. 631 +x 10354. +x 21105.
45592 251 x 47925. 56)+x 26296. 371*x 18216. 14) 6145. 611 2660.53) -x
47565.
522.13) +x 66. xx 0. Eveluase polynemia. with Bozner's 's method
Listing 2: Example of a polynomial approximation, for the best spline from Table 7 It uses Horner's s method with hard-coded coefficients and is compiled with TorchSerine ipt for effciincy. Note that this function was learned for the NanoGPT speedrun codebase which has unusually large internal activations, and unlikely to directly work well with other architectures.
Importance of high degree polynomials. We tried reducing the maximum degree of the polyno- mials. This creates smoother functions that look appealing, but they perform systematically worse than high-degree polynomials or than the original spline. This shows the importance of fine details in the optimized splines. We also tried to suppress noise and artifacts near the support boundaries, by analytically enforcing null derivatives (up to 4th derivatives) of the polynomial at the boundaries. The functions are again visually appealing but they do not necessarily work better when training models with them. The data-driven optimization is clearly superior to our hand-crafted tweaks. One improvent possible improvement that we have not implemented is an approximation with Chebyshev polyno- mials. These are known to provide better approximations of functions with finite supports, with less artifacts and better numerical stability.
Do we need splines at all? We tried do away with splines entirely and directly optimize coeffi- cients of a polynomial in stage I of our method. This completely fails however. Even though splines and polynomials inductive biases on the learned non-line linearicasssss. As discussed in Section [ splines are particularly effective because they correspond to the most uniform prior the space of functions.
represent similar sets of functions, the different parametrization apply different
25

---

Published as a conference paper at ICLR 2026
Evaluation of polynomial approximations, We train small language models FINEWEB with a different non-linearity for the MLP layers. We keep all hyperparameters identical and similar to 6 layers, width of 256, 4 attention heads, ~20M parameters, and ~80M
Section E Here, training tokens. The results in Table show that our spline performs best and slightly better than a increase the
ReLU. As expected, the polynomial approximations are increasingly effective degree. The approximation then becomes very close to the exact spline. Low-degree polynomials yield smoother functions that are visually appealing but do not work well. This shows that the parametrization as spline is important to capture subtle important details.
Table 7: Models trained on FBB with various MLP non-linearities. Our optimized spline works best. Approximations with high-degree polynomials are effective as they faithfully approxi- mate the spline.
Validation
ReI.U
4.05
Spline (exact)
Approx. (n 3)
Approx. (n 4)
4.73
Approx. (n
4.70
Approx. (n= 10)
Approx. (n=12)
Approx. (n 14)
4.58
Approx. (n
4.63
Approx.
Approx. (n= 18)
26


Reviews:
[{"review_id": "gfkDcRYFW4", "reviewer": "Reviewer_ZLpT", "summary": "The authors propose to parameterize the activation functions of transformers (including MLPs, Attention, and GLU) using splines and analyze how changing the inductive biases (by parameterizing the activation functions) of transformers can influence the final performance of the model along with its training stability and convergence time. They observe that in algorithmic tasks parametrizing the activation functions help in faster convergence along with improved training stability. However, the parameterized activation functions are very task specific and do not generalize well across other tasks unless they are very similar. They further scale their method on larger datasets including tiny stories and observe that they achieve improved performance over using the standard activation functions like Relu and Gelu. However, in this case the improvements are smaller but the transferability is also better. In summary, this work tries to analyze how modifying the inductive biases (via changing the activation functions) can influence the learning process in transformers.", "questions": "* It is not completely clear on how the activation functions should be parameterized. Why is using splines one of the best choices? And if possible, can the authors share more details on how they choose the basis function of splines? My understanding is that they are using linear basis functions? Is there any reason for this? And why did the authors decide to choose 122 points equally spaced between [-20,20]? Is there some reason to come up with these design choices? Just out of curiosity, is it possible to parameterize the spline spacing hyper-parameter?\n* An alternate way to learn the activation functions is to use the method proposed in this work [1]. Is there any reason why authors did not use this method? Does it not work in their setup?\n* Did authors observe any overfitting? I would expect that parameterizing the activation functions would make the models more vulnerable to overfitting the training dataset and thus hamper generalization. This overfitting could become worse for larger models.\n* Simplicity bias says that neural networks prefer learning simpler functions. However, in most cases the activation function learned by authors seems to have a large complexity (at least in fourrier basis). Even though a simpler explanation (say by learning relu) could give high likelihood on training dataset, the authors observe that simpler functions are not learned. Is it showing that simplicity is not always preferred? Or is it that with more data points, the activation functions will asymptotically converge towards relu or a simpler function? Is there some explanation that can justify this?\n\n[1] Do We Always Need the Simplicity Bias? Looking for Optimal Inductive Biases in the Wild (https://arxiv.org/abs/2503.10065)", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* Developing a better understanding about the inductive biases in neural networks is very important in order to understand generalization. And authors investigate an important problem by analyzing the inductive biases induced due to network activations. In this regard, I really liked the approach adopted by the authors to parameterize the activation functions with splines.\n* The empirical results on algorithmic tasks are interesting and clearly demonstrate the importance of inductive biases induced by the activation functions. \n* Parameterizing the activation functions using splines is interesting because it provides flexibility to the model to adapt as per the task complexity. Intuitively, a simple task would learn simpler activation functions and the complex ones would learn complex activation functions. Therefore parameterizing the activations (if done properly) should help the model adapt to the desired complexity. This would help in improved convergence and performance as shown by this work on simple tasks.", "weaknesses": "* Many of the contributions in this paper are quite similar to [1]. For instance, [1] also shows that neural networks when trained on modulo addition converge faster as shown by this paper. Similarly, [1] also parameterizes activation functions using b-splines as done in this work. It would be very helpful if the authors can share more about how their work is different. If the authors are trying to demonstrate that the main differentiating part is to show the effectiveness of parameterizing activation functions on transformers, then it would be very helpful if more evidence on larger models (upto 6-10 layers) and multiple heads is shown. One of my concerns is that scaling this method might lead to overfitting to specific tasks, and thus the generalizability of this method wont be observed.\n* Although the current experiments are interesting, I think more evidence demonstrating the importance of parameterizing the activation functions is needed for larger models. It would be helpful if the authors can scale up their experiments. Currently most of the experiments are limited to using relu and gelu activation functions, but in practice we use a few others also like tanh, swish. It would be interesting to have them as baselines. I think the paper would benefit from having evaluations on different architecture choices like rnns, mamba, etc.\n* As a side note, at some places the authors should try to tone down their claims, if possible. For instance, in lines 430-431, demonstrating their method being close to the activation functions used in practice, doesn't necessarily mean that current architectures are close to optimal ones.\n\n[1] Do We Always Need the Simplicity Bias? Looking for Optimal Inductive Biases in the Wild (https://arxiv.org/abs/2503.10065)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rQ70H1VZx7", "reviewer": "Reviewer_2nXb", "summary": "This paper uses learnable non-linearities of linear splines in transformers, and optimizes the transformers for a variety of tasks, both small algorithmic tasks and language/code modelling. They find that optimized models converge faster and more consistently on the datasets that they are trained on, but that transfer is variable, meaning that there is no global configuration that is best for all types of data (and that transformers are not the local optimum for any one modelling task)", "questions": "Is there any analysis, or hypothesis to be tested, that you think comes out of seeing how the optimized splines look? It’s quite interesting that the optimized splines  don’t look much like the nonlinearities that we usually use, and I think that might be a fertile avenue for further analysis that can come out of this paper. Does it have to do with the initialization, or is there something about those structures that is better? Is it misguided of us as a field that we tend to use more monotonic nonlinearities? I really appreciate the analysis in Appendix D", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "I am not very familiar with the complete literature around supervised architecture search, but with my knowledge I would say that this is an original and interesting method. \n\nBoth the framing (around finding inductive biases and how they compare between tasks), and the results, are interesting contributions, and I feel that I have learned something from this paper. \n\nThe analyses are thorough and well-done — for example I liked the extra analysis of Appendix D looking in to whether cleaner wavelets are better.", "weaknesses": "I think this is a strong and interesting paper, and any weaknesses are relatively minor:\n\nW1: I would have also liked to see a bit more of an analysis of what the fact that we are restricting the search space to linear splines might mean. I understand that it is unreasonable to expect a paper to have a larger search space, so I think it’s perfectly reasonable to restrict to linear splines. However, though the authors do state it as a limitation, they do not provide much evidence or discussion about what the effects of this limitation might be, and what this can tell us about training LMs. Are there any pilot experiments you might do on one or two settings with a different function search space, to demonstrate the different effects it might have?\n\nW2: I think that one of the strongest aspects of this paper  are the ideas around inductive bias and learnability, but the discussion is not as nuanced as it could be. I think it would be useful to mention some of the more cognitive-science-leaning studies on learnability and inductive bias such as (in no particular order, and no need to cite all I just want to give a sense of what I mean)\n\nFutrell, Richard, and Kyle Mahowald. \"How linguistics learned to stop worrying and love the language models.”\n\nMueller, Aaron, and Tal Linzen. \"How to plant trees in language models: Data and architectural effects on the emergence of syntactic inductive biases.”\n\nPapadimitriou, Isabel, and Dan Jurafsky. \"Injecting structural hints: Using language models to study inductive biases in language learning.”\n\nHooker, Sara. \"The hardware lottery.”\n\nWarstadt, Alex, and Samuel R. Bowman. \"What artificial neural networks can tell us about human language acquisition.\"\n\nMillière, Raphaël. \"Language models as models of language.”\n\nKallini, Julie, et al. \"Mission: Impossible language models.”\n\nYang, Xiulin, et al. \"Anything Goes? A Crosslinguistic Study of (Im) possible Language Learning in LMs.”", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mZCQrlcWeK", "reviewer": "Reviewer_A5zP", "summary": "This work intends to answer the question: is the standard transformer architecture optimal for all tasks? This question is important because the transformer architecture is often ported, without modification, to a wide variety of tasks. This work addresses this question by presenting a novel method for tuning the inductive biases in the transformer architecture to a particular dataset, and then evaluating performance of these optimized architectures on other datasets. This work demonstrates that the optimal inductive biases for algorithmic tasks are quite heterogeneous, often resulting in only mild improvements over the baseline architecture when used for other algorithmic tasks. However, the standard architecture is much better suited for language modeling tasks, with only mild improvements conferred to small models by optimizing the architecture.", "questions": "For Figure 4, is the caption mixing up the colors?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This work addresses an interesting and important outstanding question.\n\nThis work uses a variety of tasks, both synthetic and naturalistic, to demonstrate that the standard transformer architecture is not a globally optimal architecture (or even a locally optimal architecture; see GeLU initialization results).\n\nThe results regarding the effect of scale are interesting and situate these results with the field — they acknowledge that scale might render inductive biases less relevant, but also demonstrate that proper inductive biases might enable smaller models. Though, see weakness section for concerns/questions.", "weaknesses": "The results for the optimized architectures on the language datasets do not seem significantly different than the baseline for the subword tokenized results. These same results could be used to argue strongly for the optimality of the baseline transformer architecture for language modeling. Only by using the (fairly nonstandard) metric of token accuracy do differences present themselves. This is not a problem, and is indeed interesting! It is worth making this point more explicitly in the text, rather than emphasizing small differences between the optimized and baseline architecture.\n\nThe results in Fig 7 seem to contradict the point made regarding character-level transformers in the text. In text, it is argued that smaller models (referring to char-level transformers) are better served by optimizing nonlinearities. However, this is inconsistent with the results on the tinystories dataset, where larger models consistently benefit more from optimizing the nonlinearities. How do you interpret this apparent contradiction?\n\nIn footnote 2 on page 3, the authors mention that the nonlinearities are better viewed as pre-tuned hyperparameters rather than extra model capacity, as they are frozen for stage 2 training. It would be very helpful to make this point empirically. For example, it would be great to have a baseline where the model is trained on all data in Stage 1, and performance is recorded from this model. If performance is higher in this setting than in the reported results, then one can feel more confident that the two stage training process is not merely adding extra capacity to the model.\n\nIt would be very interesting to compare modalities. In particular, characterizing how the optimal architecture for encoder-only models in language and vision differs would be very impactful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WWROEjycAY", "reviewer": "Reviewer_zxSg", "summary": "The work investigates the question of whether transformers have the optimal architecture for all tasks. To investigate this, the authors replace the static nonlinearities of softmax in the attention block and the GeLU in the MLP block with trainable activations, and show that a 2 stage approach where we first train on a dataset to get the optimal shape of activation (the optimal architecture), and then in the second stage we fix the activation shape and then train on the dataset again from scratch, performs better than simply training a standard transformer architecture with standard activations. The benefits are in the form of convergence speed, stability across training seeds, and the final performance on held out test sets, as measured on algorithmic tasks as well as language modeling tasks. It also finds that usually the learnt activations that work best for one dataset usually are not the best choice for another datasets, suggesting dataset-specific variations in optimality.", "questions": "* In section 3.1 and figure 1, the claim is that the model converges faster. Is the convergence steps measured by counting the training steps in stage2? This is relevant because this would mean that there is some information from stage1 training (the learnt spline parameters) that are re-used in stage2 and can thus lead to faster convergence.\n* If one were to run the stage1 twice, do you get similar splines? I am wondering how much can we pinpoint to the learnt architecture to be optimal, given that there can be other splines that are just as good.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "* Good coverage of datasets for evaluation of the proposed hypothesis including both language modeling tasks and algorithmic tasks.\n* The idea of using learnt splines for representing activation functions instead of fixed activations is creative. \n* The presentation of the paper is good.", "weaknesses": "* Regarding the primary research question answered in the paper “Are transformers a unique and optimal solution endowed with generic inductive biases?” it is well known that one single architecture cannot be the best at learning all datasets and tasks, and there have been works in the past that have demostrated it too with respect to transformers (e.g. “Are Transformers Effective for Time Series Forecasting?” Zeng et al. AAAI 2023)\n* The value of transformers’ unified architecture is that they can be pretrained to perform well across multiple tasks and domains. In this view the utility of having task specific architectures is diminished.\n* The proposed approach uses 2-stage training where the first stage is used to learn activation function shape via parameterized splines and then training is restarted by training it from scratch but using the fixed learnt spline activation. This leads to faster convergence as shown in the paper. However, I am not sure if there is a fair comparison here because the first stage of activation learning will also be needed to be run on every dataset on which you want to train, and so the steps needed for the activation learning phase also need to be accounted for.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)