STAR: Synthesis of Tailored Architectures

Thank you for your thoughtful review and for highlighting the strengths of our work. We appreciate your recognition of the novel aspects of our approach, including the use of a hierarchical search space based on the theory of linear input-varying systems and the simultaneous optimization of multiple metrics. We're also glad that you found our experimental analysis to be well-analyzed. Below, we address your comments and provide further clarification:

Weakness 1 / Question 1: Comparison to pre-trained LLMs such as those of the Llama series.

Thank you for this insightful suggestion. STAR is designed to develop deep learning architectures tailored to specific objectives. To evaluate candidate architectures, and to provide a fair comparison, all architectures–baselines and STAR synthesized–are parameter-matched and trained from scratch on the same dataset. We emphasize that the baseline Transformer++ architecture in our paper is identical to the architecture used by the Llama model series. The StripedMamba baseline is based on Jamba (Lieber et al., 2024), another state-of-the-art pre-trained LLM.

We have further extended our analyses to include comparisons at the 1B-parameter scale, showing that a STAR-synthesized backbone optimized for quality and cache size outperforms a parameter-matched Transformer++ baseline and matches a StripedMamba baseline in quality when trained for 40B tokens on the same data—while reducing cache size by 90% and 37%, respectively. For this analysis, we scaled a 125M-parameter STAR backbone by increasing depth from 24 to 48 operators and width from 768 to 2048 dimensions. This finding suggests that the advantages of STAR synthesized backbones observed at smaller scales translate to larger model scales akin to those of the referenced pre-trained LLMs.

Question 2: Comparison to other KV-cache optimization methods

Thank you for this valuable question. Our results demonstrate that STAR can design cache-optimized architectures that match Transformer++ and StripedMamba in quality while offering significant cache size reductions. This is achieved solely by optimizing the choice of operators and their compositions, such as feature group sharing (e.g., KV-cache sharing) and shared featurizer weights. Importantly, existing cache optimization methods like quantization can be applied to STAR-synthesized architectures just as they are to other architectures like Transformer++. These methods thereby provide additional tools to further reduce cache sizes for STAR models and are complementary rather than competitive approaches.

Reference:

Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., ... & Shoham, Y. (2024). Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887.

Thank you for your thoughtful review and for highlighting the strengths of our work. We appreciate your recognition of the empirical strengths, versatility, and efficiency of our hierarchical design space, and that you find our paper well-written and clear. We're also glad that you found our experimental analysis to provide convincing evidence in favor of the STAR’s ability to effectively synthesize high-quality architectures tailored to various objectives. Below, we address your comments and provide further clarification:

Weakness 1: Synthesizing improved residual streams.

Thank you for this valuable suggestion. We have extended the backbone genome to support residual streams beyond the conventional pre-norm structure. Appendix A.7 details this extension and includes an ablation study comparing STAR with and without it. The results show that this extension enables STAR to synthesize architectures with slightly smaller parameter counts while maintaining the same quality. Based on these promising findings, we plan to further explore residual composition strategies in future work.

Weakness 2: Does the mutation process require expert intuition?

Thank you for this insightful comment. We have not observed any synthesized architectures exhibiting unstable training behavior. While STAR corrects invalid genome edits during evolution, these corrections are fully determined by the STAR genome itself and do not require expert knowledge, since every position of the genome has a specific function that does not change no matter to which objectives STAR is applied. To clarify this, we have added a complete description of the STAR genome and all its possible values in Appendix A.6. We have also enhanced our explanation of how invalid genome edits are repaired in the main text (Section 4.2).

Weakness 3: Underbrace in the attention operator math.

Thank you for the helpful suggestion. We have incorporated all of the reviewer's recommendations into our notation and have updated the math notation throughout to improve overall clarity for the reader.

Question 1: What are the differences between STAR 1–8, and how were they selected?

Thank you for these valuable questions. We have now included visualizations of all evaluated STAR backbones—eight per STAR evolution—in Figures B.1-25 and expanded our discussion of recurring architecture motifs (Appendix B.1). We selected these backbones from the pool of synthesized candidates as follows: after the evolution process, we chose the eight backbones that achieved the lowest perplexity among those with lower parameter counts (for the quality and quality-size optimizations) or smaller cache sizes (for the quality-cache optimization) compared to our baseline models. We now explicitly state this selection criterion in the main text (Section 5) to clarify it for readers.

Question 2: Why do you synthesize at a fixed depth of 24 LIVs?

We selected a depth of 24 LIVs at a width of 768 dimensions to match the commonly used "125M-parameter" model scale in language modeling experiments. Typically, each of the 12 layers in a 125M architecture comprises a sequence-mixing operator (e.g., attention) and a channel-mixing operator (e.g., SwiGLU), totaling 24 operators. While STAR can optimize variable-depth and variable-width architectures, this introduces challenges: shallower genomes are computationally cheaper and converge faster but may lack necessary complexity, whereas deeper genomes offer greater representational power but expand the search space, risking suboptimal convergence due to overfitting or inefficient sampling. In line with the reviewer's suggestion, we have added a new section discussing this topic (Appendix A.1).

Thank you for your thoughtful review and for highlighting the strengths of our work. We appreciate your recognition of the significance of the research problem we are tackling, that you view our hierarchical design space as clear and theoretically well-grounded, and find our choice of evolutionary optimization to search this space as well-justified. We're also glad that you found our experimental analysis thorough and the results promising. Below, we address your comments and provide further clarification:

Weakness 1 / Questions 1–2: Clarification on search space design choices and their relation to existing implementations.

Thank you for these valuable suggestions. Our revised manuscript clearly distinguishes between token- and channel-mixing structure (Section 2). Additionally, we now provide a full description of the STAR genome—and thus the hierarchical design space—in Appendix A.6, outlining the channel- and token-mixing structures of each LIV and how each LIV class relates to existing implementations. We focused on this set of operators because they encompass all empirically relevant operator types currently used in state-of-the-art language modeling. We have further extended the design space to include novel differential variants of all applicable LIV classes, similar to the "Differential Transformer" (Ye et al., 2024) (see Section 5) .

Weakness 2 / Question 3: Clarification of the meaning and parameter ranges of each STAR genome entry.

Thank you for this insightful comment. Our revised manuscript now includes a full description of each entry in the STAR genome, the possible values they can take in our analyses, and how these translate to specific architectures (Appendix A.6). We believe this greatly enhances the reproducibility and clarity of our work.

Weakness 3 / Question 4: Clarification on why the evolutionary algorithms are considered gradient-free.

We thank the reviewer for raising this insightful question. In STAR, models are optimized for quality by evaluating and comparing their performance after being trained on a predefined dataset shared across all candidate architectures. The evolutionary optimization algorithm, which is used to optimize the genomes representing the trained architectures, relies solely on the performance score of the trained model, along with other metrics such as parameter count and cache size. Importantly, it does not access the gradients used during model training. For this reason, we refer to the evolutionary algorithms employed in STAR as gradient-free—a term commonly used in the evolutionary optimization literature. To clarify this for readers, we have updated Figures 1.1 and 4.1 to include detailed schematics of the optimization process, clearly illustrating the roles of model training and evaluation within the evolutionary framework.

Weakness 4 / Questions 5–6: Details on recombination, mutation operations, and repair of invalid genome entries.

Thank you for pointing out the need for clarity in this area. We have revised the manuscript to specify that we employ uniform random sampling from the set of valid options for mutations. The valid options for each genome entry are detailed in Appendix A.6. We use the k-point crossover method for recombination, outlined in detail in Appendix A.4. As you correctly noted, these operations can result in invalid genome entries. Our revised manuscript now clearly explains how we repair any invalid entries that may occur in the backbone genome (Section 4.2). Specifically, if mutations or recombinations result in invalid configurations—such as incompatible sharing strategies—we detect and repair them by removing the invalid connection for entries 2 and 4 of the backbone genome or by re-sampling valid values for entries 1, 3, and 5. Importantly, these repair operations are infrequent during STAR optimization, occurring less than 20% of the time.

Weakness 5 / Question 7: Benefits and drawbacks of holistic versus multi-level genome optimization.

Thank you for this excellent question. We have expanded our discussion section (Appendix A.1) to address the advantages and drawbacks of treating the genome as a unified entity versus using a multi-stage optimization strategy targeting each hierarchical level. The holistic approach benefits rapid iteration and resource-limited scenarios but may not fully exploit the hierarchical design space, potentially overlooking dependencies across genome levels. In contrast, a multi-stage optimization strategy could improve convergence by leveraging the genome's modularity and addressing interactions incrementally, though it increases algorithmic complexity and computational costs. Notably, even with our current holistic approach, 90–100% of the synthesized architectures outperform baseline models on the respective metrics (quality, quality and size, quality and cache size), thus meeting our objectives.

Weakness 6 / Questions 9–10: Clarification on the choice of STAR baseline models.

Thank you for this insightful suggestion. STAR is designed to develop deep learning architectures tailored to specific objectives. To evaluate candidate architectures, we train them from scratch and compare them to parameter-matched, state-of-the-art baseline architectures trained on the same data. We emphasize that the baseline Transformer++ and StripedMamba architectures in our paper are identical to/based on those used in state-of-the-art pre-trained models like Llama, which uses a Transformer++ architecture, and Jamba (Lieber et al., 2024), which is based on a StripedMamba architecture. We did not include a pure Mamba baseline architecture, as employed by Falcon-Mamba, as it has been empirically shown to be inferior in quality to a respective striped variant incorporating a small set of attention operators (e.g., Poli et al., 2024). As suggested by the reviewer, we train parameter-matched variants of these architectures on the exact same data as STAR-synthesized backbones to ensure a fair comparison.

We have further extended our analyses to include comparisons at the 1B-parameter scale, showing that a STAR-synthesized backbone optimized for quality and cache size outperforms a parameter-matched Transformer++ baseline and matches a StripedMamba baseline in quality when trained for 40B tokens on the same data—while reducing cache size by 90% and 37%, respectively. For this analysis, we scaled a 125M-parameter STAR backbone by increasing depth from 24 to 48 operators and width from 768 to 2048 dimensions. This demonstrates that the advantages observed at smaller scales translate to larger models akin to those referenced.

Weakness 7 / Question 8: Improving the reproducibility of STAR.

Thank you for highlighting this important aspect. Our revised manuscript now provides a full description of all entries in the STAR genome, the possible values they can take, and how these relate to existing algorithm implementations (Appendix A.6). We also clearly outline each step of the evolutionary optimization process, detailing how genomes are mutated, recombined, and how we repair any invalid entries (Section 4 and Appendix A.4). We believe these updates significantly enhance the reproducibility of our work.

References:

• Ye, T., Dong, L., Xia, Y., Sun, Y., Zhu, Y., Huang, G., & Wei, F. (2024). Differential transformer. arXiv preprint arXiv:2410.05258.
• Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., ... & Shoham, Y. (2024). Jamba: A hybrid transformer-mamba language model. arXiv preprint arXiv:2403.19887.
• Poli, M., Thomas, A. W., Nguyen, E., Ponnusamy, P., Deiseroth, B., Kersting, K., ... & Massaroli, S. (2024). Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844

I appreciate the authors' comprehensive and well-articulated responses, which have successfully addressed several of my concerns. I strongly encourage the authors to integrate these clarifications into the revised manuscript to further strengthen its presentation. Sharing the codebase would also be a valuable addition, promoting transparency and enabling reproducibility. I strongly recommend the authors to open-source their code in the revised version of the paper. Also, there are missing graphs in Figure B.29, the authors should update this figure by including all of the graphs. Overall, Considering the significance of the paper's contributions and findings, specifically in neural networks design automation I am happy to revise my score and recommend the paper for acceptance.

Thank you for your valuable review and feedback. We are pleased that you found the presented results promising and that you appreciated the novelty of our architecture genome encoding. Below, we address your comments and provide further clarification:

Weakness 1: Evolutionary algorithms for architecture design are not new.

We thank the reviewer for this insightful comment. We acknowledge that evolutionary algorithms (EAs) are well-established in architecture optimization. However, STAR distinguishes itself through a novel hierarchical search space for computational units and their composition, along with a unique numerical encoding–the STAR genome–tailored for evolutionary methods. The STAR design space, grounded in the theory of linear input-varying systems (LIVs), provides a new framework for creating architecture components. LIVs generalize key computational units in deep learning, including attention, recurrences, convolutions, and other structured operators. The framework characterizes architectures at three hierarchical levels: (a) featurization, modulating linear computation based on input context; (b) operator structure, describing the structure for token and channel mixing; and (c) backbone, detailing the interplay of LIVs. Unlike prior search spaces, the LIV search space is both comprehensive and stable, with most sampled candidates training successfully.

Weakness 2: Unclear which levels are affected by mutation.

We thank the reviewer for pointing out that it is unclear whether the mutation operations performed by STAR in fact affect the underlying LIV components. We would like to clarify that these mutations affect the specific characteristics of the LIVs and the way in which individual LIVs contained in a backbone are connected, for example, determining if they share specific feature groups (akin to KV-cache sharing), or weights. The mutation operation also alters the specific characteristics of the LIV classes, such as switching between softmax attention variants with varying KV-cache compression ratios (similar to multi-query attention). By incorporating these mutation strategies and focusing on optimizing the backbone genome, STAR is able to synthesize architectures that outperform strong Transformer++ and StripedMamba baselines on multiple objectives and downstream language modeling benchmarks.