Pretrained Hybrids with MAD Skills

Thank you for noting that Manticore is a good alternative to existing methods!

On downstream task metrics. We appreciate this feedback. Table 3 in our draft reports accuracy scores for Long Range Arena (LRA) tasks. For our language modeling evaluations, we follow standard practice of reporting loss/perplexity as these directly measure model quality for generation tasks. We would be happy to add perplexity values (exp(loss)) alongside loss values if this helps interpretability. Our focus is on comparing Manticore to its component models, which our results demonstrate effectively.

On the significance of scores. Indeed, while we do not have multiple seeds for all experiments due to compute constraints, our experimental design uses architecture parameter sweeps (Figures 2, 3) that show performance surfaces rather than point estimates. These sweeps demonstrate consistent improvements over component models across different mixture weights, providing evidence for robustness.

On the selection of component models. We initially selected the GPT-Neo and Pythia family of models to represent transformers because they are fully open source and provide many different sizes to choose from. We chose Mamba because it is a popular alternative to the transformer family of architectures. However, our framework is not limited to these! We outline in Section 2 that Manticore is general enough to support any language model architecture family that is organized in blocks.

On the number of component models and blocks. Our analysis is not limited to two component models and two Manticore blocks. In fact, in 3.2, we actually used three component models, but for our other experiments, we chose to use two component models for simplicity and because most hybrids only combine components from two different architectures. For our LRA experiments, we actually used up to 6 Manticore blocks–we will clarify this in the final version. For many of our other experiments, we chose to use up to two Manticore blocks for simplicity and because they often performed well. We will include a discussion of these choices in the final version.

On presentation. Thank you for pointing out these issues. We have fixed the typo, cosmetic issues, and missing arxiv links. We are also happy to clarify any specific statements you found unclear; your feedback is appreciated.

On calling the shared representations a common language. Our use of the term is through an analogy. We are not referring to formal languages, and we will clarify.

On the models used in Section 3.2. For the experiments on the MAD tasks, we use default architecture settings from the MAD codebase [1], which we train from scratch. We will clarify this in the final version.

On the green dotted lines in Figure 2. The green dotted lines represent the loss values throughout training. Since this is not central to our analysis, we will remove these for clarity in the final version. Thank you for pointing this out!

On support for the statement in lines 288-289. The following line, 290, notes that this is provided in the Appendix. Specifically, this is shown in Table 5 of Appendix B. There, we show that post-search discretization and retraining yields similar results, but since the Manticore hybrids are discretized (meaning that the convex combinations are rounded to the nearest one-hot vector), the model size and inference cost can be decreased by roughly half.

[1] https://github.com/athms/mad-lab

Thank you for noting that posing hybrid design as a search problem is interesting and meaningful!

On metrics and the breadth of tasks. Table 3 already reports accuracy rather than loss. Furthermore, Table 3 specifically shows accuracies on Long Rage Arena (LRA) tasks, which span different modalities (ListOps for hierarchical reasoning, IMDb for text classification, CIFAR10 for vision, Pathfinder for long-range dependencies). The MAD tasks (Table 2) specifically test basic reasoning capabilities like associative recall and selective copying. Together, these cover reasoning and long-sequence scenarios.

On analyzing specific architectures. Our primary architectural insight is that for the majority of our experiments, Manticore produces architectures that mix its component models, rather than defaulting to the component models themselves. Crucially, there are cases in which MAD can be used as a proxy task to produce architectures that do well on tasks that are similar to the MAD tasks. Figures 2 and 3 both depict the resulting architectures from our experiments in relation to component models. In particular, we discuss in lines 317-336 how the architecture derived using MAD relates to both of its component models—GPT-Neo and Mamba—and two naive frankenmerged hybrids comprising these two models. The optimal model in this case, is one that partly still uses at least one of the GPT-Neo blocks, rather than defaulting entirely to Mamba. This suggests that successful hybrids leverage complementary capabilities rather than simply selecting the better architecture for each block.

On analyzing the computational and parametric efficiency of Manticore. We analyze the parametric efficiency as well as the computational costs associated with pretraining and inference in Appendix E.

On model families and larger models. Our experiments demonstrate generality across diverse families: Transformers (GPT-Neo, Pythia), State Space Models (Mamba), and alternative architectures (Striped Hyena). We test sizes from 125M to 410M parameters, which is a parameter range that is particularly relevant as it represents the scale where many novel architectures (SSMs, linear attention variants) are first validated before scaling up. Our framework (Section 2.1) supports any architecture following the standard LM recipe. We will clarify in the final version that computational costs scale linearly with component model size, making larger models feasible.

Thank you for noting the advantages of Manticore, including the reuse of existing pretrained models/avoiding costly pretraining, the tractability of NAS usage, and our experiments with programming pretrained hybrids!

On the reliance on gradient-based NAS. It is true that in the NAS literature, where the objective is often to derive a discrete architecture from a learned convex combination of components, DARTS can often fail to produce a performant architecture. However, our search space is continuous, which makes it particularly well-suited to gradient-based NAS techniques. While DARTS or other search methods may not always find global optima, Figure 3 shows it consistently finds regions that achieve our goal of outperforming individual component models, and Table 4 shows DARTS outperforms DASH in our setting.

On inference cost. This is an explicit tradeoff: doubled inference for avoiding expensive pretraining of new hybrids. All other considerations equal (e.g. total training cost and parameter count), a reasonable baseline would be an ensemble, and Manticore outperforms ensembles. We will make this point more clear in the final version. Additionally, NAS is well-suited to problems of trading off considerations like inference cost or model size with performance – exploring NAS algorithms that navigate this tradeoff for Manticore is an interesting avenue that we shall also consider.

On inference, knowledge distillation, and sparsification. This is a great point! We measured the inference throughput for the IMDb task in LRA, and obtained the following results:

• Mamba: 283.508 samples/s
• GPT-Neo: 87.719 samples/s
• Manticore: 68.853 samples/s.

This demonstrates that in practice, while the inference cost of Manticore is naturally going to be higher than its component models, the overhead remains quite manageable, especially when combining transformers with subquadratic architectures such as Mamba. Note that we also include an analysis of the inference FLOPs in Appendix E.3. We agree that techniques such as knowledge distillation might be useful for reducing this cost further, and adapting techniques such as MOHAWK [1] to our setting might be a useful next step in this research direction. We perform an ablation over post-search discretization and retraining (which is a form of sparsification; this can be found in Appendix B), and we find that it can help in certain cases. Exploring this direction further to improve the inference cost of Manticore is an intriguing direction for future work.

On projector architecture. We did not find that deeper projectors with nonlinearities led to improved performance, and they came at the cost of more expensive projector pretraining. For this reason, we kept the projector design simple, while still including a gating mechanism which would allow for the component architectures to still be part of the search space. We will include a discussion of this in the final version.

[1] https://arxiv.org/abs/2408.10189

Thank you for praising the importance of our problem, the significance of the steps that Manticore takes toward developing efficient LMs, the novelty of our approach, efficiency, and of our results supporting our claims! Your other points are aligned with clarifications that we are already in the progress of integrating into the manuscript, and we respond to them below.

On the types of tasks that Manticore helps with. This is a great point that we would love to make more explicit in the final version! Several of our experiments in Figure 3 were on datasets that combine data from multiple sources (e.g. NI Spanish QA and Alpaca). Our empirical results reflect some of the known early set of theoretical findings on hybrid models. Specifically, it is known that hybrids excel at compositional tasks like finding a token arbitrarily far in the past and then performing a local copy operation – this for instance, necessitates a tradeoff in SSM state size and transformer context. Results like these motivate the study of tools like Manticore. For this reason, we suspect that Manticore excels at tasks in which the component models are specialized for certain data sources or aspects of the dataset, which we briefly mention on line 251. We will discuss this further in the final version.

On the cost of inference FLOPs. We discuss our FLOPs analysis and additional experimental results related to this in Appendix E. All other considerations equal (e.g. total training cost and parameter count), a reasonable baseline would be an ensemble, and Manticore outperforms ensembles. Although latency-sensitive applications are not the focus of our work, we note that for such applications where the inference cost of Manticore (or ensembles) is prohibitive, it might still be beneficial to train a hybrid from scratch.

On DARTS as the default search algorithm. We include experimental results that compare DARTS to other NAS algorithms – this is shown in Figure 3. We found that between DARTS, DASH, and GAEA, there was no clear winner across all of the tasks, but that for all but one task, at least two of them were effective. In short, we agree with your point (i) that we should discuss the limitations of DARTS for Manticore, but we note that our framework is flexible and supports other NAS algorithms, while leaving room for specialized algorithms to be developed for pretrained hybrid search spaces. We will highlight this in the final version. For (ii), we also include an ablation comparing DARTS to the DASH search algorithm, both with and without alternating architecture parameter/shared weights updates in Appendix B, along with whether or not to perform post-search discretization and retraining. These ablations help characterize the desirable traits of NAS search algorithms for Manticore.

On the choice of projector architecture. This is a great point! We wanted to keep the projector design as simple as possible to reduce the cost of projector pretraining, while including a gating mechanism that would allow for the component architectures to still be part of the search space. Conceivably, one could use multi-layer projectors with nonlinearities, but we found that this introduced higher cost and complexity without necessarily improving performance. We will include a discussion of this in the final version.