Composer: A Search Framework for Hybrid Neural Architecture Design
Authors: Bilge Acun
Prasoon Sinha
Newsha Ardalani
Sangmin Bae
Alicia Golden
Meghana Madhyastha
Chien-Yu Lin
Fei Sun
Neeraja J. Yadwadkar
Carole-Jean Wu
Venue: ICLR 2026 Poster (2026)
Keywords: Neural architecture search
hybrid models
efficient ML
Primary area: foundation or frontier models, including LLMs
TL;DR: An efficient search framework for hybrid neural architecture design

Abstract: Hybrid model architectures that combine computational primitives (e.g., Attention, MLP) in different ratios have shown promising performance beyond Transformers. Some studies have shown that different interleavings of primitives can affect model quality as well. However, prior works explore the hybrid model architecture design space manually. Due to the large design space and training costs, discovering hybrid models that combine key computational primitives for pre-training is challenging. In this work, we take a principled approach in designing a modular hybrid model architecture search framework — Composer. Composer explores model architectures at a small scale and extrapolates the top-performing model architectures to a larger scale using our proposed scaling strategies. Using Composer, we discover new hybrid LLM architectures that outperform Llama 3.2. Compared to Llama 3.2 and previous state-of-the-art baselines, the new model architectures consistently reduce validation loss at parameter scales of 350M-8B and improve evaluation accuracy on the downstream tasks by up to 2-2.1% on average while improving both training and inference efficiency.

Published as a conference paper at ICLR 2026
COMPOSER: A SEARCH FRAMEWORK FOR HYBRID NEURAL ARCHITECTURE DESIGN
Bilge Acun", Prasoon Sinha' Newsha Ardalani, Sangmin Bae, Alicia Golden, Chien-Yu Lin, Meghana Madhyastha, Fei Sun, Neeraja J. Yadwadkar', Carole-Jean Wu FAIR at Meta, The University of Texas at Austin Joint first authors, work done at Meta acun@meted com, prasoon. sinna Sethaa@tteeeeseee..em
ABSTRACT
Hybrid model architectures that combine computational primitives (c.g., Atten- tion, MLP) in different ratios have shown promising performance beyond Trans- formers. Some studies have shown that different interleavings of primitives can affect model quality as well. However, prior works explore the hybrid model ar- chitecture design space manually. Due to the large design space and training costs, discovering hybrid models that combine key computational primitives for pre- training is challenging. In this work, we take principled approach in designing a modular hybrid model architecture search framework - Composer. Composer explores model architectures at small scale and extrapolates the top-performing model architectures to a larger scale using our proposed scaling strategies. Using Composer, discover new hybrid LLM architectures that outperform Llama 3.2. Compared to Llama 3.2 and previous state-of-the-arr baselines, the new model ar- chitectures consistently reduce validation loss at parameter scales of 350M-8B and improve evaluation accuracy on the downstream tasks by 2-2.1% on average while improving both training and inference efficiency.
INTRODUCTION
Transformers (Vaswani et al., 2017) have long served as the foundation of large language mod- els (LLMs), powering mainstream models like BERT (Devlin al., 2019) and GPT (Radford et al., 2019). The standard Transformer architecture features a fixed sequential interleaving of self-attention and multi-layer perceptron (MLP) layers. While this design remains effective, recent works demonstrate that hybrid LLM architectures-wv which deviate from the conventional Trans- former stack can further improve model quality. For example, unlike Transformer- based archi- tectures which stack 1:1 ratio of computational primitives, approaches like Qwen3-Next (Qwen, 2025), Mamba-2 (Dao & Gu, 2024), and MAD (Poli et al., 2024) adjust the ratio of Transformer and State Space Model (SSM) primitives within stackable blocks, skewing the composition toward one type or another. Mcanwhile, other approaches rearrange the primitives in more sophisticated patterns by breaking the stacked structure. For example, DeepSeek's V3 MoE model (DeepSeek- AI, 2025) incorporates a few dense MIPs in the initial layers followed by sparsely activated MoEs. FistVIT (Vasu et al., 2023) uses convolution in the beginning stages of the model and attention in the later stages. Sandwich Transformer (Press et al., 2020) reorders the interleaving of attention and MLP layers without changing the ratio.
Despite promising advances in hybrid LLM architectures, the model architecture design process manual and based on intuition no systematic framework exists today to enable automatic, efficient discovery of hybrid LLM architectures that perform well at scale. An effective search framework is essential given the vast model architecture design space-for example, a 32-layer hybrid LLM consisting of simply attention and MLP layers already yields over 4 billion (232) possible architecture configurations. The Nemotron model family builds a Post Neural Architecture Search (PostNAS) framework that prunes and replaces blocks of pre-trained models (Bercovich et al., 2025; Gu et al., 2025). STAR (Thomas et al., 2025) presents an initial attempt towards a framework targeting pre- pre-training hybrid LLMs from scratch; however, its design assumes conducting search on the target dataset for edge use cases. We find that conducting search with web-scale datasets is either ineffective impractical for performance evaluation.

---

Published as a conference paper at ICLR 2026
Composer
80
Computationel Primitives
Model condicistn
Search Engine
Evalustor
Large-scale
Search
60
Acturaay
LLM
(Modet:Aoc)
Aggregetor
Small sosie
Extrapola for
Moving Average
Transformer
model
20
40
60
80
100
Search Trial Index
Figure 1: Design of the Composer search framework (left) and the accuracy convergence over search trials (right).
In this work, we take principled approach to answer the key research question: how do we design model architecture search framework that automatically and efficiently discovers novel hybrid LLM architectures for pretraining that outperform state-of-the- th-t--rrt- models at scale? We find that when scaling down dataset and model size via Chinchilla scaling laws (Hoffmann et al., 2022), the model quality of small explored hybrid LLMs does not reflect large scale performance. Therefore, require new methodologies for small-scale neural architecture search that reflective of at-scale performance. We design Composer, a hybrid neural architecture search (HNAS) framework. Unlike traditional neural architecture search that assumes fixed interleavings of computational primitives when searching over model hyperparameters (e.g., model width, number of layers) (Wang et al., 2025; Javaheripi et al., 2022b), Composer rearranges the interleaving pattern and ratio of computa- tional primitives to automatically discover perforerfonnin hybrid LLM architectures. Leveraging our modular framework, we conduct an extensive exploration to answer the following key design questions:
1. What is an effective search algorithm for conducting small-scale model architecture search? Searching the vast design space requires efficient search algorithms and effective techniques for scal- ing down model dimensions. Composer's s Search Engine generates and searches through candidate hybrid LLM architectures bounded size. We propose both Bayesian Optimization and iterative search techniques for efficient search, and provide key architectural design principles that outper- form standard Transformer architectures.
2. What datasets should
to evaluate scaled-down hybrid LLM architectures? For ef- ficient search, candidate architectures need to be trained and evaluated efficiently with either small scale proxy datasets or sampled-down versions of target datasets. We investigate the efficacy of three different datasets for small-scale search with Composer's ' Evaluator.
3. How do we synthesize the model candidates from search results into final hybrid LLM? Small-scale search can produce multiple high performing model candidates. Composer's Aggregator synthesizes the search results into final hybrid LLM using clustering technique with the goal of generalizing to large-scale architectures.
4. How do we extrapolate discovered small hybrid LLM architectures to large-scale model sizes? Once the small model candidates are found, the architecture needs to be scaled. Com- poser's Extrapolator scales up the discovered architecture to an arbitrary desired size (we scale up ~ 1000x and show results up to 8B) with two proposed techniques: stacking and stretching.
As our focus in this work LLMs with Attention/ML.P hybridization. After conducting our framework's design exploration, we demonstrate that the top-two performing hybrid LLMs discovered with Composer, termed Com- posite architectures, outperform Llama 3.2 and several other ttate-oftthe-art model architectures. Our Composite architectures, consisting of advanced interleavings and a 1:2 ratio of Grouped Query Attention (GQA) SwiGLU layers, reduce validation loss by 0.03 and increase accuracy down- stream LLM evaluation tasks by 2-2.1% on average. across a variety of model sizes and training budgets. Meanwhile, we increase training throughput by 1.25x, reduce KV cache size by 1.69x, and reduce inference latency by 1.33x x on average.
build a HNAS framework in principled manner, we design
COMPOSER DESIGN
We introduce Composer, an efficient and automatic hybrid neural architecture search framework, as illustrated in Figure 1. The framework ingests set of computational primitives and uses Bayesian
2

---

Published as a conference paper at ICLR 2026
Optimization to perform hybrid model architecture search small scale (e.g., million parameter models) and discover architectures that will perform well when extrapolated to ~0000 scale (e.g., 3B model size). Composer has four components: the HNAS Engine, Evaluator, Aggregator, and Extrapolator. After, Composer outputs a new hybrid LLM at a given size (e.g. 3B) ready for pre-training. We describe each of Composer's core components next.
2.1 HYBRID NEURAL ARCHITECTURE SEARCH ENGINE
The HNAS Engine systematically discovers high-quality, novel hybrid LLM architectures. We de- fine a hybrid LLM as sequence of computational primitives (c.g.. Attention, MLP) from the set P = {P1. P2,.P2,... ................................ :PZ), where Z is the number of unique computational primitives, For a fixed num- ber of layers N, a hybrid LLM is formally defined as a = (A1, 42 ... aN) ) E PN, where a E P specifies the primitive at layer The discrete search space contains all possible primitive arrange- ments AN
uNy) a; E P, Vi E [N]], yielding |Avl z candidate architectures.
The discrete search space exponentially grows with the target model size, since the depth of the target size N grows with model size (e.g.. Llama 3.2's depth increases from 32 to 72 layers from 1B to 8B model size). To efficiently navigate the large design space, we propose three search methodologics.
(1) One-Shot Search: This methodology performs one-shot n-layer search, where S N. n < N. the discovered model is then extrapolated, or scaled up, to target size via one of proposed extrapolation techniques in 2.4. We leverage Bayesian Optimization with Gaus- sian Process surrogate modeling to navigate through the hybrid architecture design space of size 7 We build upon the Ax framework (Olson al., 2025; Meta Platforms, 2024), using BoTorch SingleTaskGP model with an RBF Kernel, dimension-scaled priors, and qLogNEI ac- quisition function for single objective optimization. Bayesian Optimization offers better sample efficiency and uncertainty modeling compared to Reinforcement Learning or evolutionary search. The optimizer targets a black-box function An + I measuring validation accuracy after pre-training: f(a) Accuracy(Pred PreeTrain(a, Durin), Deat), where D and Dyal are the training and validation dataset used by the HNAS Evaluator ($ 2.2). We solve the optimization problem: a" arg max (a))
(2) End-Layer Incremental Search: This methodology prunes the design space and incrementally builds hybrid LLM to target size. Specifically, we perform n-layer iterative search, starting with layers, where n N and n > 1. Then, we progressively increase the architecture depth by n layers at each step, fixing computational primitives for the previous layers and searching only the last n. Hence, at any given iteration, there are only 2n unique architectures to evaluate. This process repeats up to N layers with total of N/n iterations.
(3) Middle-Layer Incremental Search: This methodology follows a similar process to (2), but searches over middle rather than end layers. The Search Engine first conducts n-layer search. Then, for each subsequent n-layer expansion (e.g., n to 2n, 2n to 3n), we fix the beginning and end layers by splitting the previously found architecture down the middle. We then only search over the middle n layers. This approach ensures that at each stage, the central layers are optimized while the outer layers remain fixed from the prior iteration. This process repeats up to N layers.
Further improving search efficiency via width scaling. To reduce search cost, the one-shot search method searches over fewer layers than the target model size (i.e. <<< N) while the incremental search techniques prune the design space with their iterative nature. However, without reducing the width of computational primitives, the search cost remains prohibitive ($ 3.3). Hence, we also scale down the widths of the primitives compared to the target model size. We find that this not only reduces search cost, but also enables Composer discover higher quality hybrid LLMs ($ 3.3).
2.2 HYBRID NEURAL ARCHITECTURE EVALUATOR
During search, the HNAS Evaluator trains and evaluates candidate hybrid LLMs with small dataset to provide fast, reliable signals the potential quality of the architecture at scale. The dataset used during search is crucial for accurately identifying architectures that perform well at scale when pre- trained with the target dataset; use DCLM (Li et al., 2024), a large web-scale text dataset, for pre- training. We empirically evaluate three datasets for small-scale search in $ 3.2: randomly sampled- down DCLM (Li al., 2024), MAD (Poli al., 2024) which is synthetic token-manipulation
3

---

Published as a conference paper at ICLR 2026
Composer Component
Default Methodology
Details
HNAS Search Engine
One-Shot Search
Both and 16-layer search
HNAS Evaluator
MAD's synthetic tasks
Details in Appendix E.2
HNAS Aggregator
No clustering
K-means with 5 clusters
HNAS Extrapolator
Stacking and stretching
Stack 6-layer search, stretch 16-layer search
Table 1:
Default methodology used during ablation study of each of Composer's 's components.
dataset designed to probe different capabilities of LLMs, and finally BabiStories, which is a synthet- ically generated children's 's story dataset (Zhang et al., 2025) (an OSS model generated version of the TinyStories dataset (Li & Eldan, 2024)). Appendix E.2 provides further details about each dataset.
2.3 HYBRID NEURAL ARCHITECTURE AGGREGATOR
The Aggregator post-processe search results, either after the search process completes with One- Shot Search or after each iteration of End-LLayer/Mddde-- small-scale hybrid LLM. We propose N, clustering to select the primitive at each layer conditioned
Incremental Search, to finalize the
on the sequence of the c previously selected primitives. Formally, let C =aa() a
a (
denote the set of top candidate architectures, where each architecture (*)) (a 22 an is
a sequence of computational blocks indexed by layer. We use K-means clustering, with validation accuracy of the candidate models from search as the target metric, to populate C. For given layer i, let - denote the selected sequence of blocks at the previous e layers. Then,
Ne: à = mode mmde
{{a (m))
(m) a (m) C, ai-ei-1 a Vi E [L,n]
(I)
where à; denotes the selected block at layer i, and mode(.) returns the most frequent computational primitive among the filtered candidate architectures that match the selected prefix.
When c = 0, No clustering selects the dominant block at each layer among the top candidate archi- tectures independently without conditioning on prior layers. When c Ni clustering conditions the block choice at each layer on the block selected at the immediate preceding layer. Finally, for the entire sequence of pre-
1, the block selected at each layer index is conditioned viously selected blocks, enforcing full prefix consistency. We also consider simply using the best architecture discovered during the search process. We find that No clustering (c 0) produces the best hybrid LLM at scale; clustering over all oopperformmng candidate LLMs smoothes out noise or overfitting that may occur during amall-scale search. We include this analysis in Appendix D.5.
2.4 HYBRID NEURAL ARCHITECTURE EXTRAPOLATOR
With final small hybrid LLM, the HNAS Extrapolator scales up the architecture the desired model size (e.g., 3B). the depth of the hybrid LLM n matches the depth of target model size N, the Extrapolator simply scales the width back up to that of Llama 3.2. Otherwise, the Extrapolator also scales up the depth using one of two techniques, stretching or stacking. We describe each next.
Extrapolation via stretching. Stretching scales up the depth of the hybrid LLM while keeping the interleaving pattern and ratio of computational primitives the same as the discovered, small hybrid LLM, thereby "stretching" the architecture. Formally, we partition hybrid architecture a into G contiguous groups, where each group g contains identical primitives p:
G Σ
h h [[(1, [(11, (p2, 12).................. ha)]: where Pg E, hig E Et
h1
Given the desired and current model sizes M and m. respectively we define scaling factor M The scaled architecture is then
h large
[p1, [s- h1 1), (p2, [s h2]), ........... (pc. [s ha
Intuitively, the scaling factor represents the ratio by which the model size is increased, scaling the number of layers in each group proportionally to reach the target model size. Appendix ?? provides visualization of stretching.
4

---

Published as a conference paper at ICLR 2026
Extrapolation via stacking. This technique considers the discovered small hybrid LLM stack- able block. To scale up to the desired parameter size, we stack s copies of this block sequentially, where is the scale factor We use this term to shrink the discovered L.M following Equation 3 and append these layers to the end of the stacked architecture. This ensures the final architecture elosely matches the target size.
mm J- We also define remainder scaling term M mod (m x s).
Figure depicts how each extrapolation methodology works.
ALP
Smal-scole hybriil model architacture (e-gu 1A paranaters)
MLP
Hybrial wlack الaهي sc search
Attantion
Attestion
"Stratck out"
Scaled hyporld wodel arckitecture via stretching (eig 3B)
Scaled bybrid model architacture stocking (e.g, 3B)
Figure 2: Illustration of extrapolating small-scale hybrid architecture to large-scale via stretching (left) or stacking (right).
DESIGN EXPLORATION OF COMPOSER'S S CORE COMPONENTS
We perform a design exploration study with different methodologies for each of Composer's 's com- ponents proposed in $ 2. For each component under test. fix the methodology of the other components (Table 1) based on our findings which methodology performs best. Each methodol- ogy produces different small hybrid LLMs. To evaluate each methodology's efficacy, we scale up the hybrid LLMs to IB parameters, pre-train them with the DCLM dataset (Li et al., 2024), and compare validation loss. We use the same width for all LLMs, ensuring performance differences arise due to differences in the ratio and interleaving of computational primitives between hybrid LLMs. We set width dimensions to that of Llama 3.2 1B (2048 8192. 32 attention heads, 8 KV heads). Appendix B details the architectures produced by each methodology. We provide further details of our pre-training setup in Appendix C.3.
3.1 EXPI.ORATION OF SEARCH METHODOL.OGIES
We first evaluate the model quality versus search cost of the three search methodologies we propose in 2. For One-Shot Search, we evaluate two variants with n 6 and n 16, performing 6-layer and 16-layer search, respectively. For End-Layer and Middle-Layer Increment Search, we set 2 and n 4 to perform 2-layer and 4-layer incremental searches, respectively, up N=32 (Llama 3.2 1B, the target model size, has 16 Transformer blocks or 32 layers). Figure 3-left presents the at-scale DCLM validation loss and search cost for each methodology.
Observation I:
All three proposed search methodologies produce hybrid LLMs that outperform Llama 3.2 across compute budgets, showing that breaking the standard Transformer architecture with more advanced interleavings of attention and MLP layers improves model quality.
The End-Layer Iterative search methodology produces hybrid LLM with (roughly) a 1:1 Attention- to-MI.P ratio, like Llama 3.2. However, during search, Composer learns that breaking the standard sequential Transformer interleaving for more intelligent arrangements with variable sized groups of Attention and MLPs interleaved can greatly improve model quality.
Observation In addition to the ordering of the layer types, we find that 1:2 Attention-toMMLP ratio can further improve model quality compared to 1:1 ratio.
Both Middie-Inccemmnntal and One-Shot Search outperform End-Inncremenal Search (0.01-0.02 re- duction in validation loss across training budgets). Only 40% of consists of attention layers. Both One-Shot searches further reduce the number of attention layers to only 33% of the hybrid LLM's depth while maintaining model quality with lower search costs compared to End-Innrrementl (1.4-2.1 reduction in search cost). Hence, despite all three proposed search methodologies producing high quality LLMSs compared to Llama, we leverage One-Shot
's hybrid LLM
5

---

Published as a conference paper at ICLR 2026
Liama 3.2
One-Shot Layer
Liama 3.2
MAD Stacking
End-Layer Incremental
One-Shot Layer
Small-Scale DCLM
MAD Stretching
3.4
Middle-Layer Incremental
Search Cost (GPU Hrs)
3.4
Large-Scale DCLM
Search Cost (GPU His)
LLLM Val Loss
3.2
75
3.2
600
3.0
80
3.0
400
2.8
2.8
200
2.6
2
Training Budget (FLOPs x 1e20)
Training Budget (FLOPs 1e20)
Figure 3: Design exploration of (1) search methodology for Composer's HNAS Search Engine (left side) and (2) datasets for Composer's HNAS Evaluator (right side). We report the model quality at 1B scale and search cost for each technique.
search throughout the rest of this paper as the methodology for Composer's Search Engine due to better model quality and search cost trade-off.
3.2 EXPI.ORATION OF EVALUATION METHODOLOOGY
is crucial that the dataset used for small-scale search enables Composer to identify representative hybrid LLMs that perform scale. However, searching target model sizes (e.g., 3B) with web- scale datasets is prohibitive. As DCLM is our target large-scale pre-training dataset (Li et al., 2024), evaluate two approaches for leveraging it during search: (I) scaling down both the model and the dataset via scaling laws, (2) scaling down only the dataset while keeping the model large.
Observation 3: Compared to using the traditional scaling laws methodology scale down web- scale datasets and model sizes, we find small-scale proxy datasets can be more effective in guiding the search process.
We reduce the model and data size, following Chinchilla scaling laws (Hoffmann et al., 2022), down to 4-layer 4M parameter model trained on 98.6M tokens (12K samples). After the 4-yyyer search, we extrapolate to IB size by stacking. Figure 3-right shows that this methodology, Small- Scale DCLM, only slightly outperforms Llama 3.2 and improvements diminish with larger training budgets (4e20 FLOPs). Therefore, as an alternative approach, we sample down only the dataset size 12K samples while keeping model size large. We perform 16-layer search and extrapolate IB size via stretching. Figure 3-right shows that this methodology, labeled Large-Scale DCLM, yields high quality models that consistently outperform Llama 3.2; however, the search cost large (>25 GPU days). Ultimately, we find that DCLM is either ineffective or impractical for performance evaluation.
We also evaluate the efficacy of small-scale synthetic datasets as the evaluation mechanism: MAD (Poli et al., 2024) and BabiStories (Zhang et al., 2025) (Appendix D.6 details our experience with BabiStories). With MAD, we conducted 6-layer and 16-layer searches, extrapolating the pro- duced LLMs to IB parameters via stacking or stretching, respectively. MAD reduced search cost by 8x compared to Large-Scale DCLM while producing hybrid LLMs that consistently outperform Llama 3.2 at scale. We suspect this for two reasons: its token-manipulation tasks are (1) learnable by small models since they have small vocabulary size, and (2) representative of large-scale LLM tasks. This enables efficient small-scale search with strong large-scale performance. While further study of the proxy dataset quality is required, Composer's Evaluator uses MAD throughout this work.
3.3 EXPLORATION OF EXTRAPOLATION AND SCALING METHODOLOGIES
We study the efficacy of stacking versus stretching to extrapolate small hybrid LLMs to larger sizes. We perform search with varying the number of layers Composer searches over from 4 to 32. We then stretch and stack the discovered hybrid LLMs to IB scale, and pre-train them with DCLM for 4e20 FLOPs. Figure 4-top reports DCLM validation loss for each model variant.

---

Published as a conference paper at ICLR 2026
Stacking is an extrapolation mechanism that consistently produces high-quality hybrid LLMs across different n-layer searches. However, stretching hybrid LLMs discovered from larger n-layer search configurations enables Composer to discover more creative interleavings of computational primitives, resulting in higher quality hybrid LLMs.
Observation 4:
Generally, stacking as a mechanism for extrapolation works well regardless of the number of layers searched. Stretching does not work well when conducting search over small number of layers (e.g.. 2A + 4M), as creates hybrid LLMs where Attention dominates beginning lay- ers and MLP dominates end layers without any good interleaving pattern in the middle (e.g.. 10A + 20M at scale). However, an inflection point occurs where stretching consistently out- performs stacking beyond 16-layer searches. Expanding the search space with 16-layers allows Composer to creatively find new interleaving patterns that it cannot explore with the small, re- stricted search space from 6-layer search. Moreover, stretching preserves information by main- taining transitions from one computational primitive to another, enabling the hybrid architec- ture to better capture global information (i.e., propagate signals and gradients across transition points to capture more complex dependencies). However, conducting search beyond 16 lay- ers produces lower quality hybrid LLMs; we hypothesize this is due to larger number of lay- exponentially increasing the search space, making it challenging for Composer to discover higher quality models within the fixed trial budget used for search. Per Figure 4, we choose two different depths to perform search with: stack 6-layer searches and stretch 16-layer searches.
DCLM Vel. Loss
2.0
Stacking Stretching
Observation 5: Scaling down both model width and depth greatly reduces search cost by >6x while also discovering hybrid LLMs that perform better at scale. This is because con- ducting small-scale search with similar width- to-depth ratios as target model sizes preserves model characteristics during search and, thus, increases the likelihood of high model quality at the target scale.
2.8
2.7
10
30
# of Layers Searched
Liama 3.2
No Width Scaling
Width Scaling
3.4
Search Cost (GPU Hre)
DCUM Val. Loss
3.2
200
3.0
130
2.8
Figure 4-top shows the importance of scaling down depth from the target model size to re- duce the design space and discover quality hybrid LLMs (e.g., 32 versus 16 layers stretching). We also study the importance of scaling down width. Figure 4-bottom compares the model quality and search cost of 6-layer search with width scaling (128x258, 16 heads) and without width scaling (Llama 3.2 IB width).
2.8
Training Budget (FLOPs 1e20)
Figure 4: [Top] Model quality of different hy- brid LLMs produced after stacking and stretch- ing with various search depths. [Bottom] Model quality and search cost with and without width scaling.
Scaling down the width reduces search cost by 6.38x. Without width scaling, Composer conducts search with model sizes 100M-200M parameters. This increases the time to train/evaluate each candidate hybrid model, inflating the total search cost. Moreover, scaling down width enables Com- poser to produce higher quality models (validation loss reduces by 0.02-0.04). Without scaling down the width, Composer's candidate hybrid architectures during search have greatly skewed width to depth ratios: too wide and shallow. Hence, the hybrid architectures with the best performance at small scale do not translate the best performance at scale. For example, without width scaling, Composer finds that the architecture 3A + 3M (3 Attention followed by 3 MLPs) performs best at small scale. However, when also scaling down width, Composer finds a better architecture with 1:2 iitention-to-Mil ratio (2A + 4M) that outperforms the 1:1 ratio model at large scale.
EVALUATION
With finalized methodologies (Table 1), we conduct a 6-layer and 16-layer search and discover two unique hybrid LLMs:
6-Layer Search Hybrid L.M 2.4 + 4M
7

---

Published as a conference paper at ICLR 2026
Liama 3.2
Composer: Stacking
Composer: Stretching
Model 3 350M
Model Size:
Model
DCLM Val. Loss
3.2
3.25
3.5
3.00
3.0
3.0
2.75
Training Budget
Training Budget
Training Budget
(FLOPs 1e20)
(FLOPs 1e20)
(FLOPs 1e20)
Figure 5: Validation loss across model sizes (350M-3B) and training budgets (2e19-4c20 FLOPs).
Llama
Composer Stacking
Composer Stretching Inf Tput (tps)
Train Tput
Time/Step (s)
KV Cache (GB)
Inf Lat (s)
100
1.0
toli's x 1K
150
1000
2
0.5
100
500
50
0.0
0
20 Seq Len (x1K)
Model Size (B)
Model Size (B)
Seq Len (x1K)
Seq Len (x1K)
Figure 6: Composite LLMs' training and efficiency improvements compared to Llama 3.2.
16-Layer Search Hybrid LLM 2A 5M 2.4 3M 1A 3M
(5)
We extrapolate these architectures to various sizes via stacking and stretching, respectively. We refer to any variant of Equation 4 and Equation as Stacked and Stretched Composite LLMs, respectively.
4.1 ISOFLOP ANALYSIS AND EVALUATION ON DOWNSTREAM TASKS
We evaluate our Composite LLMs against Llama 3.2 across wide range of model sizes (350M-B)) -BB)) and training budgets (2c19 4c20 FLOPs) show our hybrid LLMs maintain predictive perfor- mance as scale increases. Figure 5 presents DCLM validation loss across model sizes and training budgets. Appendix D.4 reports performance across popular downstream tasks. We also include scaling analysis of our Composite models in Appendix D.I.
Key Result 1: Composite models are robust across model sizes, training budgets, and downstream
tasks. They consistently reduce loss over Llama 3.2 by 0.05-1.0 and outperform Llama 3.2 all the downstream tasks with performance improvements to 2.8-8.3% (1.1-3.1% avg stacked and stretched).
4.2 TRAINING AND INFERENCE EFFICIENCY EVALUATION
Figure 6 analyzes the training and inference efficiency of Composite LLMs. We include inference efficiency metrics at IB scale with batch size Appendix D.3 details inference data across model sizes and batch sizes. Composite models have fewer layers than Llama 3.2: stacking has 27 layers and stretching has 29 layers, while Llama 3.2 has 32 layers (16 Transformer blocks). Moreover, Composite models have 1:2 Attention-to-MLP ratio (9-10 Attention layers), unlike the 1:1 ratio with Transformer architectures (16 Attention layers). These two properties greatly improve the efficiency of Composite models compared to Transformer based LLMs.
Key Result 2: Composite models consistently improve training and inference efficiency. Com- pared to Llama 3.2, across model sizes, we increase training throughput by 1.25x, reducing per- step training time by 1.32x average. Across sequence lengths, at IB scale, improve infer- ence latency by x.33x on average. Moreover, as a byproduct of having fewer Attention layers, we also reduce KV cache size by 1.69
8

---

Published as a conference paper at ICLR 2026
4.3 COMPARISON AGAINST OHATE-OF-THE- PREVIOUS WORKS
Next, we compare our hybrid architectures against three previous state-of-the-art works at 1B size: Sandwich Transformer (Press et al., 2020), Striped Attention (Poli et al., 2024), and the best per- forming architecture from STAR (Thomas et al., 2025). Sandwich Transformer has 1:1 Attention- to-MLP ratio with a rearranged "sandwich" interleaving pattern: 8 Attention, followed by 8 MLP and Attention sequentially interleaved, ending with 8 MLP. The Striped Attention models apply 1:2, 1:4, or 1:8 ratio of Attention to MLP layers, stacking multiple blocks of 1A + 2M (or 4M/8M) to desired sizes. STAR's best model consists of mix of gated convolution, MLP, self-attention, and recurrent layers. Further architectural details are provided in Appendix B.6.
Model
Loss
Are C.
Hella.
Wino.
SciQ
PIQA
Arc E.
Avg.
Llama 3.2 Sand. Transformer
2.80
29.8
53.1
55.8
80.6
71.8
61.03
58.69
2.77
30.8
54.93
55.25
83.4
71.5
63.43
59.88
1:2 Striped Attn.
2.81
29.0
52.9
56.4
80.0
72.6
62.92
58.97
1:4 Striped Attn.
2.82
29.8
51.9
53.8
78.3
71.9
63.09
58.13
1:8 Striped Attn.
2.85
30.7
50.9
52.3
75.7
71.9
62.58
57.35
STAR*
27.9
52.6
53.9
87
71.8
60.8
59
Composite: Stacked
2.77
28.84
54.56
55.72
87.6
73.56
64.73
60.83
Composite: Stretched
2.77
32.25
54.96
53.9
87.9
72.3
63.26
60.76
Table 2: Comparison of Composer's stacked and stretched hybrid architectures to LLM architectures from previous works: Sandwich Transformer (Press et al., 2020), Striped Attention (Poli et al., 2024), and STAR (Thomas et al., 2025) at 1B scale. Note that for Wino. and Are E., accuracy is reported. For other tasks, normalized accuracy is reported.
We pre-train all LLMs with DCLM, except STAR, which could not be pre-trained since the hybrid model is not open-sourced. To make as fair of a comparison with STAR, we match its pre-training setup by training all models with 37.5B tokens (see Appendix C.S for details). Table reports each model's performance the same downstream tasks STAR, with STAR's 's results (which is pre-trained on a different dataset) taken directly from their paper.
Key Result 3: Composite LLMs outperform state-of-the-ard hybrid LLMs when trained with fixed number of tokens, reducing loss by 0.03 while increasing accuracy across the downstream tasks by up to 3.7% (1-2% avg).
4.4 ROBUSTNESS OF COMPOSER'S SEARCH FRAMEWORK
(1) Comparison against randomly generated hybrid LLMs.
We show Composer's robustness by comparing Composite LLMs against five randomly generated 16-layer hybrid LLMs stretched to IB scale. Ap- pendix B.7 details each architecture. The Composite LLMM consistently outperforms the randomly gener- ated architectures across training budgets (Figure 7- top). Those beginning with MLP layers (RI/R3) perform poorly with larger budgets. The remain- ing three architectures perform relatively well, but heavily skew with Attention layers. Composer's 1:2 Attention-to-MLL ratio provides the best model quality while also improving training and inference efficiency ($ 4.2).
6/ 3.00
3.25
R1
R3
R5
R2
R4
Com
poser
2.75
Val. Loss
3.5
p25
p200
p50
Composer
3.0
p75
Training Budget (FLOPs x 1e20)
(2) Comparison of relative rankings at small and large scale. We compare the relative rankings of candidate architectures evaluated during small-scale search and large-scale pre-training show Com- poser's ss robustness. We rank the candidate LLMs after search and pre-train the p0-p100 L.MSs 1B
Figure 7: [Topp] Model quality of the Com- posite LLM versus five randomly generated ones at IB scale. [Bottom] Performance of ranked candidate architectures during search (p25-p100) at IB scale. Relative rankings between small and large scale hold.

---

Published as a conference paper at ICLR 2026
size. Figure 7 shows each architecture's validation loss after pre- pre-training. Appendix B.8 details each architecture.
The Spearman rank correlation between small-scale 6-layer search and IB scale is 0.97 on average across compute budgets, indicating high positive correlation (nearly identical rankings). The pO (not shown, loss is >55) and p25 LLMs perform poorly: they consist of only MLP or only Attention layers, respectively. The p50 and p000 LLMs are all Attention-heavy (2:1 Attention-to- toML ratio); these architectures do not perform well as the best stacked architecture Composer produces with 1:2 ratio. The p75 has a 1:2 Attention-to-MLP ratio, however, the layer interleavings are sub-optimal: the model begins with MLP layers and ends with Attention. Beginning with multiple Attention layers enables deep contextual understanding and feature extraction, while ending with MLP refines and projects features into accurate outputs. Composer's L.MM after Ng clustering adheres to both properties and provides superior model quality, demonstrating the advantage of No clustering in smoothing out noise or overfitting that may occur during small-scale search.
DISCUSSION
Our study demonstrates Composer's efficacy in producing high-quality hybrid LLMs that achieve strong performance at scale on standard natural language understanding downstream tasks, such as PIQA (Bisk et al., 2020), WinoGrande (Sakaguchi et al., 2021), etc. Further work remains to understand Composer's s efficacy for settings that require extended contexts or complex reasoning.
For example, Composite LLMs outperforms Llama on long-context or reasoning tasks, and how our HNAS framework should be augmented to target these specific types of tasks needs further study. Specifically, we may require new small scale long-context or reasoning-speciffe datasets to effectively probe the performance of small hybrid LLMs for these tasks at large scale. However, as we find that sampled-down web-scale datasets are either impractical or ineffective for efficient small-scale search, may require new roken-manipuland specifically targeting long-context or reasoning tasks. Ultimately, determining the best approach for dataset design and model evaluation in these domains remains an open question and promising direction for future research.
datasets, like MAD (Poli et al., 2024),
CONCLUSION
We present Composer, a search framework that systematically discovers novel hybrid neural archi- tectures that improve model quality scale. We propose and evaluate the efficacy of several method- ologies for each step in Composer's flow to build the framework in a principled manner. Specifically, study different search, evaluation, aggregation, and extrapolation techniques, to efficiently duct small-scale search and produce hybrid model architectures that continue to perform well at 1000x larger than their searched size. Composite architectures outperform several state-of-the-an hybrid LLM architectures and Llama 3.2 in validation loss, downstream benchmarks, and model efficiency. Composer's efficient and extensible framework opens up possibilities to incorporate dif- ferent computational primitives-See such as Gated Delta Net, Mamba, Sliding Window Attention- in the search process and uncover new high-performance model architecture designs.
ACKNOWLEDGMENTS
We thank Samuel Hsia, Michael Kuchnik, and Geet Sethi from the FAIR SysML team for their valuable discussions and feedback on this work.
REFERENCES
Joshua Ainslie, James Lee-Thorp, Michiel De Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. arXiv preprint 233.v:2305.1322 2023.
Sangmin Bae, Yujin Kim, Reza Bayat, Sungnyun Kim, Jiyoun Ha, Tal Schuster, Adam Fisch, Hrayr Harutyunan,, Ziwei Ji, Aaron Courville, and Se-Yo Young Yun. M Learning
10

---

Published as a conference paper at ICLR 2026
dynamic recursive depths for adaptive token-level computation, 2025. URL https: /artps. org/ hat//bb//a 2507 10524.
Akhiad Bercovich, Itay Levy, Izik Golan, Mohammad Dabbah, Ran El-Yaniv, Omri Puny, Ido Galil, Zach Moshe, Tomer Ronen, Najeeb Nabwani, et al. Llama-nemotronn: Efficient reasoning models. arXiv preprint FrXx::5055000499 2025.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the A12 reasoning challenge. CoRR, 0tss/8033005557 2018. URL http: //tpxvaavvvvvvaa///////../aaaav. org/ abs / 8038 05457.
Team Cohere, Arash Ahmadian, Marwan Ahmed, Jay Alammar, Milad Alizadeh, Yazeed Alnumay, Sophia Althammer, Arkady Arkhangorodsky, Viraat Aryabumi, Dennis Aumiller, et al. Command a: An enterprise-ready large language model. arXiv preprint arXiv:2504.0068 2025.
Tri Dao and Albert Gu. Transformers are ssms: generalized models and efficient algorithms through structured state space duality. In Proceedings of the 41st International Conference on Machine Learning. ICML'24. JMg..org. 2024.
DeepSeek-AL. Deepseek-v3: contig 67 1w...s...nnn.........ooooooooond
ai/
joo.: 2025. Accessed: 2025-08-17.
Jacob Devlin, Ming-Woo Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume I (Long and Short Papers), pp. 4171- 4186. Association for Computational Linguistics, 2019. doi: 10.19653/ URL 1423.
https: /toi/iioooo// wwwoorggo//..//. 1tt./11.116665 1ee.111661lll1
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. The language model evaluation harness, 07 2024. URL htt.s: www.do.eoooooeee//eee www.or/ccrrr///10000000 502.
Yuxian Gu Qinghao Hu, Shang Yang, Haocheng Xi, Junyu Chen, Song Han, and Han Cai. Jet-nemotron: Efficient language model with post neural architecture search. arXiv preprint arXiv:2508. 15884, 2025.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Oriol Vinyals, Jack W. Rae, and Laurent Sifre. Training compute-optimal large language models. Proceedings of the 36th International Con- ference on Neural Information P Processing Systems, NIPS '22, Red Hook, NY, USA, 2022. Curran Associates Inc. ISBN 978171387 1088.
Andrew Howard, Mark Sandler, Bo Chen, Weijun Wang, Liang-Chieh Chen, Mingxing Tan, Grace Chu, Vijay Vasudevan, Yukun Zhu, Ruoming Pang, Hartwig Adam, and Quoe Le. Searching for MobileNetV3, In 2019 IEEF/CVF International Conference on Computer Vision (ICCV), pp. 1314-1324, Los Alamitos, CA, USA, November 2019. IEEE Computer Society. doi: 10. 1109/CCVV.011.000440 URL http:: ///oi//iiii ieeecom wwwp:t........1.1.1 1109/ ICCV 2019.00140.
Mojan Javaheripi, Gustavo de Rosa, Subhabrata Mukherjee, Shital Shah, Tomasz Religa, Caio Cesar Teodoro Mendes, Sebastien Bubeck. Farinaz Koushanfar, and Debadeepta Dey. Litetransform- ersearch: Training-fre neural architecture search for efficient language models. Advances in Neural Information Processing Systems, 35:24254-24267. 2022a.

---

Published as a conference paper at ICLR 2026
Mojan Javaheripi, Gustavo H. de Rosa, Subhabrata Mukherjee, Shital Shah, Tomasz L. Religa, Caio C.T. Mendes, Sebastien Bubeck. Farinaz Koushanfar, and Debadeepta Dey. Litetransform- ersearch: training-free neural architecture search for efficient language models. In Proceedings of the 36th International Conference on Newral Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2022b. Curran Associates Inc. ISBN 9781713871088.
Ganesh Jawahar, Subhabrata Mukherjee, Xiaodong Liu, Young Jin Kim, Muhammad Abdul- Mageed, Laks VS Lakshmanan, Ahmed Hassan Awadallah, Sebastien Bubeck, and Jianfeng Gao. of-experts with adaptive computation for efficient neural ma- chine translation. arXiv preprint arXiv:2210.07533 2022.
Automoe: Heterogeneous
Nikita Klyuchnikov, Ilya Trofimov, Ekaterina Artemova, Mikhail Salnikov, Maxim Fedorov, Alexander Filippov, and Evgeny Burnaev. Nas-bench-nlp neural architecture search benchmark for natural language processing. IEEE Access,
Jeffrey Li. Alex Fang, Georgios Smyrnis, Maor Ivgi, Matt Jordan, Samir Yitzhak Gadre, Hritik Bansal, Etash Kumar Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muen- nighoff, Reinhard Heckel, Jean Mercat, Mayee F Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bitton, Marianna Nezhurina, Amro Kamal Mohamed Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Joshua P Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah M Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham M. Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander T Toshev, Stephanie Wang, Dirk Groeneveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alex Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-LM: In search of the next generation of training sets for language models. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. URL https: opwnreveeeiieeeeeeeeeeeeeeeeeeee nnt/ /tee/foooz forum C-CNNWWin IE.
Sheng Li, Mingxing Tan, Ruoming Pang, Andrew Li, Liqun Cheng, Quoc V. Le, and Norman P. Jouppi. Searching for fast model families on datacenter accelerators. In 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8081-8091, 2021. doi: 10. 10.1109/CVPR46437 000.772222.0000 1.0099.
Yuanzhi Li and Ronen Eldan. Tinystories: How small can language models be and still speak coherent english, 2024. URL
n http:///oorrmmiiiiiiiiiseee
Wanchao Liang, Tianyu Liu, Less Wright, Will Constable, Andrew Gu, Chien-Chin Huang. Iris Zhang, Wei Feng, Howard Huang, Junjie Wang, Sanket Purandare, Gokul Nadathur, and Stratos Idreos. Torchtitan: One-stop pytorch native solution for production ready LLM pretraining. In The Thirteenth International Conference on Learning Representations, 2025. URL https: // oppenreview. net fott//foummmm wwwoormmmddsSSNNMNAA SFN66N5
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz. et al. Jamba: A hybrid transformer- mamba language model. arXiv preprint arXiv:2403.19887 2024.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv
preprint arXiv:1806.0900 2018.
Lotfi Abdelkrim Mecharbat, Hadjer Benmeziane, Hamza Ouarnoughi, and Smail Niar. Hyt-nas: Hybrid transformers neural architecture search for edge devices. In Proceedings of the 2023 Workshop on Compilers, Deployment, and Tooling for Edge AI, pp. 41-45, 2023.
Abhinav Mehrotra, Alberto Gil CP Ramos, Sourav Bhattacharya, Lukasz Dudziak, Ravichander Vipperla, Thomas Chau, Mohamed Abdelfattah, Samin Ishtiaq, and Nicholas Donald Lane. Reproducible neural architecture search for speech recognition. In International Conference on Learning Representations, 2020.
Meta. The llama 4 herd: The beginning of a new era of natively multimodal innovation, 2025. URL https: / //ai wwwa. www.boooo////oocoooccooooooooooo llama 4 multimodal- wer.iiodal-inteligeeeee inteeligenee/.......
Inc. Meta Platforms. Adaptive experimentation platform, 2024. URL https: / //x. dev/.
12

---

Published as a conference paper at ICLR 2026
Miles Olson, Elizabeth Santorella, Louis C Tiao, Sait Cakmak, Mia Garrard, Samuel Daulton, Zhiyuan Jerry Lin, Sebastian Ament, Bernard Beckerman, Eric Onofrey, et al. Ax: a platform for adaptive experimentation. In AutoML 2025 ABCD Track, 2025.
Hieu Pham, Melody Y. Guan, Barret Zoph, Quoe V. Le, and Jeff Dean. Efficient neural architecture search via parameter sharing. In ICML, 2018.
Michael Poli, Armin W. Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, Ce Zhang, and Stefano Massaroli. Mechanistic design and scaling of hybrid architectures. In ICML, 2024. URL https: https://cppnrevieww net forum? dd=GDpGGddnn..
Ofir Press, Noah Smith, and Omer Levy. Improving transformer models by reordering their sublayers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 2996-3005, Online, July 2020. Association for Computational Linguisties. doi: URL https: hwwp://wwwwww.....////w.wow..... aclweb. http//annooooloy//2222
10. 1/22655//12222.allmm.aa acl ael-main.22 270.
Qwen. Qwen3-next:
Towards ultimate training & inference efficiency, 2025. URL https:
//gpee..aabblolg2200440472022000 20t.:/ccc8800911150228888568 (ee.:0000880000828888888308 7ddoddd27cccll from=
research. Lasearchh addaeettttaddaanneeeeentee. list.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: adver- sarial winograd schema challenge at scale. Commun. ACM, 64(9):99-106, August 2021. ISSN 0001-0782. doi: 10.1145/3474381. URL atttps://dot org/ 1145/34743881
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.0520,, 2020.
J Su, Y Lu, Pan, Murtadha, B Wen, and Liu Roformer. Enhanced transformer with rotary position embedding.. 2021. DOI: hrips://doi. org/10. 1016/.. neucom, 2023.
Mingxing Tan and Quoe V. Le. Efficientnetv2: Smaller models and faster training. In International Conference on Machine Learning. 2021. URL https: : /apii wwe.emnniccccoo...oooooo .org/
Armin W Thomas, Rom Pariickunn, Alexander Amini, Stefano Massaroli, and Michael Poli. STAR: Synthesis of tailored architectures. The Thirteenth International Conference on Learning Rep- resentations, 2025. URL https: http///peeeveieeeeeee/eeeeeeeeee net wwwtteeruuteeeeee forum? Hos:HohhHH32::::
Pavan Kumar Anasosalu Vasu, James Gabriel, Jeff Zhu, Oncel Tuzel, and Anurag Ranjan. Fastvit: A fast hybrid vision transformer using structural reparameterization. In ICCV, 2023. URL https: a230 14189.
//aaxivvorrg.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Éukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31 Inter- national Conference on Neural Information Processing Systems, NIPS' 17, pp. 6000-6010, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781 7781510860964
Irene Wang, Newsha Ardalani, Mostafa Elhoushi, Daniel Jiang, Samuel Hsia, Ekin Sumbul, Divya Mahajan, Carole-Jean Wu, and Bilge Acun. Carbon aware transformers through joint model- hardware optimization. arXiv preprini arXiv:2505.01386, 2025.
Johannes Welbl, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. W-NUT 2017, pp. 94, 2017.
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 preprint rrXiiv2301.08722 2023.
Chen Xing, Devansh Arpit, Christos Tsirigotis, and Yoshua Bengio. A walk with sgd, 2018. URL https://arxiie wwtp://wwwwg/o////oo/oo/ooooossp 802. 08770.
13

---

Published as a conference paper at ICLR 2026
Jin Xu, Xu Tan, Renqian Luo, Kaitao Song, Jian Li, Tao Qin, and Tie-Yan Liu. Nas-bert: Task- agnostic and adaptive-size bert compression with neural architecture search. Proceedings of the 27th ACM SIGKDD Conference Knowledge Discovery & Data Mining, pp. 1933-1943, 2021.
Chris Ying, Aaron Klein, Eric Christiansen, Esteban Real, Kevin Murphy, and Frank Hutter. Nas- bench-101: Towards reproducible neural architecture search. International conference on ma- chine learning, pp. 7105-7114, PMLR, 2019.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in neural infor- mation processing systems, 32, 2019.
Jianyu Zhang, Niklas Nolte, Ranajoy Sadhukhan, Beidi Chen, and Leon Bottou. Memory mosaics.
In The Thirteenth International Conference on Learning Representations, 2025. URL https:
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, and Shen Li. Pytorch fsdp: Ex- periences on scaling fully sharded data parallel. Proc. VLDB Endow., 16(12):3848-38)) -8638, August 2023. ISSN 2150-8097, doi: (0.14777/3611540.3661 1569. URL https: /ttps://ddot www.ddi..//o///eeeeeeeeeoooeeeee 14778/3611540 3611569.
Barret Zoph, Vijay Vasudevan, Jonathon Shlens, and Quoc V Le. Learning transferable architectures for scalable image recognition. In Proceedings of the IEEE conference on computer vision and pultern recognition, pp. 8697-8710, 2018.
14

---

Published as a conference paper at ICLR 2026
A RELATED WORK
Hybrid LLM Architectures. Hybrid LLM architectures have become increasingly popular to im- prove the quality and efficiency of LLMs. Several hybrid LLMs build new stackable blocks (anal- ogous to Transformer block (Vaswani et al.. 2017)) by adding new computational primitives and adjusting the ratio of primitives within a block. For example, Qwen3-Next (Qwen, 2025), Mamba- 2 (Dao & Gu, 2024), Jamba (Lieber et al., 2024) and MAD (Poli et al., 2024) all adjust the ratio of Transformer and State Space Model (SSM) primitives within a block, skewing the composition to a single computational primitive. Other hybrid LLMs consist of more sophisticated interleavings of computational primitives by breaking the conventional stacking-base
structure. Command-A (Co- here et al., 2025) introduces bybridization by interleaving sliding- window and full attention layers in 3:1 ratio. Similarly, LJama-4 models (Meta, 2025) use interleaved attention layers without positional emheddings. DeepSeek- V3 671B model (DeepSeek-Al, 2025) incorporates three dense MIPs in its initial layers followed by sparsely activated MoEs. FistViT (Vasu et al., 2023) leverages convolu- tions in the beginning layers and attention in later stages. Sandwich Transformer (Press et al., 2020) maintains 1:1 Attention-to-MLP ratio, but reorders the interleavings of the primitives.
Neural Architecture Search Frameworks. Neural Architecture Search more generally encom- passes the art of automating the time consuming process of neural architecture design for given task. There is wealth of literature that studies how to (1) design an effective search space for a task (2) search through the space efficiently and (3) estimate the performance
given architec- ture quickly. A popular class of search spaces are cell-based search spaces (White et al., 2023). In cell based search spaces, the search is performed over small cells and stacked several times in sequence to form an architecture (Zoph et al., 2018). Bayesian Optimization, reinforcement learn- ing, and evolutionary search are classes of search strategies that have been shown to perform better than random search. A performance predictor is defined any function which predicts the accu- racy or relative accuracy of architectures, without fully training them. Various kinds of proxies exist including: learning curve extrapolation, zero-cost proxies and subset selection methods. NAS has been studied extensively for older and smaller models including CNN (ResNet style) models, NLP, and ASR models (Pham et al., 2018; Liu et al., 2018; Howard et al., 2019; Li al., 2021: Tan & Le, 2021; Mecharbat et al., 2023). There have been various benchmarks released incluuding NAS- Bench- 101 (Ying et al., 2019), NAS-BERT (Xu et al., 2021), MAS-Bench-ASSH (Mehrotra et al., 2020), MAS-Bench-NLP (Klyuchnikov et al., 2022), to make NAS more reproducible. However, NAS has not been as well studied for modern Large Language Models (LLMs). One of the reasons is that the scale of the models makes it more difficult to design effective performance proxies that extrapolate to large scale. Some attempts include AutoMOE (Jawahar et al., 2022) that uses supernet training combined with evolutionary search to design MoE architecture for machine translation and LiierraasooomerScaccch (Javaheripi et al., 2022a) that uses a zero cost performance proxy (the number of model parameters). However, the models searched over are smaller and older (OPT) models. More importantly, these traditional neural architecture search frameworks assume fixed interleavings/ratios of computational primitives when searching over model hyperparameters such as model width, number of layers, attention heads, output dimensions. Such zero-cost performance proxies may not apply when searching on different model architecture variants.
Neural Architecture Search Frameworks for Hybrid Models. There are only few recent works that consider the hybrid model architecture design space in their search. The Nemotron model family builds a Post Neural Architecture Search (PostNAS) framework that prunes some of the modeling blocks or replaces Global-Attentic blocks of pre-trained models with efficient Attention variants (Bercovich et al., 2025; Gu et al., 2025). This optimization in post-training setting targets a completely different problem than designing new hybrid architecture for pre-training and, it does not reduce the large pre-training and experimentation costs of LLMs. STAR (Thomas et al., 2025) presents an initial attempt towards a framework targeting pre-training hybrid LLMs from scratch, however, its design assumes conducting search on the target dataset for edge use cases. We find that conducting search on web-scale datasets for LLMs is either ineffective or impractical for per- formance evaluation.
15

---

Published as a conference paper at ICLR 2026
B HYBRID LLM ARCHITECTURE DETAILS
With our design exploration of Composer's 's components in $ and evaluation against Llama 3.2 and other state-of-the-ard previous works, evaluate hundreds of unique LLM architecture designs. In this section, we provide architectural details for each evaluated LLM.
Any LLM produced by Composer consists of Grouped-Query Attention with rotary positional em- beddings (RoPE) (Su et al., 2023) (referred to as Attention or A for brevity) as the sequence mixing layer and SwiGLU (Shazeer, 2020) (referred as MLP or M) as the channel mixing layer. We use RMSNorm (Zhang & Sennrich, 2019) for normalization and no linear bias term. All models tie the embedding layers.
Model Size
Dim
Hidden Dim
Num Heads
Num KV Heads
350M
1536
4096
24
1B
2048
8192
32
8
3B
3072
8192
24
8
8B
4096
14336
32
Table 3: Model dimensions for stretched and stacked Composite architectures, along with Llama 3.2 which we pre-train from scratch.
Table 3 describes the width dimensions of our models under different sizes. Below, we detail each architecture. We organize the descriptions section by section following and 4 in our paper.
B. SEARCH METHODOLOGY EXPLORATION: HYBRID LLM ARCHITECTURES
In $ 3.1, explore the efficacy of three search methodologies: One-Shot Search, End-Layer In- cremental Search, and Middle-Layer Incremental Search. We describe the Composite architecture produced by each search methodology before and after extrapolation to 1B scale.
One-Shot 6-Layer Search
Small scale (400K parameters): 2A 4M
Large scale (IB parameters): 4x(2A44M) + IA + 2M
One-Shot 16-Layer Search
Small scale (1. 1M parameters): 2A + 5M + 2A + 3M + 1A 3M
Large scale (1B parameters): 4A + 9M + 4A + 5M + 2A 5M
End-Layer Incremental Search
Small scale (1M parameters): 2A + 2M 2A 2M 4A IM 3A + 2M 1A 1M+
1A + 6M + 1A + 3M + 1A
Large scale (1B parameters): 2A 2M 2A + 2M 4A + IM 3A 2M IA IM+
1A 6M 1A + 3M + 1A
Middle-Layer Incremental Search
Small scale (1.1M parameters): 3A + 1M 1A 3M + 2A 1M+ 1A 3M 1A 2M
3A + 2M 1A + IM + 1A + 6M
Large scale (1B parameters): 3A IM 1A 3M 2A + IM 1A 3M 1A 2M
3A + 2M + 1A + IM + 1A + 6M
B.2 DATASETS EXPLORATION: HYBRID LLM ARCHITECTURES
In $ 3.1, we explore the efficacy of three datasets for small-scale search: MAD (Poli et al., 2024), Sampled-DCLM (Li al., 2024), and BabiStories (Zhang et al., 2025). We describe the hybrid LLM architecture produced by each dataset before and after extrapolation to IB scale.
MAD Stacking
Small scale (400K parameters): 2A + 4M
Large scale (IB parameters): (x 1A 2M
16

---

Published as a conference paper at ICLR 2026
MAD Stretching
Small scale (1.1M parameters): 2A + 5M + 2A +M IA 3M
Large scale (1B parameters): 4A + 9M 4A + 5M + 2A 5M
Small-Scale DCLM
Small scale (2.9M parameters): 1M 1A 2M
Largg scale (IB parameters); 6x + 1A + 2M) + IM + IM
Large-Scale DCLM
Small scale (500M parameters): 1A + 2M + 3A + 3M + 1A+ IM + 2A 3M
Large scale (IB parameters): 2A + 4M + 6A + 6M + 2A 2M 4A
Small-Scale BabiStories
Small scale (1M parameters): 14M + + IM
Large scale (IB parameters): 19M 2A 2M
Large-Scale BabiStories
Small scale (150M parameters): 4A + 1M 4A IM + +3 3M
Largge scale (IB parameters): 11A 3M 1IA + 3M + 9A + 9M
B.3 AGGREGATION EXPLORATION: HYBRID LLM ARCHITECTURES
In Appendix D.5, explore the efficacy of different aggregation techniques with Ne clustering. Here, we detail the hybrid LLM architecture produced by each technique for a 6-layer (400K pa- rameters) and 16-layer search (1M parameters).
6-Layer Search
No architecture: 2A 4M
N1 architecture: 2A + 1M 1A 2M
N architecture: 2A 1M +3A
p100 architecture: 4A + 2M
16-Layer Search
No architecture: 2A 5M 2A 3M 1A + 3M
Ni architecture: 2A 5M + +A + 6M 1A
Ni-1 architecture: 2A 4M +3A + 2M + 4A + IM
p100 architecture: 3A + 7M + IA 5M
B.4 EXTRAPOLATION METHODOLOGGY: HYBRID L.MM ARCHITECTURES
In $ 3.3, we study the efficacy stretching and stacking as we ablate n, the number of layers for search. We also study the efficacy of with and without width scaling. We describe the Composite architectures produced by each methodology before and after extrapolation to IB scale below.
Layer Search
Small scale (270K parameters): 2M
Large scale via stacking B1B parameters): 8x(2A 2M) 1A 1M
Large scale via stretching (1B parameters): 18A + 18M
Layer Search
Small scale (330K parameters): 2A + 3M
Large scale via stacking (1B parameters): 6x(2A + 3M)
Large scale via stretching (IB parameters): 12A 18M
n Layer Search
Small scale (400K parameters): 2A 4M
17

---

Published as a conference paper at ICLR 2026
Large scale via stacking (1B parameters): 4x(2A+MM) 1A + 2M
Large scale via stretching (IB parameters): 10A+20M
Layer Search
Small scale (465K parameters): 2A + 5M
Large scale via stacking (IB parameters): 3x(2A 5M) 2A 4M
Large scale via stretching (IB parameters): 8A + 20M
7 Layer Search
Small scale (530K parameters): 2A + 4M 1A 1M
Large scale via stacking (IB parameters): (2A + 4M + 1A IM) + 2A + 3M + IA
1M
Large scale via stretching (IB parameters): 8A + 16M 4.A + 4M
n Layer Search
Small scale (600K parameters): 2A 4M + 1A + 2M
Large scale via stacking (IB parameters): 3x(AA + 4M + IA + 2M)
Large scale via stretching (IB parameters): 6A + 12M 3A + 6M
10 Layer Search
Small scale (660K parameters): 2A 4M IA 3M
Large scale via stacking (IB parameters): 2x(2A 4M + 1A 3M) + 2A 3M + 1A
2M
Large scale via stretching (1B parameters): 6A 12M 3A 9M
71 = 11 Layer Search
Small scale (730K parameters): 2A 4M 1A 2M 1A 1M
Large scale via stacking (1B parameters): 2x(2A 4M + 1A 2M + 1A + IM) + 2A +
3M + 1A 2M 1A + 1M
Large scale via stretching (IB parameters): 6A + 11M + 3A + 6M 3A +3 3M
n 12 Layer Search
Small scale (795K parameters): 2A 5M+ + 1A + 4M
Large scale via stacking (1B parameters): 2x(2A 5M + 1A 4M) + 1A + 1M IA
IM
Large scale via stretching (IB parameters): 4A+ 11M + 2A + 8M
n 13 Layer Search
Small scale (860K parameters): 3A 7M 1A 2M
Large scale via stacking (1B parameters): 2x(3A 7M + 1A + 2M) + 1A I
1M
Large scale via stretching (1B parameters): 7A + 15M + 3A + 5M
14 Layer Search
Small scale (925K parameters): 2A 5M 4A 1M + 1A IM
Large scale via stacking (IB parameters): 2x 5M + 4A IM + IA + IM) + 1A
2M + 2A 1M IA + IM
Large scale via stretching (IB parameters): 5A + 12M + 10A + 3M 3A + 3M
15 Layer Search
Small scale (990K parameters): 3A + 2M + 2A + 5M 1A + 2M
Large scale via stacking B1B parameters): 2x(3A + 2M + 2A 5M 1A + 2M)
Large scale via stretching (1B parameters): 6A + 4M 4A 10M + 2A 4M
72 = 16 Layer Search
Small scale (1. parameters): 2A + 5M + 2A 3M 1A + 3M
18

---

Published as a conference paper at ICLR 2026
Large scale via stacking (IB parameters): 2A 5M 2A 3M IA 3M + 2A 4M
2A + 2M 1A + 2M
Large scale via stretching (IB parameters); 4A + 9M + 4A 5M 2A 5M
20 Layer Search
Small scale (1.3M parameters): 2A + 9M + IA 1M 1A 1M 1A 2M 2A
Large scale via stacking (IB parametens): 2A 9M + 1A IM IA + IM + 1A 2M
3A 4M + IM + + IM + IA IM 1A
Large scale via stretching (1B parameters): 3A 13M 2A 2M 2A 2M 2A 3M
+ 3A
24 Layer Search
Small scale (1.6M parameters): 2A + IM + 2A + IM + 1A IM+ 1A 2M + 3A IM+
6A 3M
Large scale via stacking (1B parameters): 2A 1M + 2A 1M 1A + 1M + 1A 2M
3A + 1M + 6A 3M 2A + 1M + 2A + 1M + 1A 1M + 2M 2A + IM + 4A
2M
Large scale via stretching (IB parameters): 4A + 2M + 4A + 2M + 2A + 2M + 2A + 4M 6A 2M 12A 6M
7 28 Layer Search
Small scale (1.9M parameters): 2A 3M 1A 2M 3A 4M 3A 1M 66 2M
1A
Large scale via stacking (1B parameters): 2A 3M + 1A 2M + 3A + 4M + 3 1M+
6A 2M 2A + IM 1A + IM + 1A 2M 1A +M 1M 1A
Large scale via stretching (1B parameters): 3A + 4M + 2A +3M + 4A + 6M + + 2M
8A 3M 2A
7 =
32 Layer Search
Small scale (2.1M parameters): 2A + 3M 1A 1M 1A 1M 2A 5M 1A 3M
IA + 1M 2A + IM + A + 2M + 1A IM + 2A
Large scale via stacking (IB parameters): 2A 3M + IA IM LA + 1M +2 5M+
1A 3M 1A + 1M + 2A + 1M + 1A 2M 1A 1M 2A
Large scale via stretching (1B parameters): 2A + +3M + IA + 1M 1A 1M + 5M
IA 3M 1A + 1M 2A 1M 1A 2M 1A IM 2A
Width Scaling
Small scale (400K parameters): 2A 4M LLarge scale (IB parameters): 4x(2A+4M) + 1A + 2M
No Width Scaling
Small scale (200M parameters): 3A + 3M
Large scale (1B parameters): 5x(3A+3M) 2A 2M
B.5 COMPOSER'S BEST HYBRID LLM ARCHITECTURES ACROSS MODEL SIZES
With finalized methodologies from our design exploration, we conduct a 6-layer and 16-layer search and discover two unique architectures, presented Equation and Equation 5. We extrapolate these architectures to various sizes via stacking and stretching, respectively. Below, we describe each architecture per size.
Stacked Composite Architecture
Small-scale 6-layer search: 2A 4M 350M: (2A 4M)
1B: 4x(2++4M) + 1A + 2M
3B: 8x(2A+4M) + 2M
8B: 10x(2A 4M) 1A 1M
19

---

Published as a conference paper at ICLR 2026
Stretched Composite Architecture
Small-scale 16-layer search: 2A + 5M + 2A 3M IA + 3M
350M: 3A 8M 3A 5M 2A +M
. 1B: 4A 9M + 4A 5M 2A + 5M
3B: 7A + 16M + 7A + 10M + 4A + 10M
8B: 8A 19M + 8A + 12M 4A + 12M
B.6 BASELINES
We compare our two best architectures against several baselines: Llama 3.2, Sandwich Trans- former (Press al., 2020), and Striped Attention (Poli et al., 2024). We describe each of their architectures below.
Llama 3.2
350M: 14x(1A + IM)
1B: 16x(1A + 1M)
3B: 28A(1A + 1M)
8B: 36x(1A + IM)
Striped Attention (1B)
1:2 ratio: 9x(1A 2M)
1:4 ratio: 5x(1A 4M)
1:8 ratio: 2x(1A 8M) 1A + 4M
Sandwich Transformer (IB): 8A + IM + IA+ IM IA IM IA IM+ IA + IM IA IM 1A + 1M + 1A 8M
B.7 RANDOMLY GENERATED MODELS
In 4.4, assess Composer's robustness by comparing the model quality of our Composite LLMs with five randomly generated LLMs. We detail the architecture each randomly generated LLM at 1B scale below.
Random 1: 3M 2A 7M 2A + 10M
Random 2: 2A + 2M + 4A + 4M + 7A + 4M 3A + 4M 3A + 3M
Random 3: 7A + 2M + 2- + 2 + +A 2M 5A 2A + 2M
Random 4: 5M 9A +2M +M +2
9M
Random 5: 5A + 3M + +33 +3A+7M++7 +7 +M 3A +3M+ 5A
B.8 RELATIVE RANKINGS AT SMALL-SCALE SEARCH AND LARGE-SCALE PRE-TRAINING
$ 4.4 and Appendix D.2, we asses Composer's robustness by comparing relative rankings of candidate architectures evaluated during small-scale search and large-scale pre-training. We detail the architecture each randomly generated LLM before and after scaling to IB size.
p0 6-Layer Search
Small-scale (400K parameters): 6M
Large-scale (IB parameters): 20M
p25 L-yayer Search
Small-scale (400K parameters): 6A
Large-scale (1B parameters): 96A
p50 6-Layer Search
Small-scale (400K parameters):3M + 1A + IM+ 1A
Large-scale (1B parameters): (3M + LA + IM + IA) +IM+ 1A + IM + 1A
p75 6-Layer Search
20

---

Published as a conference paper at ICLR 2026
Training Parameter
MAD
Sampled-DCIM
Babistories
Optimizer
AdamW
Adam
AdamW
Optimizer Momentum
81, 82
0.9, 0.98
B1, 02 = 0.9, 0.98
B1, B2 = 0.9, 0.98
Dropout
None
0.05
None 128
Batch Size
128
Vocab Size
16-64
128K
50K
Training Epochs
200
Learning Rate Schedule
Cosine Decay
Cosine Decay 10000
Cosine Decay 927158
Number of Training Samples
800
Number of Evaluation Samples
1280
9275
9275
Parallelism
None
FSDP GPUs
None
Base Learning Rate
[le-4, Se-4, le-3]
5e-3
[le-4, Se-4, le-3]
Weight Decay
[0.0, 0.1]
0.1
[0.0, 0.1]
Table 4: Composer Evaluator's training setup per dataset.
Small-scale (400K parameters): 1M + 2A + 1M + 2A
Large-scale (IB parameters); 7x(IM + 2A + IM + 2A) IM IA
p100 6-Iaaer Search
Small-scale (400K parameters): 4A 2M
Large-scale (IB parameters): 7x(4A 2M) 1A IM
p0 16-Layer Search
Small-scale (IM parameters): 2M 4A +M 4A + IM+ 1A IM
Large-scale (1B parameters): 5M + 9A + 7M + 9A + 3M + 3A 3M
p25 16-Layer Search
Small-scale (1M parameters): 1M 2A + 2M 5A + 1M + 2A + 3M
Large-scale (IB parameters): 3M 5A + 5M + 12A + 3M 5A + 7M
P50 16-Layer Search
Small-scale (IM parameters): 2A + IM + 1A 5M 2 IM+3A- IM
Large-scale (IB parameters); 5A + 3M + 3A 10M + +5 3M + 6A 3M
p75 16-L.ayer Search
Small-scale (400K parameters):2A- 3M 4A + 2M + 1A + 2M + 1A 1M
Large-scale (1M parameters): 5A + 6M + 8A + 5M + 3A + 5M +3A 3M
p100 16-Layer Search
Small-scale (IM parameters): 3A + 7M + 1A + 5M
Large-scale (IB parameters): 5A + 11M 2A + 8M
C EXPERIMENTAL DETAILS
C. SEARCH TRAINING/VALIDATION SETUP
During small-scale search, the HNAS Evaluator evaluates each candidate hybrid LLM architec- ture provided by the HNAS Engine. We provide details on the training and validation setup the HNAS evaluator uses. This setup changes depending the dataset used for evaluation. Table 4 details the training setup for each dataset. We explore three datasets: MAD (Poli et al., 2024), Sampled-DCLie (Li et al., 2024), and BabiStories (Zhang al.. 2025). For MAD and BabiStories, we follow the same training setup (Poli et al., 2024). For Sampled- DCLM, Sampled-DCLM, we reduce the sample size 12K samples.
21

---

Published as a conference paper at ICLR 2026
Model Size
DP
Local Batch Size
Acc
Seq Len
Global Batch Size
Learning Rafe
350M
8
8
8192
0,5B
3e-3
1B
16
4
8192
0,5B
6e-4
3B
16
4
8192
0.5B
3e-4
8B
16
4
8192
0,5B
6c-4
Table 5: Pre-training setup details per model size.
C.2 MODEL SIZE AND FLOPS CALCULATION
We describe size (number of parameters) and FLOPs calculations per computational primitive: Grouped Query Attention (GQA) (Ainslie al., 2023) and SwiGLU (Shazeer, 2020).
Attention: We provide general formulation for calculating FLOOPs (forward and backward pass) and parameter size of single Attention layer that be applied to both GQA MHA layers. We follow DeepMind's calculations (Hoffmann et al., 2022), calculating FLOPs for the forward and backward pass.
dim x num ku heads
Attn Parameter Count
dim x
(6)
num heads
dim
Attn FLOPs = 6 x dim. + x x num_heads
x num kv heads
+ 6 dim² + 2 2 x numhheads x dim
MLP (SwiGLU):
MIP Parameter Count = 3 x dim x hidden_din
(8)
MLP FLOPs = 6 x 3 x dim x hiddeenddi _dmm
(9)
C.3 SCALING LAWS PRE-TRANE TRA-TIIING SETUP
For fair comparison between models, we keep training setups between architecture similar as possible per model size. We pre-train all models using TorchTitan (Liang et al., 2025), adding support for custom LLM architectures. Table details the pre-training setup. To study how our models scale, we compare across four model sizes: 350M, IB, 3B, and 8B models. We use FSDP parallelism (Zhao et al., 2023) to train all models. Models of sized 350M are trained on H200 GPUs, while remaining sizes are trained on 16 H200 GPUs. We linearly interpolate learning rates from common settings, obtaining linear inverse relationship with model size. For models, we train with the DCLM dataset (Li et al., 2024) with 0.5B global batch size and sequence length 8192. We use the Adam W optimizer. Following (Bae et al., 2025). we use a Trapezoidal learning rate scheduler (Xing et al., 2018), with linear warm up and decay phases using 20% the number of training steps each.
To determine the length of training, we use the IsoFLOP methodology, like DeepMind's Chin- cilla (Hoffmann et al., 2022), and fix the training budget (number FLOPs to train with) across hybrid architectures: five budgets of 2e19, 4e19, 8e19, 2e20, 4e20 FLOPs. As the training bud- get and model sizes between architectures are the same, the only differences in training between the models is the number of training tokens. This is because the FLOPSs per token differ between models, since the ratio of computational primitives can vary. Appendix C.2 details equations for cal- culating the FLOPs of single Attention or MLP layer in the forward and backward pass processing one token. The FLOPPs per token of given architecture is then given by:
FLOPs per token (M) (Num Attn layers) x (Attn FLOPs)
(10)
+ (Num MLP layers) x (MLP FLOPs)
The number of training tokens D) for a given model M under given training budget C' FLOPs is
C
D FLOPs per token (M)
(11)
22

---

Published as a conference paper at ICLR 2026
C.4 INFERENCE EFFICIENCY MEASUREMENT PROTOCOL
To evaluate Composer's 's improvements in inference efficiency with its discovered hybird LLMs, we measure the per-token generation time. We measure the generation time under various model configurations using random weights and inputs. We evaluate throughput and inference latency with a variety of sequence lengths - 2K, 4K, 8K, 16K, 32k-where 25% of the sequence length the prompt. We three prompts as warm up phase, and then mcasure the average throughput (tokens per second) and latency across an additional five prompts.
C.5 COMPARISON AGAINST PREVIOUS SOTA WORKS PRE- TRAINING SETUP
In $ 4.3, we follow the pre-training setup of STAR's (Thomas et al., 2025) IB experiment results to fairly compare against them. Rather than fixing the training FLOP budget, we pre-train every model with the same number of tokens instead: 37.5B tokens (see Table 5 of 1B setup). All models are pre-trained for 71564 training steps.
C.6 DOWNSTREAM TASKS
We evaluate validation loss on the test set from the DCLM dataset (Li et al., 2024). Additionally, use the Language Model Evaluation Harness framework (Gao et al., 2024) evaluate accuracy on six few-shot tasks: HellaSwag (Zellers et al., 2019), PIQA (Bisk et al., 2020), WinoGrande (Sak- aguchi et al., 2021), ARC-easy and ARC-challenge (Clark al., 2018), and SciQ (Welbl et al., 2017). We adhere to the standard number of shots specified by the evaluation framework for each task, evaluate performance measurements were conducted on a single H200 GPU.
D ADDITIONAL RESULTS
D.1 SCALING LAWS ANALYSIS
2e+119 FLOPs
4e+19 FLOPs
8e+19 FLOPs
2e+20 FLOPs
4e+20 FLOPs
Llama 3.2
Composer Stretching
Composer Stacking
4.6
4.2
4.2
4.2
Validation Loss
Validation Loss
Validation Loss
3.8
3.8
3.8
3.4
3.4
3.0
3.0
3.0
2.6
2.6
2.6
350M
1B
38
8B
350M
18
3B
8B
350M
1B
3B
8B
Model Size (B)
Model Size (B)
Model Size (B)
Figure 8: Compute-optimal sealing analysis for Llama 3.2 and our Composite LLMss Each star indicates the optimal model size for a given compute budget.
Figure and 9 illustrates that our Composite LLMs exhibit a very similar compute-optimal scaling behavior compared to Llama 3.2 under IsoFLOPs constraints. In Figure 9, the Composite models have nearly identical slopes (Llama has slope of -0.63, where Stretching and Stacking have slopes of -0.628 and -0.612, respectively), indicating stable gains in performance
scale model sizes. While further investigation is required to validate that Composer's ' extrapolation techniques scale to larger model sizes (c.g., tens hundreds of billions of parameters), ciently discovering new hybrid LLMs that can scale 1000 the search size while still outperforming Transformer- based architectures.
work lays blueprint for effi-
23

---

Published as a conference paper at ICLR 2026
3.4
Validation Loss
Composer Stretching
Composer Stacking
3.0
2.6
350M
1B
3B
Model Size (B)
Figure 9: Compute-optimal scaling analysis for Llama 3.2 and our Composite LLMs.
D.2 COMPOSER'S ROBUSTNESS: RELATIVE RANKINGS
In $ 4.4, we show Composer's ' search framework is robust, as the relative rankings of explored 6- layer candidate architectures during small-scale during search and at 1B scale after pre-training are nearly identical. We provide additional results detailing the relative rankings for 16-layer search Figure 10. The Spearman rank correlation between small-scale 16-layer search 0.9 on aver- age across compute budgets, indicating high positive correlation (nearly identical ranking). Ap- pendix B.& details each architecture. This high correlation in relative-ranken rlaakie--ankkngs between scall-scale search and at-scale performance showcases our search framework's robustness.
p0
p50
p100
p25
p75
Composer
3.2
Loss
2.8
Training Budget (FLOPs 1e20)
Figure 10: Performance of ranked candidate architectures during 16-layer search (p0-p100) at 1B scale stretched out. Relative rankings between small and large scale hold.
D.3 EXTENDED INFERENCE EFFICIENCY RESULTS
In $ 4.2, we detail the inference efficiency improvements of our Composite LLM's compared to Llama 3.2 at IB scale with a batch size of 1. We provide further analysis across other model sizes 350M to 3B with batches of 1, 4, 8, and 16 prompts. Figure 11- 13 present the KV cache size (top), inference throughput (middle), and inference latency (bottom) sequence length varies. Across model sizes and batch sizes, our Composite LLMs' continue to reduce KV cache size, increase infernece throughput, and reduce inference latency as sequence length increases.
D.4 PERFORMANCE ON DOWNSTREAM TASKS ACROSS MODEL SIZES
In $ 4.1, we report the DCLM validation loss our Composite LLMs against Llama 3.2 across four sizes: 350M, IB, 3B, and 8B. We provide additional results in Table 6 that shows that across model sizes, our Composite LLMs consistently improve accuracy on downstream tasks (tasks are detailed in Appendix C.6). We pre-train each model with the max training budget (4e20 FLOPs) using the DCLM (Li et al., 2024) dataset. We outperform Llama 3.2 all downstream tasks with performance improvements up 2.8% (1.1-3.1 1..1% avg).
24

---

Published as a conference paper at ICLR 2026
Llama 3.2
Composer Stacking
Composer Stretching
Lateney (sec) Tput (tokens/sed) seensleee)) Cache (GIB)
Batch Size
KV Cache (GIB)
Batch Size
KV Cacha (GIB)
Batch Size 8
Batch Size 16
10
0.5
N
KV Cache
0.0
0
0
20 Batch Size 4
20 Batch Size 8
Batch Size
Lalency (Bec) Tput (lokens/se))
Lalency (880) put (Boken)/Bee)
Latency (sec) Tpu: (lakens//sc)
Batch Size 16
400
1000
1500
2.00
1000
1000
500
Batch Size
Batch Size 4
Batch Size 8
Batch Size 16
500
2000
1000
500
10 Seq Len (x1K)
Seq Len (x1K)
Seq Len (x1K)
Seq Len (X1K)
Figure 11: Comparison of the inference efficiency of Composite LLMs versus Llama 3.2
350M scale. We report KV cache size (top), inference throughput (middle), and inference latency (bottom) as prompt lengths and batch size change.
Composer Stacking
Composer Stretching
kv Cache (GBB
Batch Size
Tput (to Rokens/sac)jy Cache (GBB)
Batch Size
KV Cache (GB)
Batch Size 8
KV Cache (GBB)
Batch Size 16
2.5
10
0
20
20
Latency (sec) Tput (tokensteao)
Batch
Batch
Laternd Stency (Sec) Tput (Bokens/se)) (ec)
Batch 8
Latency (sec) Tput (tokone/sac)
Batch
3.00
1000
200
1000
500
100
500
500
20
20
10
15
2.5
5.0
7.5
Batch Size
Latency (sec)
Batch
Batch
500
2000
1000
500
20 Seq (x1K)
20 Len (x1K)
10
15
2.5
5.0
7.5
Seq Len (x1K)
Seq
Len (x1K)
Figure 12: Comparison of the inference efficiency of our Composite LLMs versus Llama 3.2 at 1B scale. We report KV cache size (top), inference throughput (middle), and inference latency (bottom) as prompt lengths and batch size change.
D.5 EXPLORATION OF AGGREGATION TECHNIQUES
After the HNAS Engine and Evaluator finish searching and evaluating the design space of LLM architectures, the HNAS Aggregator leverages N, clustering to synthesize the final hybrid LLM architecture from the search results. 2.3 explains the algorithm for N. clustering.
In this section, we evaluate three different values for c. When = 0, No clustering selects the dominant block at each layer among the top candidate architectures independently. resulting in no conditioning based prior layers. When 1, N₁ clustering conditions the block choice at each layer on the block selected the immediate preceding layer. Finally, for = 1, the primitive selected at each layer index is conditioned on the entire sequence of previously selected blocks, enforcing full prefix consistency.
To convincingly understand which aggregation technique is best, exhaustively evaluate all hybrid n-layer architectures at small scale using MAD (Poli et al., 2024), from n = 4 to 10. Then, we apply No. N1. and N,-1 for each value of synthesize the final hybrid architecture and scale up IB scale via stacking. We also scale up the p100 (best) model during small-scale evaluation. Finally, include results for 16-layer search, however, do not exhaustively evaluate all 216 hybrid
25

---

Published as a conference paper at ICLR 2026
Liama 3.2
Composer Stacking
Composer Stretching
KV Cacha (GIB)
Batch Size
KV Cacha (GIB)
Batch Size
Tput (tokens/sec) KV Cache (GiB)
Batch Size
Latency (sec) Tput (takens/sec) KV Cache (GiB)
Batch Size
Tput (tokenstaed (sec)
Lateney (sec) Tput (lokens
Batch Size
B Batch Size
Batch Size
Batch Size 16
500
150
750
500
t00
500
50
250
250
Batch Size
Batch Size 4
Batch Size
Batch Size 16
Latency (sec)
1000
2000
1000
2500
Latency
0
2.5
Seq Len (X1K)
Seq Len (x1K)
Seq Len (x1K)
Seq Lan (x1K)
Figure 13: Comparison of the inference efficiency of our Composite LLMs versus Llama 3.2 3B
scale. We report KV cache size (top), inference throughput (middle), and inference latency (bottom) prompt lengths and batch size change.
Size
Model
Are C.
Are E.
Hella.
Wino.
SciQ
PIQA
Avg.
Llama
27.7
53.2
49.2
54.1
78.3
70.1
350M
Composite Stacked
27.9
53.9
49.6
54.8
78.7
71
56.0
Composite Stretched
30,5
55.4
50
54.4
79.6
70.3
56.7
Llama 3.2
30
56.8
52.9
55.9
81.7
70.7
58.0
Composite Stacked
33.5
59.5
56.7
59.4
83.2
73.1
60.9
Composite Stretched
32.5
61.5
56.1
56.1
88.1
72.5
61.1
Llama 3.2
30.1
58.6
54.3
56.4
83.1
72.8
59.2
3B
Composite Stacked
31.5
60.3
57.1
57.2
85.2
74.4
61.0
Composite Stretched
66.9
57.9
57.5
72.3
Llama 3.2
42.2
70.9
70.2
66.0
90.5
76.5
69.3
8B
Composite Stacked
42.7
71.5
70.6
64.8
91
77.4
69.7
Composite Stretched
43.6
71.8
70.4
66.5
91.3
78.5
70.4
Table 6: Evaluation of Composer's Stacked and Stretched hybrid architectures compared to Llama
3.2 on 6 downstream tasks.
architectures during small scale, it is prohibitive. Instead, results presented aggregated from 100 search trials. Figure 14 details the DCLM validation loss per hybrid LLM produced by each aggregation technique across
10-layer search.
Across all n-layer search, No clustering synthesizes the highest quality hybrid LLM. Increasing c reduces the number of top candidate architectures. N₂ clustering aggregates over synthesize hy- brid LLM, thereby reducing the generalizability of the produced LLMs to at-scale performance. No clustering also outperforms the p100 model, it smooths over any noise or overfitting phenomena that occurs during search. Hence, Composer leverages No clustering to produces its two Composite LLMs.
The following explanation provides further intuitive reasoning as to why No clustering works well. Consider each architecture as sequence blocks across layer indices. The search process re- turns set of top-performing architectures, which we filter based on validation accuracy with MAD. We interpret this collection as representative Monte Carlo sample from distribution over high- performing designs. This distribution is implicitly shaped by the architecture search process in our case, Bayesian optimization (BO). The BO surrogate model evaluates full architectural config- urations and captures interactions between blocks at different layers, learning correlations between structural choices and validation performance. As a result, the search does not sample randomly, but instead includes architectures that exhibit useful inter-layer dependencies.
Given this, we can define an empirical frequency for how often each block appears at given layer across the sampled architectures. This frequency estimates how likely that block is to appear that
26

---

Published as a conference paper at ICLR 2026
4-Layer Search
S-Layer Search
6-Layer Search
7-Layer Search
3.0
DCLM Val Loss
2.8
2.6
8-Layer Search
9 9-Layer Search
10-Layer Search
16-Layer Search
3.0
DCLM Val Loss
2.8
2,6
p100
No
N
p100
No
p100
Agg. Technique
Agg. Technique
Agg. Technique
Agg.
Technique
Figure 14: Comparison of Ne aggregation techniques across different n-layer searches.
layer in a high-performing design. Nn clustering selects, for each layer, the block that appears most frequently. The final architecture is then assembled by independently choosing the most common block at each layer across the top-performing candidates.
Selecting the block with the highest marginal frequency each layer equivalent computing the architecture that maximizes the product of these marginal frequencies across all layers. In other words, Ng clustering selects a configuration that is most probable under the assumption that each layer's 's block choice is independent of the others. This corresponds to a Naive Bayes-style estimator well-known technique in probabilistic modeling which maximizes the product of marginal likelihoods for individual variables while ignoring conditionals.
While this formulation ignores explicit inter-layer dependencies during aggregation, we argue this is appropriate because such dependencies are already accounted for during the search phase. That is, the BO search process acts as dependency-awarr sampler to construct the set of top-performing architectures. No clustering aggregates across structurally pre-filtered space, smoothing out noisy or overfitting samples while preserving dominant patterns.
Under this framing, No clustering is statistically consistent estimator of the most likely block sequence under a marginal model of high perforriff architectures. By prioritizing frequent blocks and avoiding conditional noise overfitting, Ng clustering yields architectures that generalize well when extrapolated to larger scales.
D.6 EXPLORATION OF DATASETS: BABISTORIES
In main paper, explored the efficacy of using different datasets (MAD (Poli al., 2024) and sampled-down version of DCLM (Li et al., 2024)) for evaluating candidate hybrid architecture with the HNAS Evaluator. We also include our findings when using BabiStories (Zhang al., 2025) to evaluate candidate architectures. As we leverage MAD (Poli et al., 2024) throughout our paper, include MAD's results here well for comparison.
Similar to our experiments with DCLM in $ 3.2, we explored two approaches for leveraging BabiS- tories during small-scale search. First, we randomly sample down the BabiStories dataset while keeping model size large (150M parameters). We perform a 16-layer search and extrapolate the discovered hybrid L.M to 1B parameters via stretching. Figure shows that this methodology, Large-Scale Stories, yields a relatively high quality model that outperform Llama 3.2 across train- ing budgets. However, the search cost is large (> 100 GPU-hours). Hence, as a second approach, we reduce the width of the model and keep model size between MM-MM parameters, Small-Scale Babis- tories. However, the stretched model at 1B scale consistently performs worse than Llama 3.2 across all training budgets. Ultimately. like DCLM, we found that using BabiStories either impractical or produced poor-performing models. Thus, Composer's 's HNAS Evaluator uses MAD throughout this work quickly but accurately evaluate the efficacy of candidate LIM architectures.
27

---

Published as a conference paper at ICLR 2026
Llama 3.2
MAD Stacking
Small-Scale BabiStories
MAD Stretching
Large-Scale BabiStories
3.4
Search Cost (GPU Hrs)
DCLM Val Loss
3.2
Cost
100
3.0
50
2.8
2.6
2
Dataset
Training Budget (FLOPs x 1e20)
Figure 15: Design exploration of datasets for HNAS Evaluator. We report the model quality at IB scale and search cost for each technique.
E COMPOSER DESIGN DETAILS
E.1 VISUALIZATION OF BAYESIAN OPTIMIZATION SEARCH JOURNEY
Layer Index
MLP
10
Attention
20
40
60
80
100
Search Trial Index
Figure 16:
Pictoral representation Composer' Search Engine converging on the computational primitives per layer index over 100 search trials for a 16-layer search.
Figure 16 shows Composer's Search Engine converging on computational primitives per layer index over 100 search trials for 16-layer search. discovers begin the hybrid architecture with atten- tion layers and end with MI.P layers. It then explores different interleavings of layers in between to find top performing hybrid architecture candidates.
E.2 DESCRIPTION OF DATASETS USED BY THE HNAS EVALUATOR
MAD Dataset (Poli al., 2024): The MAD dataset is set of synthetic tasks. The dataset consists of simple pretext token manipulations that provide quick and cheap performance estimates for given LLM arheitecture. MAD consists of six categories of tasks that LLMs should be capable of performing: in-context recall, fuzzy in-context recall, noise in-context recall, selective copying, compression, and memorization. For each candidate hybrid architecture in our search, we train the architecture across 800 samples task and then evaluate the architecture on 1280 samples to obtain cross-entropy loss for each task. We average the cross-entropy loss across six tasks as the final loss for the given architecture.
BabiStories (Zhang al., 2025): The BabiStories dataset synthetic benchmark children's stories, meant to train and evaluate the reasoning and generalization capabilities of language models in narrative contexts. The dataset is generated by LLMs using set rules that define entities, actions, and events, allowing for the systematic creation of diverse storylines. The simple gram- mar and smaller vocabulary size of the dataset allows small models of millions parameter scale to produce coherent English.
DCLM (Li et al., 2024): The DCLM dataset a large-scale, curated corpus designed for training and evaluating language models scale. It comprises over million documents spanning diverse
28

---

Published as a conference paper at ICLR 2026
domains, including news articles, scientific papers, and web content, each annotated with rich meta- data (e.g., publication date, source, topical categories). Unlike traditional sentence- or paragraph- level datasets, DCLM preserves the full document structure, enabling models to capture onngrragee dependencies and discourse-level phenomena. The dataset is preprocessed to ensure high textual quality, with non-linguistic artifacts and duplicates removed.
29


Reviews:
[{"review_id": "IvrhRpt4F1", "reviewer": "Reviewer_USsZ", "summary": "Overall, this work addresses important issues in NAS and demonstrates a fairly high level of research, providing extensive experimental results through numerous experiments. The writing is relatively clear, and the technical details are complete, though it lacks deeper theoretical analysis. \n\nIf the theoretical analysis and intuitive understanding of the experimental results and proposed methods could be strengthened, I would be willing to raise the score.", "questions": "1. I feel that the main focus of this work is not clearly highlighted. I suggest the author clarify what the central focus of this work is. Is it the discovered Composite architecture or the proposed Composer framework? \n\n2, Although four key questions are raised, they are not well answered. I recommend that the author provide deeper and more insightful conclusions for one or two of these questions. This might be better.\n\n3. The author claims, 'During search, the HNAS Evaluator trains and evaluates candidate hybrid LLMs with a small dataset to provide fast, reliable signals on the potential quality of the architecture at scale.' Please explain why a small dataset can provide reliable quality signals for large-scale models. How dependent is this signal on the dataset? I am very curious about this.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This article proposes a search framework for hybrid neural architecture design, conducts a variety of experiments, and achieves significant performance improvements.", "weaknesses": "1. The fourth key issue raised by the authors is important among the four issues. This paper introduces the approach of neural architecture extrapolators and empirically demonstrates its effectiveness, but the rationale behind this extrapolation method lacks analysis and justification. For how to extrapolate from smaller-scale networks to larger-scale networks, especially the possible explanation, it is suggested to refer to the following literature for a theoretical explanation.\n\nMathNAS: If Blocks Have a Role in Mathematical Architecture Design, Sihai Zhang et al, NeurIPS, 2023\n\n2. The first key issue raised in this paper is not that important, unless this small-scale model architecture search differs significantly from other small-scale architecture searches when scaling up, which is not discussed in this paper.\n\n3. The second key issue proposed in this paper is also difficult to explain. The paper only completes experimental comparisons and does not, and I believe cannot, provide insightful explanations as to why this dataset is used rather than that one.\n\n4. Perhaps the third key issue raised in this paper is the true contribution of the work. Using the idea of clustering to form the final network seems interesting. However, the related content included in the appendix does not discuss why clustering works. It is suggested that the authors provide an explanation for this.", "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": "ysudAit547", "reviewer": "Reviewer_ziXT", "summary": "This paper introduces Composer, a modular framework for automatically discovering hybrid large language model architectures that outperform standard Transformers. While standard transformers are composed by a uniform 1:1 ratio of alternating attention and mlp layers, the authors found that a 1:2 ratio is a/ more efficient to compute and b/ (marginally) better performing on multiple choice benchmarks", "questions": "please address above's weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- very relevant topic, other works also propose that the standard 1:1 ratio might be suboptimal\n- LLM's can be more efficient (computation) and potentially more performant.\n- solid framework and experimental design", "weaknesses": "W1 unfortunately only multiple choice downstream benchmarks were tested. you propose to alter 'the core' of current LLM's, it is pretty unclear whether this architecture change will also perform on disciplines as long context, coding, math, extended reasoning, or massive multi-linguality. at least for the final 3b models you should do one more ablation in these directions, baseline vs your 1:2 mix, to also validate this finding still holds.\n\nW2 isn't it correct that your 1b MLP has a different scale-up factors? shouldn't that skew search results?\n\nW3 will you publish your code before the conference?", "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": "blwZbmeWOE", "reviewer": "Reviewer_romG", "summary": "This paper introduces a framework for optimising the combination of different computational primitives (like attention or MLP blocks) in a neural network, towards hybrid architectures. It uses scaling strategies beyond scaling laws to search on small models that can later be scaled up. They focus on finding architectures for pretraining and ablate over choices for search, evaluation, aggregation and extrapolation in a comprehensive manner. They report consistent improvements in their found hybrid architectures on LLM benchmarks, including improved efficiency and downstream performance.", "questions": "- To quantify the robustness of the Composer framework for other types of pretraining data, it would be good to have an experiment that swaps DCLM for another corpus. Would the models found on different pretraining data also outperform Llama 3.2?\n- How many trials are performed of Bayesian Optimisation for the One-Shot 6-layer setup? The search space is 2^6 = 64 and it seems as 100 trials were performed for the 16-layer setup. Can a comprehensive evaluation of all architectures in the search space be done with an analysis how well the best can be found?\n- What task is the search accuracy evaluated on for Figure 1 (right)?\n- What are the details on the Bayesian Optimisation approach for One-Shot searches? What kernel/acquisition function/hyperparameters? There seems to be no details on this.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The experimental methodology is strong, with different components tested in isolation, and the claims of the paper appear well evidenced.\nThe paper is presented well overall, notably figure 1 and the 4 key design questions neatly structure the paper. This paper includes a strong ablation study across several dimensions of the proposed framework, which is presented in a clear structure with highlighted observations/key findings and further discussion.\nThe overall contribution of this paper is to provide a well designed framework for searching and scaling transformer architectures for LLM pretraining that can be useful for the broader ICLR community.", "weaknesses": "While the presentation of the method and experiments are good, the main paper is missing a Related Work section, which would significantly help in clarifying the surrounding literature and the contribution of this work, especially compared to the STAR method. Moving this in part from the appendix to the main paper would improve the clarity here.\n\nMore ablation could be done on the clustering method that is claimed to smooth out noise. Only Figure 6 (bottom) and the short discussion on it go into the effect of this clustering approach, which seems to be a core reason for the strong performances achieved. I also found the clustering explanation the hardest section to understand. Perhaps a visual aid could clarify it, e.g. a diagram of N_0 vs N_1 vs N_all or similar.\n\nIssues in text and tables:\n* L265-267: I don’t see how both One-Shot methods have lower search cost than End-Incremental. Figure 2 (left) suggests it is the same for 6-layer and more costly for 16-layer.\n* L320-321: What is meant by better capturing global information here?\n* Key results 1 and 2: Are these results computed across both the Composite - Stacked and Composite - Stretched models? A clarifiction is needed here\n* Table 6, final row reports an incorrect Avg. value for the Composite - Stretched model. It should be 62.0 instead of 66.9. I would point out to make sure the percentage claims in the abstract and elsewhere are not based on this mistake. \n* C.6 Baselines: If I understand correctly, the Striped attention should have its 1:4 ratio as 5x(1A+4M) but it says 5x(1A +2M). Same for 1:8.", "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": "ATFTf4MZHD", "reviewer": "Reviewer_naoJ", "summary": "This paper presents Composer, a systematic framework for automating the design of hybrid neural architectures that combine computational primitives (e.g., Attention, MLP) in novel ratios and interleavings. Key contributions include: A modular search framework with four core components: HNAS Engine (search algorithm), Evaluator (dataset selection), Aggregator (candidate synthesis), and Extrapolator (scaling to large models) .", "questions": "Please see the weaknesses.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.Composer addresses a critical gap in hybrid architecture design by automating the exploration of vast design spaces (e.g., 4 billion configurations for 32-layer models) . The integration of modular components (e.g., clustering-based aggregation  and synthetic dataset evaluation) ensures both efficiency and robustness.\n\n2.The proposed stacking and stretching techniques enable seamless scaling of small architectures to large models. For example, stretching 16-layer configurations preserves global interleaving patterns, outperforming stacking in larger models .", "weaknesses": "1. While MAD enables efficient search, its token-manipulation tasks may not fully represent real-world complexity. The paper lacks a comparison with other proxy datasets.\n2. Results are validated up to 3B parameters, but the framework’s applicability to larger models (e.g., 7B+) remains untested. \n3. The paper does not fully analyze the final searched structures, e.g., why specific interleavings (e.g., 1:2 Attention/MLP ratios) outperform others.", "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": "LOu0Bi86KK", "reviewer": "Reviewer_PWjK", "summary": "The paper addresses the limitations of traditional Transformer architectures in large language models (LLMs), which rely on a fixed 1:1 interleaving of self-attention and multi-layer perceptron (MLP) layers.\nThe authors note that hybrid architectures, using varied arrangements of computational primitives such as attention and MLPs, have demonstrated superior performance in prior manual designs.\nHowever, the vast design for a 32-layer model using only attention and MLP—coupled with high training costs, makes manual exploration intractable.\nComposer introduces an automated framework for hybrid neural architecture search (HNAS) tuned to pre-training from scratch, aiming to discover architectures that outperform baselines like Llama 3.2 at scales from 350M to 3B parameters.\n\nThe framework comprises four core components:\n* The Search Engine employs Bayesian Optimization with Gaussian Processes for one-shot searches or incremental methods (end-layer and middle-layer) to explore small-scale models with reduced widths.\n* The Evaluator assesses these candidates using small proxy datasets to provide rapid feedback predictive of large-scale performance.\n* The Aggregator generates top candidates through Nc clustering, selecting primitives layer-by-layer based on conditional frequencies among clustered architectures.\n* The Extrapolator scales up the models to the target size using stacking (repeating blocks) or stretching (proportional expansion of primitive groups).\n\nExperiments yield hybrids surpassing Llama 3.2. The authors remark the benefits of 1:2 attention-to-MLP ratios and non-standard interleavings over the standard 1:1 Transformer structure.\n\nProxy datasets like MAD (and other synthetic datasets like BabiStories) enable efficient evaluation, outperforming downscaled DCLM in cost.\n\nDuring aggregation, Stacking is robust across search depths, but stretching excels for larger small-scale searches (e.g., 16 layers).\n\nThe resulting \"Composite\" architectures, consisting of 1:2 interleavings of Grouped Query Attention and SwiGLU layers, reduce validation loss by 0.05-0.08 and improve downstream task accuracy by 1.1-3.1% on average, while improving training throughput by 1.25x.", "questions": "Will the code be publicly available?\n\nWhat is the performance of Composer-discovered hybrids on long-context tasks like needle-in-a-haystack or multi-document QA? These tasks have been shown to benefit from attention architecture improvements.\n\nHas the framework been applied to other primitives, such as state-space models (e.g., Mamba) with attention?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper introduces a principled, module framework for NAS. The framework's four components (Search Engine, Evaluator, Aggregator, and Extrapolator) provide a blueprint that is extensible to other primitives.\n\nThrough experiments, the author find \"creative\" 1:2 attention-to-MLP ratios and non-standard interleavings that represent useful alternative architectures for future exploration.", "weaknesses": "The novelty of the method is minimal, as the Composer framework consists of common building blocks.\n\nOne critical weakness of the paper is that the proposed method proposes no theoretical or empirical justification of optimality. Through experiments, the authors are able to find hybrid architectures that improve upon the LLaMA3.2 baseline, however in the absence of theoretical grounding, we would like to see the Oracle of these experiments, i.e. a large scale experiment to train all candidates to completion, in order to validate the use of proxy evaluations.\n\nSimilarly, there is no evidence that the method scales to models that are larger than 3B parameters, how do the authors justify the use of the Composer framework for these?\n\nThe paper is also limited in scope, since it only mentions in the end, but does not test the method on state-space models, or less conventional attention mechanisms.\n\nEmpirical results, while positive (e.g., 0.05-0.08 validation loss reductions and 1.1-3.1% average downstream accuracy gains), show modest and context-specific advantages. Some variants, like those from downsampled DCLM or BabiStories, only marginally outperform or even underperform Llama 3.2 at higher training budgets.", "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)