# Pretrained Hybrids with MAD Skills

Nicholas Roberts111Corresponding author<EMAIL_ADDRESS>UniversityofWisconsin-Madison  Samuel Guo  UniversityofWisconsin-Madison
Zhiqi Gao  UniversityofWisconsin-Madison  Satya Sai Srinath Namburi GNVV
UniversityofWisconsin-Madison  Sonia Cromp  UniversityofWisconsin-Madison
Chengjun Wu  UniversityofWisconsin-Madison  Chengyu Duan
UniversityofWisconsin-Madison  Frederic Sala  UniversityofWisconsin-Madison

###### Abstract

While Transformers underpin modern large language models (LMs), there is a
growing list of alternative architectures with new capabilities, promises, and
tradeoffs. This makes choosing the right LM architecture challenging.
Recently-proposed _hybrid architectures_ seek a best-of-all-worlds approach
that reaps the benefits of all architectures. Hybrid design is difficult for
two reasons: it requires manual expert-driven search, and new hybrids must be
trained from scratch. We propose Manticore,222The Manticore is a fearsome
human/lion/scorpion hybrid from Persian mythology. a framework that addresses
these challenges. Manticore automates the design of hybrid architectures while
reusing pretrained models to create pretrained hybrids. Our approach augments
ideas from differentiable Neural Architecture Search (NAS) by incorporating
simple projectors that translate features between pretrained blocks from
different architectures. We then fine-tune hybrids that combine pretrained
models from different architecture families—such as the GPT series and
Mamba—end-to-end. With Manticore, we enable LM selection without training
multiple models, the construction of pretrained hybrids from existing
pretrained models, and the ability to _program_ pretrained hybrids to have
certain capabilities. Manticore hybrids outperform existing manually-designed
hybrids, achieve strong performance on Long Range Arena (LRA) tasks, and can
improve on pretrained transformers and state space models.

##  1 Introduction

Transformers are the workhorse architecture for large language models and
beyond, powering a vast collection of foundation models. While for years it
appeared that the Transformers family would remain the undisputed standard, a
recent _Cambrian explosion_ of proposed architectures has taken place. Many of
the new architectures achieve subquadratic complexity—in contrast to the
quadratic complexity of self-attention in Transformers—by using local or
linear attention [9, 6, 3, 45], or resurrecting and scaling recurrent networks
[6, 9, 31], or by building on state-space modeling principles [13, 33, 32, 11,
14]. These approaches potentially promise to overturn the dominance of
Transformers through more efficient training and inference.

However, no single new model is a clear overall winner when varying data
modalities, tasks, and model sizes. Comparing architectures on a fixed task is
fraught with difficulties [2]. Even if these issues are overcome,
practitioners would have to experiment with and evaluate every architecture
for each new task—an expensive proposition. Instead, seeking a best-of-all-
worlds approach, researchers have proposed the use of _hybrid models_ that mix
multiple architectures. These hybrids, such as the MambaFormer [29]—a mix of
the popular Mamba with a Transformer—have shown potential in maintaining the
desirable properties of multiple model classes.

While promising, hybrids suffer from two main obstacles that stymie their
adoption:

  * •

Manual Design. Hybrid architectures are hand-crafted, either by manually
exploring the large search space of hybrids or by relying on often unreliable
intuition and heuristics.

  * •

Failure to Use Pretrained Models. It is unclear how to integrate _pretrained_
model components from models with different architectures. Pretrained models
are a key advantage of foundation models. However, due to compatibility
issues, hybrids are often trained from scratch, leading practitioners to
resort to small hybrids in limited settings or incur high costs.

A potential solution to the latter challenge is the use of _model merging_
techniques [42, 44, 41, 15, 8, 16], some of which can operate cross-
architecture [1, 12]. Unfortunately, such tools are embryonic–they are
expensive and it is unclear how well they work with the diverse types of
architectures a user may seek to build a hybrid from.

We propose a framework for _automatically designing hybrid architectures_ that
overcomes these obstacles. Our approach is inspired by principles from neural
architecture search (NAS), but applies these at the level of _LM blocks_
rather than convolutional cells [23, 21] or operations [36, 35]. The resulting
framework is simple and tractable. It sidesteps merging different
architectures by using simple linear projectors to translate between the
“languages” spoken by various architectures. This enables us to include blocks
from many different architectures with little to no changes required. In
addition, inspired by the mechanistic architecture design (MAD) framework
[34], we show how our framework can learn hybrid architectures via MAD that
transfer to new tasks.

Concretely, our proposed system, Manticore:

  1. 1.

Automatically selects language models, without training several models from
scratch,

  2. 2.

Automatically constructs pretrained hybrids without evaluating the entire
search space,

  3. 3.

Provides a technique for programming hybrids to have certain skills without
full training.

Experimentally, our automatically designed hybrids compete with existing
hybrids and models on the MAD tasks [34] and Long Range Arena (LRA) [38], we
produce pretrained hybrids that can improve downstream fine-tuning
performance, and we can program hybrids using the MAD tasks.

##  2 Related work

Language Model Architectures: Transfomers and Beyond. Transformers are
currently the dominant LM architecture. The success of the “vanilla”
architecture introduced by Vaswani et. al. [40] has led to many proposed
variations. The quadratic complexity of the base self-attention operation has
inspired the search for alternative architectures that offer comparable
performance with subquadratic complexity. One line of work builds off _state-
space models_ , with variations made to enable language modeling [32, 33, 13,
3]. Another line of work involves linear attention by formulating transformers
as RNNs and expressing self-attention as a kernel dot-product [17]. Other new
approaches increase the expressivity of this formulation with data-dependent
gating [43].

Our work does not propose a new architecture. Instead, we focus on the idea
_that practitioners should be able to take advantage of new architectures in a
transparent way_.

Neural Architecture Search & Mechanistic Search. Neural architecture search
(NAS) techniques are used to automatically search for optimal architectures.
These techniques have produced state-of-the-art models in several domains.
Much of the challenge in NAS is the complexity of search; in the most standard
form, NAS involves a difficult bilevel optimization over a large search space.
Much effort has been aimed at reducing costs, often via continuous relaxations
of search spaces, with techniques like DARTS [23] and DASH [36].

Using NAS to discover architectures for language modeling—and especially those
that may rival Transformers—has thus far been hard. A promising approach is
the MAD framework [34] , which uses “ _mechanistic_ tasks” (synthetic tasks
organized around simple principles) to search for high-quality architectures.
While we do not seek to discover new architectures, we are inspired by this
approach in our effort to search for hybrid architectures.

Hybrid Architectures. Perhaps unsurprisingly, there is no single dominant
architecture among either standards, like Transformers, or emerging
subquadratic architectures. While there are some insights that can be
converted into heuristics for model selection, generally, to take advantage of
new models, practitioners must exhaustively evaluate all of them on each of
their tasks. The cost of doing so has inspired the idea of crafting hybrid
architectures that mix components from different approaches, with the goal
being to obtain best-of-all-worlds behavior.

Unfortunately, the space of hybrid architectures is already large and only
grows with each new proposed approach. Manually crafting hybrids is costly;
users must either brute-force the enormous search space or alternatively hand-
craft a small candidate set of hybrids in the hope that it includes a
reasonably performant choice. Our work provides an efficient alternative to
this process.

Model Merging. A final prospective approach to using multiple models is
_merging_. Merging pretrained models (of the same architecture) has shown
promising results [42, 44, 41, 15, 8, 16], creating powerful large-scale
merges such as SOLAR-10.7B [18] and
Goliath-120B333https://huggingface.co/alpindale/goliath-120b from two fine-
tuned Llama2-70B [39] models. The former two were produced using a trial-and-
error-based technique called ‘frankenmerging,’ introduced in MergeKit [12].
Frankenmerging involves stitching together different fine-tuned versions of
the same model or, hypothetically, different models. This has inspired efforts
to merge models of different architectures using large-scale evolutionary
search [1]. However, such efforts are still embryonic, with substantial
computational drawbacks, requiring many training runs. Manticore, on the other
hand, does not require training a large number of models.

##  3 Methods

We now describe Manticore, our framework for automatically designing hybrid
architectures by mixing components of pretrained models. Manticore relies on
projectors to align features across architectures, then applies a convex
combination to the aligned features, as summarized in Figure 1. In Section
3.1, we discuss and formally define the structure of Manticore hybrids,
including the projectors and convex combination mixture weights, as well as
how both of these components are used within Manticore. In Section 3.2, we
detail the search procedures (inspired by NAS) and training routines involved
in pretraining, fine-tuning, and programming hybrids.

Figure 1: Our proposed Manticore framework, which enables: (1) cross-
architecture LM selection, (2) the construction of pretrained hybrids, and (3)
the ability to program hybrids to have certain skills.

###  3.1 The Structure of Manticore Hybrids

Our framework comprises three main parts: the individual LMs that we combine
to produce our overall hybrid, projectors that translate feature
representations between LMs of different architectures, and convex combination
mixture weights that specify how much the hybrid will use the features of each
component architecture. We detail each of these in the following.

####  3.1.1 Component Models

We refer to a model that is used in Manticore as a component model. Any modern
LM can be used as a component model in our framework. In this section, we will
formally define the general high-level structure of the component models that
we support. For an LM $M$ with model embedding dimension $d_{M}$ on a sequence
of $t$ tokens from a set $\mathcal{V}$, denoted
$x=(x_{1},...,x_{t})\in\mathcal{V}^{t}$, a forward pass $M(x)$ is typically
computed using the following recipe:

  1. 1.

Apply an embedding function,
$M_{\text{embed}}:\mathcal{V}^{t}\to\mathbb{R}^{t\times d_{M}}$ to the tokens,
resulting in a sequence of embeddings denoted
$x_{\text{embed}}=M_{\text{embed}}(x)$.

  2. 2.

Take forward passes through $L_{M}$ ‘blocks’–we denote the $\ell^{\text{th}}$
block as $M_{\text{Block}}^{(\ell)}:\mathbb{R}^{t\times
d_{M}}\to\mathbb{R}^{t\times d_{M}}$. Specifically, for all $\ell\in[L_{M}]$,
we obtain $x_{\ell+1}=M_{\text{Block}}^{(\ell)}(x_{\ell})$, where
$x_{1}:=x_{\text{embed}}$.

  3. 3.

Finally, we pass $x_{L_{M}+1}$ into a language modeling head,
$M_{\text{head}}:\mathbb{R}^{t\times d_{M}}\to(\Delta^{|\mathcal{V}|-1})^{t}$,
where $\Delta^{|\mathcal{V}|-1}$ is the probability simplex of dimension
$|\mathcal{V}|$.

This recipe applies to virtually all modern transformer-based LMs, recurrent
models, and state-space models. Our framework supports all of these, and any
other architecture that follows this recipe.

####  3.1.2 Projectors

Suppose we have two pretrained component models, $M$ and $M^{\prime}$. In
general, even if the model dimensions are the same for both models
($d_{M}=d_{M^{\prime}}$), blocks from $M$ and $M^{\prime}$ may not be directly
compatible, as their input and output features are likely to be very
different. It is also possible that $d_{M}\neq d_{M^{\prime}}$, in which case
composing blocks from $M$ and $M^{\prime}$ is not even well-defined.

Figure 2: Examples of component models used in Manticore. Transformer and
Mamba component models are shown.

Figure 3: The projectors with residual connections used in Manticore, used for
translating features between pretrained blocks of different component models.

To overcome this issue, we apply projectors to both the inputs and the outputs
of a block (or a sequence of blocks, discussed in Section 3.1.4) that we wish
to combine in Manticore hybrids. Overall, our goal in designing projectors is
to enable the blocks of $M$ and $M^{\prime}$ to _speak a common language_ ,
such that their features are compatible and can be reused in the resulting
hybrid model. This is conceivably challenging—the mapping between feature
spaces could be highly nonlinear and might require a lot of task-specific data
to adequately learn the mapping. So do projectors need to be heavyweight,
data-hungry, highly nonlinear objects? Fortunately, the answer is no—we find
that a simple linear transformation with a gated residual, pretrained on
general language data, is sufficient.

Suppose that we want to create a Manticore hybrid from $K$ different
pretrained component models, denoted $M_{(1)},...,M_{(K)}$ with model
dimensions $d_{M_{(1)}},...,d_{M_{(K)}}$. We define
$d_{\max}:=\max_{k\in[K]}d_{M_{(k)}}$, then want input and output projectors
for the blocks of each model that convert their features to a common feature
space of dimension $d_{\max}$. For any sequence of blocks of length
$(n+1)<L_{d_{M_{(k)}}}$ from model $M_{(k)}$ and length-$t$ input,

$\left(M_{(k)\text{Block}}^{(\ell+n)}\circ...\circ
M_{(k)\text{Block}}^{(\ell)}\right):\mathbb{R}^{t\times
d_{M_{(k)}}}\to\mathbb{R}^{t\times d_{M_{(k)}}},$

we want functions $\text{Proj-in}^{(\ell)}_{(k)}:\mathbb{R}^{t\times
d_{\max}}\to\mathbb{R}^{t\times d_{M_{(k)}}}$ and $\text{Proj-
out}^{(\ell+n)}_{(k)}:\mathbb{R}^{t\times d_{M_{(k)}}}\to\mathbb{R}^{t\times
d_{\max}}$, so

$\displaystyle\bigg{(}$ $\displaystyle\text{Proj-out}^{(\ell+n)}_{(k)}\circ
M_{(k)\text{Block}}^{(\ell+n)}\circ...$ $\displaystyle\circ
M_{(k)\text{Block}}^{(\ell)}\circ\text{Proj-
in}^{(\ell)}_{(k)}\bigg{)}:\mathbb{R}^{t\times
d_{\text{max}}}\to\mathbb{R}^{t\times d_{\text{max}}}.$

For input $x\in\mathbb{R}^{t\times d_{M_{(k)}}}$ we parameterize each
projector as a linear transformation with gated residual:

$\text{Proj-
in}^{(\ell)}_{(k)}(x;\alpha):=(1-\alpha)\cdot\text{Linear}_{d_{\text{max}}\to
d_{M_{(k)}}}(x)+\alpha\cdot\text{Trunc}(x;d_{M_{(k)}})$ $\text{Proj-
out}^{(\ell)}_{(k)}(x;\alpha):=(1-\alpha)\cdot\text{Linear}_{d_{M_{(k)}}\to
d_{\text{max}}}(x)+\alpha\cdot\text{Pad}(x;d_{\text{max}}).$

Respectively, $\text{Trunc}(\cdot;d)$ and $\text{Pad}(\cdot;d)$ truncate and
zero-pad input to dimension $d$, and $\text{Linear}_{d\to
d^{\prime}}:\mathbb{R}^{d}\to\mathbb{R}^{d^{\prime}}$ is a learnable linear
transformation. Gating weights are parameterized as $\alpha\in[0,1]$.

In total, where $\alpha\in\Delta^{K-1}$ and $I_{k}$ is a length-$n_{k}$ vector
of block indices from component model $k$, we define the output of the block
sequence defined by $I_{k}$ as

$h_{k}(x;\alpha_{k},I_{k})=\left(\text{Proj-out}^{(I_{k,n_{k}})}_{(k)}\circ
M_{(k)\text{Block}}^{(I_{k,n_{k}})}\circ...\circ
M_{(k)\text{Block}}^{(I_{k,1})}\circ\text{Proj-
in}^{(I_{k,1})}_{(k)}\right)(x;\alpha_{k}).$

####  3.1.3 Mixture Weights

Next, we would like to mix the activations of different component models’
block sequences, in a way that allows us to learn how much _influence_ the
blocks from each component model will have on the overall hybrid model.
Learning the amount of influence that each block sequence should have on the
overall hybrid is critical—if certain blocks produce less helpful features, we
need a way to down-weight them. Conversely, we want to use the best blocks in
our hybrid as much as possible—we want to up-weight these helpful blocks.
Overall, a parameterization that allows us to learn these weights should lead
to better hybrids.

We do this by taking a convex combination of the projectors’ outputs: given
the projected features $h_{k}(x;\alpha_{k},I_{k})$ for each component model
$k\in[K]$, we output a convex combination of projected features

$\text{Mix}_{\alpha}(x;I_{1},...,I_{K})=\sum_{k\in[K]}\alpha_{k}h_{k}(x;\alpha_{k},I_{k}).$

Figure 4: Mixture weights used in Manticore, which learn how much influence
component model blocks should have.

We reuse the convex combination weights as the gating weights in the
projectors. This choice yields the convenient property that when the mixture
weights $\alpha$ are set to one in index $k$ and zero everywhere else, the Mix
function exactly computes a sequence of blocks from component model $k$ while
ignoring the projectors and the blocks from other component models. We adopt a
popular parameterization for mixture weights from the NAS literature [23]: we
parameterize $\alpha$ using a softmax of scalars. That is, we define
$\alpha_{k}:=\frac{\exp(a_{k})}{\sum_{j\in[K]}\exp(a_{j})}$ for all $k\in[K]$.

####  3.1.4 Manticore

We are now ready to define our overall hybrid architecture. We seek to create
a hybrid from $K$ component models, $M_{(1)},...,M_{(K)}$, each with a
potentially different number of blocks, denoted $L_{M_{(k)}}$ for component
model $k$. We fix $L$ to be the number of Manticore blocks, where $L$ is a
common factor of each of the depths $L_{M_{(k)}}$, for all $k\in[K]$—we treat
this choice of factor as a hyperparameter. For each of the $L$ Manticore
blocks, we want to mix a sequence of blocks from each of the $K$ component
models. We also want the number of blocks from each model $k\in[K]$ that are
allocated to a single Manticore block to be evenly spread out throughout the
$L$ Manticore blocks—this is why we require $L$ to be a factor of
$L_{M_{(k)}}$.

For each component model $k\in[K]$, divide the indices of the blocks
$[L_{M_{(k)}}]$ evenly into $L$ contiguous parts, denoted as
$[L_{M_{(k)}}]=(I_{k,1},...,I_{k,L})$. Then, adopting the notation from our
component models, a Manticore block is defined as

$\text{Manticore}_{\text{Block}}^{(\ell)}(\cdot):=\text{Mix}_{\alpha^{(\ell)}}(\cdot;I_{1,\ell},...,I_{K,\ell})$

with $\text{Manticore}_{\text{Block}}^{(\ell)}:\mathbb{R}^{t\times
d_{\text{max}}}\to\mathbb{R}^{t\times d_{\text{max}}}$, for each $\ell\in[L]$,
and $\alpha^{(\ell)}$ being the mixture weights at $\ell$. Next, we initialize
a new set of embedding weights and a new task specific (or language modeling)
head, and we can finally illustrate a forward pass with a Manticore hybrid
model, denoted using the shorthand notation
$\text{Manticore}(\cdot):=\text{Manticore}[M_{(1)},...,M_{(K)}](\cdot)$. Let
$x=(x_{1},...,x_{t})\in\mathcal{V}^{t}$ be a sequence of $t$ tokens from a set
$\mathcal{V}$. The forward pass is computed as follows:

  1. 1.

Apply the new embedding function
$\text{Manticore}_{\text{embed}}:\mathcal{V}^{t}\to\mathbb{R}^{t\times
d_{\text{max}}}$ to the tokens, resulting in a sequence of embeddings denoted
$x_{\text{embed}}=\text{Manticore}_{\text{embed}}(x)$.

  2. 2.

Take forward passes through $L$ Manticore blocks, each with dimension
$d_{\text{max}}$, concretely, we compute
$x_{\ell+1}:=\text{Manticore}_{\text{Block}}^{(\ell)}(x_{\ell})$, where
$x_{1}:=x_{\text{embed}}$.

  3. 3.

Pass $x_{L_{M}+1}$ into a new task-specific or language modeling head,
$\text{Manticore}_{\text{head}}:\mathbb{R}^{t\times d_{M}}\to\mathbb{T}$,
where $\mathbb{T}$ is the appropriate output space for the learning task.

In NAS terms, our search space is over the set of $L\ni\ell$ mixture weights
$\alpha^{(\ell)}\in\Delta^{K-1}$. However, _our search space differs from
typical gradient-based NAS techniques_ in the sense that we do not require
discretization to derive a final architecture after we obtain our mixture
weights. Typically, NAS would involve selecting a single sequence of component
architecture blocks at each of the Manticore blocks, usually by taking the
$\operatorname*{arg\,max}$ of the mixture weights. Instead, the mixtures
themselves are what characterize Manticore hybrids. Nonetheless, if we were to
replace the mixture weights $\alpha^{(\ell)}$ with discrete one-hot vectors,
we could derive any of the following: the component model architectures
themselves, existing hybrid architectures, and ‘frankenmerged’ models [12].

###  3.2 How To Use Manticore

With Manticore, we can automatically select language models without training
every model in the search space, automatically construct pretrained hybrid
architectures without significant trial-and-error, and also program pretrained
hybrids without full training. In this section, we will discuss the details of
how Manticore can be used in each of these three usage scenarios.

Training hybrids from scratch. Manticore can be used to automatically select
LMs without training all of the LMs in the search space. Our selection
technique is simple: inspired by gradient-based NAS techniques [23] and
treating the mixture weights as our ‘architecture parameters,’ we proceed in
two steps: 1\. train mixture weights along with all other parameters, and 2.
freeze the mixture weights and retrain the rest of the parameters from
scratch. Unlike NAS, we found that in many pretraining settings, it was
sufficient to stop at 1. and forgo retraining. In our pretraining experiments,
we primarily use randomly-initialized GPT-Neo [5] and Mamba [13] as component
models without projectors, and separately experiment with a subset of the
blocks from MAD [34].

Fine-tuning pretrained hybrids. Manticore can be used to create and fine-tune
pretrained hybrids. We create pretrained hybrids as follows: begin with a set
of pretrained models, replace their LM heads and embeddings with a single
randomly initialized LM head and embedding layer, and pretrain the projectors
on a small amount of general language data such as FineWeb [30] while keeping
the original component model weights frozen.444 We found one billion tokens to
be sufficient for projector pretraining. To fine-tune the pretrained hybrids
on downstream task data, we first search for mixture weights by training all
of the parameters simultaneously, we freeze the mixture weights, rewind the
component models and projectors to their pretrained state, and fine-tune. This
procedure completely sidesteps large-scale pretraining of new hybrids. In our
synthetic experiments, we create pretrained Manticore hybrids from pretrained
GPT-Neo-125M [5] and Mamba-130M [13] models, while for our experiments on real
natural language data, we opt for pretrained Pythia-410M [4] and Mamba-370M
[13] as component models.

Programming hybrids. Excitingly, we can program Manticore mixture weights by
using external information to predict them.  We consider two scenarios. If we
know that a component model has blocks that are somehow incompatible with the
target task—e.g. resulting from sequence length constraints—we can omit these
blocks by setting their mixture weights to 0. Otherwise, we can predict good
mixture weights by searching on a fixed set of proxy tasks—for this, we use
the MAD tasks [34]. The MAD tasks are synthetic unit tests that are predictive
of hybrid LM scaling laws, but within our framework, we find that they are
also useful for finding general-purpose pretrained hybrids. We use the
following procedure for programming mixture weights using the MAD tasks.
First, run search on the MAD tasks using a smaller, randomly initialized
version of our pretrained hybrid. For each MAD task, our search procedure
returns a set of mixture weights—we simply average the resulting mixture
weights, freeze them, and fine-tune on the downstream task data.

##  4 Experimental Results

We provide experimental evidence that validates the following claims about
Manticore:

  * •

C1.  Pretrained hybrids can outperform their component models on fine-tuning
tasks,

  * •

C2.  Trained from scratch, Manticore hybrids are competitive with existing
hybrids and LMs, and

  * •

C3.  We can program mixture weights using external sources without search on
the task data.

###  4.1 Fine-Tuning Pretrained Hybrids

We evaluate C1, first on a synthetic fine-tuning task, and then on natural
language fine-tuning tasks.

Setup. We consider a synthetic LM dataset comprising GPT-Neo and Mamba
generated completions of text from Penn Treebank [27]. Naturally, we also use
pretrained GPT-Neo-125M and Mamba-130M models as component models, creating a
single Manticore block with projectors that were pretrained on one billion
tokens from FineWeb [30]. We perform search using DARTS and perform post-
search retraining with the model weights and projectors rewound to their
pretrained state.

Results. Our results are shown in Figure 5 (left). We compare our search
results to a sweep over a range of possible mixture weights, and find that our
search procedure returns the optimal mixture weights, outperforming both Mamba
and GPT-Neo. This confirms our claim that Manticore hybrids can outperform
their component models on synthetic fine-tuning tasks. Given that this task
comprises two slices that each of our component models should be good at—GPT-
Neo should be good at predicting GPT-Neo outputs, and vice versa—we
hypothesize that Manticore hybrids are well suited to situations in which the
component models have complementary ‘skills’ [7].

Figure 5:  Mixture weight sweeps on Penn Treebank completions using pretrained
GPT-Neo-125M and Mamba-130M as our component models. (Left) When we create one
Manticore block, there is a region of the search space where we improve over
Mamba. (Right) The same holds for two Manticore blocks, and our technique for
hybrid programming using MAD discovers this region.

Setup. We evaluate on three natural language fine-tuning datasets: Penn
Treebank [27], the Alpaca instructions dataset [37], and ELI5 [10]. We use
Pythia-410M and Mamba-370M as component models, and create a single Manticore
block from their blocks with projectors that were pretrained on one billion
tokens from FineWeb [30]. As before, we first search for mixture weights, and
then we retrain with the fixed mixture weights found by search.

Results. Our results are shown in Table 1. Manticore outperforms its component
models on Alpaca and ELI5, while it achieves performance between its two
component models on Penn Treebank. This confirms our claim that Manticore can
outperform component models on real natural language tasks. The fact that
Mamba-370M outperforms Manticore in this setting is not a failure of our
framework, as Mamba-370M is included as part of our search space. We speculate
that the use of more powerful search procedures from the NAS literature, such
as GAEA [21], could improve our search performance and help to recover or
outperform Mamba-370M.

Table 1:  Manticore on natural language tasks using Pythia-410m and Mamba-370m as component models. The best test losses are bolded and the second-best are underlined.  Task | Pythia-410M (A) | Mamba-370M (B) | Manticore [A, B]
---|---|---|---
Penn Treebank | 0.9099 | 0.8397 | 0.8600
Alpaca | 2.5011 | 2.2999 | 2.1779
ELI5 | 4.1260 | 3.9414 | 3.9331

###  4.2 Training Hybrids from Scratch

For C2, we compare to prior hybrids on MAD and non-hybrid component models on
LRA and MAD.

Setup. We compare training Manticore from scratch to training existing hybrid
architectures on the MAD tasks. We begin with two hybrid architectures from
the literature: Mambaformer [29], which combines Mamba and attention blocks,
and the striped multi-head Hyena + Mixture-of-Experts (MoE) MLP architecture
that was shown to perform well on the MAD tasks [34]. We compare these two
baselines to a Manticore hybrid combining three component models: striped
multi-head Hyena + MoE-MLP, a transformer, and Mamba. We use two blocks for
each of these architectures, creating two Manticore blocks. Again, we search
for mixture weights and then retrain.

Results. The results of this experiment are shown in Table 2. We outperform
the striped multi-head Hyena + MoE model from the MAD paper, and we approach
the performance of Mambaformer on all but one task. This validates the claim
that Manticore hybrids, trained from scratch, compete with existing hybrids.
Despite Mambaformer not being a component model, it is in our search space,
and we again speculate that improvements in search would lead to its recovery.

Table 2:  Trained from scratch on MAD tasks, Manticore beats or matches the performance of existing hybrids on all but one task. The best test losses are bolded and the second best are underlined.  Task | Striped MH Hyena | Mambaformer | Manticore
---|---|---|---
\+ MoE-MLP
In-context Recall | 3.7153 | 0.0020 | 0.0048
Fuzzy In-context Recall | 4.1714 | 4.1712 | 4.1750
Noisy In-context Recall | 4.1643 | 4.1646 | 4.1607
Selective Copying | 1.8021 | 0.0005 | 0.0171
Memorization | 8.8353 | 5.2179 | 8.9254

Setup. We compare Manticore hybrids to their component models on LRA, when
trained from scratch. We create GPT-Neo and Mamba component models of similar
sizes to those in Tay et al. [38] and create a Manticore hybrid. As a
simplified pipeline, we do not retrain model weights after search.

Results. Our results are shown in Table 3. We outperform the component models
on all tasks except for IMDb, in which case Manticore was between GPT-Neo and
Mamba. This validates the claim that Manticore hybrids, trained from scratch,
compete with existing LMs.

Table 3:  Manticore hybrids trained from scratch on LRA using GPT-Neo and Mamba component models. The best test accuracies are bolded. ∗GPT-Neo does not support the Pathfinder-X sequence length requirement, so we set its mixture weight to 0 and Manticore reduces to Mamba.  Task | GPT-Neo (A) | Mamba (B) | Manticore [A, B]
---|---|---|---
ListOps | 37.90 | 20.65 | 38.70
IMDb | 59.62 | 87.74 | 72.44
CIFAR10 | 39.37 | 20.81 | 43.15
Pathfinder32 | 89.41 | 85.76 | 91.45
Pathfinder-X | N/A∗ | $\textbf{75.50}^{*}$ | $\textbf{75.50}^{*}$

Setup. Next, we compare Manticore to non-hybrid architectures trained from
scratch on the MAD tasks. We compare two-block GPT-Neo and Mamba models to a
Manticore hybrid using a single Manticore block. Again, we report the
performance of the search procedure itself without retraining.

Results. Our results are shown in Table 4. Manticore outperforms GPT-Neo and
Mamba on all of the MAD tasks in this setting. This provides further evidence
for our claim that Manticore hybrids compete with existing LMs when trained
from scratch. It is conceivable that our larger Manticore hybrids simply
perform better than component models due to their size—however, we find that
post-search discretization and retraining tends to result in similar
performance, but reduces the model size by roughly half. We include an
ablation of post-search discretization in the Appendix.

Table 4:  Trained from scratch on the MAD tasks, Manticore improves over small two-block component models combined into a single Manticore block. The best test losses are shown in bold.  Task | GPT-Neo (A) | Mamba (B) | Manticore [A, B]
---|---|---|---
In-context Recall | 4.0771 | 4.1858 | 4.0768
Fuzzy In-context Recall | 4.4384 | 4.8097 | 4.2797
Noisy In-context Recall | 4.1843 | 4.2605 | 4.1823
Selective Copying | 1.0470 | 3.7765 | 0.9478
Memorization | 4.6110 | 5.2281 | 4.1367

###  4.3 Programming Hybrids

We evaluate C3 with two types of external data: access to task metadata such
as sequence length requirements, and the use of the MAD tasks as a proxy for
search on downstream task data.

Setup. As in many of our previous experiments, we used the GPT-Neo and Mamba
architectures as component models to our Manticore hybrid. However, this time,
we set out to train from scratch on the extremely long-range Pathfinder-X task
from LRA, which requires sequence length support greater than that of GPT-Neo.
Using this external information about the task, we set the mixture weights for
GPT-Neo to 0, which in this case, means that Manticore reduces to Mamba.555As
of writing, Mamba on LRA is open: https://github.com/state-
spaces/mamba/issues/282.

Results. The results of this experiment are shown in the last row of Table 3.
In the simple case of having access to task metadata, this validates the claim
that we can program mixture weights to exclude incompatible blocks. At the
time of writing, we are not aware of prior published Mamba results on LRA
despite community interest, which would make our evaluation in Table 3 the
first such result. Note that we did not thoroughly tune hyperparameters, so we
view this result as a starting point for the community, rather than a final
answer.

Setup. Finally, in the case in which we can run all of our component models on
our learning task, we program the mixture weights using the MAD tasks as a
search proxy. We set out to fine-tune a pretrained hybrid comprising GPT-
Neo-125M and Mamba-130M with two Manticore blocks on our Penn Treebank
completions synthetic. We train a scaled-down version of this Manticore hybrid
with randomly initialized weights and two blocks per component model on the
MAD tasks. This yields mixture weights for each of the MAD tasks—we average
them and then fine-tune our pretrained hybrid on Penn Treebank completions
using the predicted mixture weights.

Results. Our results are shown in Figure 5 (right). We superimpose the
predicted mixture weights and mean search trajectory from MAD onto the
architecture loss landscape computed on Penn Treebank completions. We find
that this procedure recovers a hybrid that outperforms the component models
(Mamba, lower right; GPT-Neo, upper left) and substantially outperforms the
naive frankenmerges in our search space (upper right and lower left) [12].
This validates the claim that we can program mixture weights using external
sources without performing search on the task data. Intriguingly, search on
the MAD tasks appears to follow the architecture gradient on the different
downstream fine-tuning task, even though the architecture is scaled-down and
trained from scratch on MAD. We suspect that the mixture weights and
architecture loss landscapes for pretrained hybrids are fairly universal
across fine-tuning tasks, and that the same procedure is likely to work on a
wide range of scenarios. Furthermore, we hypothesize that this technique could
outperform other gradient-based NAS methods directly applied to the downstream
task.

## References

  * Akiba et al. [2024] T. Akiba, M. Shing, Y. Tang, Q. Sun, and D. Ha.  Evolutionary optimization of model merging recipes, 2024.
  * Amos et al. [2024] I. Amos, J. Berant, and A. Gupta.  Never train from scratch: Fair comparison of long-sequence models requires data-driven priors.  In _The Twelfth International Conference on Learning Representations_ , 2024.  URL https://openreview.net/forum?id=PdaPky8MUn.
  * Arora et al. [2024] S. Arora, S. Eyuboglu, M. Zhang, A. Timalsina, S. Alberti, D. Zinsley, J. Zou, A. Rudra, and C. Ré.  Simple linear attention language models balance the recall-throughput tradeoff.  _arXiv:2402.18668_ , 2024.
  * Biderman et al. [2023] S. Biderman, H. Schoelkopf, Q. Anthony, H. Bradley, K. O’Brien, E. Hallahan, M. A. Khan, S. Purohit, U. S. Prashanth, E. Raff, A. Skowron, L. Sutawika, and O. van der Wal.  Pythia: A suite for analyzing large language models across training and scaling, 2023.
  * Black et al. [2021] S. Black, L. Gao, P. Wang, C. Leahy, and S. Biderman.  GPT-Neo: Large Scale Autoregressive Language Modeling with Mesh-Tensorflow, Mar. 2021.  URL https://doi.org/10.5281/zenodo.5297715.
  * Botev et al. [2024] A. Botev, S. De, S. L. Smith, A. Fernando, G.-C. Muraru, R. Haroun, L. Berrada, R. Pascanu, P. G. Sessa, R. Dadashi, L. Hussenot, J. Ferret, S. Girgin, O. Bachem, A. Andreev, K. Kenealy, T. Mesnard, C. Hardin, S. Bhupatiraju, S. Pathak, L. Sifre, M. Rivière, M. S. Kale, J. Love, P. Tafti, A. Joulin, N. Fiedel, E. Senter, Y. Chen, S. Srinivasan, G. Desjardins, D. Budden, A. Doucet, S. Vikram, A. Paszke, T. Gale, S. Borgeaud, C. Chen, A. Brock, A. Paterson, J. Brennan, M. Risdal, R. Gundluru, N. Devanathan, P. Mooney, N. Chauhan, P. Culliton, L. G. Martins, E. Bandy, D. Huntsperger, G. Cameron, A. Zucker, T. Warkentin, L. Peran, M. Giang, Z. Ghahramani, C. Farabet, K. Kavukcuoglu, D. Hassabis, R. Hadsell, Y. W. Teh, and N. de Frietas.  Recurrentgemma: Moving past transformers for efficient open language models, 2024.
  * Chen et al. [2023] M. F. Chen, N. Roberts, K. Bhatia, J. WANG, C. Zhang, F. Sala, and C. Re.  Skill-it! a data-driven skills framework for understanding and training language models.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=IoizwO1NLf.
  * Davari and Belilovsky [2023] M.-J. Davari and E. Belilovsky.  Model breadcrumbs: Scaling multi-task model merging with sparse masks.  _ArXiv_ , abs/2312.06795, 2023.  URL https://api.semanticscholar.org/CorpusID:266174505.
  * De et al. [2024] S. De, S. L. Smith, A. Fernando, A. Botev, G. Cristian-Muraru, A. Gu, R. Haroun, L. Berrada, Y. Chen, S. Srinivasan, G. Desjardins, A. Doucet, D. Budden, Y. W. Teh, R. Pascanu, N. D. Freitas, and C. Gulcehre.  Griffin: Mixing gated linear recurrences with local attention for efficient language models, 2024.
  * Fan et al. [2019] A. Fan, Y. Jernite, E. Perez, D. Grangier, J. Weston, and M. Auli.  ELI5: Long form question answering.  In A. Korhonen, D. Traum, and L. Màrquez, editors, _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 3558–3567, Florence, Italy, July 2019. Association for Computational Linguistics.  doi: 10.18653/v1/P19-1346.  URL https://aclanthology.org/P19-1346.
  * Fu et al. [2023] D. Y. Fu, T. Dao, K. K. Saab, A. W. Thomas, A. Rudra, and C. Ré.  Hungry Hungry Hippos: Towards language modeling with state space models.  In _International Conference on Learning Representations_ , 2023.
  * Goddard et al. [2024] C. Goddard, S. Siriwardhana, M. Ehghaghi, L. Meyers, V. Karpukhin, B. Benedict, M. McQuade, and J. Solawetz.  Arcee’s mergekit: A toolkit for merging large language models, 2024.
  * Gu and Dao [2023] A. Gu and T. Dao.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ , 2023.
  * Gu et al. [2022] A. Gu, K. Goel, and C. Re.  Efficiently modeling long sequences with structured state spaces.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=uYLFoz1vlAC.
  * Ilharco et al. [2023] G. Ilharco, M. T. Ribeiro, M. Wortsman, L. Schmidt, H. Hajishirzi, and A. Farhadi.  Editing models with task arithmetic.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=6t0Kwf8-jrj.
  * Jang et al. [2024] D.-H. Jang, S. Yun, and D. Han.  Model stock: All we need is just a few fine-tuned models.  _ArXiv_ , abs/2403.19522, 2024.  URL https://api.semanticscholar.org/CorpusID:268733341.
  * Katharopoulos et al. [2020] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret.  Transformers are rnns: Fast autoregressive transformers with linear attention, 2020.
  * Kim et al. [2023] D. Kim, C. Park, S. Kim, W. Lee, W. Song, Y. Kim, H. Kim, Y. Kim, H. Lee, J. Kim, C. Ahn, S. Yang, S. Lee, H. Park, G. Gim, M. Cha, H. Lee, and S. Kim.  Solar 10.7b: Scaling large language models with simple yet effective depth up-scaling, 2023.
  * Kim et al. [2020] J. Kim, D. Linsley, K. Thakkar, and T. Serre.  Disentangling neural mechanisms for perceptual grouping, 2020.
  * Krizhevsky [2009] A. Krizhevsky.  Learning multiple layers of features from tiny images.  2009\.  URL https://api.semanticscholar.org/CorpusID:18268744.
  * Li et al. [2021] L. Li, M. Khodak, N. Balcan, and A. Talwalkar.  Geometry-aware gradient algorithms for neural architecture search.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=MuSYkd1hxRP.
  * Linsley et al. [2018] D. Linsley, J. Kim, V. Veerabadran, C. Windolf, and T. Serre.  Learning long-range spatial dependencies with horizontal gated recurrent units.  In _Proceedings of the 32nd International Conference on Neural Information Processing Systems_ , NIPS’18, page 152–164, Red Hook, NY, USA, 2018\. Curran Associates Inc.
  * Liu et al. [2019] H. Liu, K. Simonyan, and Y. Yang.  Darts: Differentiable architecture search, 2019.
  * Loshchilov and Hutter [2019] I. Loshchilov and F. Hutter.  Decoupled weight decay regularization.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Bkg6RiCqY7.
  * Maas et al. [2011] A. L. Maas, R. E. Daly, P. T. Pham, D. Huang, A. Y. Ng, and C. Potts.  Learning word vectors for sentiment analysis.  In D. Lin, Y. Matsumoto, and R. Mihalcea, editors, _Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies_ , pages 142–150, Portland, Oregon, USA, June 2011\. Association for Computational Linguistics.  URL https://aclanthology.org/P11-1015.
  * Marcus et al. [1993a] M. P. Marcus, B. Santorini, and M. A. Marcinkiewicz.  Building a large annotated corpus of English: The Penn Treebank.  _Computational Linguistics_ , 19(2):313–330, 1993a.  URL https://aclanthology.org/J93-2004.
  * Marcus et al. [1993b] M. P. Marcus, B. Santorini, and M. A. Marcinkiewicz.  Building a large annotated corpus of English: The Penn Treebank.  _Computational Linguistics_ , 19(2):313–330, 1993b.  URL https://aclanthology.org/J93-2004.
  * Nangia and Bowman [2018] N. Nangia and S. R. Bowman.  Listops: A diagnostic dataset for latent tree learning, 2018.
  * Park et al. [2024] J. Park, J. Park, Z. Xiong, N. Lee, J. Cho, S. Oymak, K. Lee, and D. Papailiopoulos.  Can mamba learn how to learn? a comparative study on in-context learning tasks, 2024.
  * Penedo et al. [2024] G. Penedo, H. Kydlíček, L. von Werra, and T. Wolf.  Fineweb, 2024.  URL https://huggingface.co/datasets/HuggingFaceFW/fineweb.
  * Peng et al. [2023] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derczynski, X. Du, M. Grella, K. Gv, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, J. Lin, K. S. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, J. Wind, S. Woźniak, Z. Zhang, Q. Zhou, J. Zhu, and R.-J. Zhu.  RWKV: Reinventing RNNs for the transformer era.  In H. Bouamor, J. Pino, and K. Bali, editors, _Findings of the Association for Computational Linguistics: EMNLP 2023_ , pages 14048–14077, Singapore, Dec. 2023. Association for Computational Linguistics.  doi: 10.18653/v1/2023.findings-emnlp.936.  URL https://aclanthology.org/2023.findings-emnlp.936.
  * Poli et al. [2023a] M. Poli, S. Massaroli, E. Nguyen, D. Y. Fu, T. Dao, S. Baccus, Y. Bengio, S. Ermon, and C. Ré.  Hyena hierarchy: towards larger convolutional language models.  In _Proceedings of the 40th International Conference on Machine Learning_ , ICML’23. JMLR.org, 2023a.
  * Poli et al. [2023b] M. Poli, J. Wang, S. Massaroli, J. Quesnelle, R. Carlow, E. Nguyen, and A. Thomas.  StripedHyena: Moving Beyond Transformers with Hybrid Signal Processing Models, 12 2023b.  URL https://github.com/togethercomputer/stripedhyena.
  * Poli et al. [2024] M. Poli, A. W. Thomas, E. Nguyen, P. Ponnusamy, B. Deiseroth, K. Kersting, T. Suzuki, B. Hie, S. Ermon, C. Ré, et al.  Mechanistic design and scaling of hybrid architectures.  _arXiv preprint arXiv:2403.17844_ , 2024.
  * Roberts et al. [2021] N. C. Roberts, M. Khodak, T. Dao, L. Li, C. Re, and A. Talwalkar.  Rethinking neural operations for diverse tasks.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, _Advances in Neural Information Processing Systems_ , 2021.  URL https://openreview.net/forum?id=je4ymjfb5LC.
  * Shen et al. [2022] J. Shen, M. Khodak, and A. Talwalkar.  Efficient architecture search for diverse tasks, 2022.
  * Taori et al. [2023] R. Taori, I. Gulrajani, T. Zhang, Y. Dubois, X. Li, C. Guestrin, P. Liang, and T. B. Hashimoto.  Stanford alpaca: An instruction-following llama model.  https://github.com/tatsu-lab/stanford_alpaca, 2023.
  * Tay et al. [2021] Y. Tay, M. Dehghani, S. Abnar, Y. Shen, D. Bahri, P. Pham, J. Rao, L. Yang, S. Ruder, and D. Metzler.  Long range arena : A benchmark for efficient transformers.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=qVyeW-grC2k.
  * Touvron et al. [2023] H. Touvron, L. Martin, K. R. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. M. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. S. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. M. Kloumann, A. V. Korenev, P. S. Koura, M.-A. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, I. Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom.  Llama 2: Open foundation and fine-tuned chat models.  _ArXiv_ , abs/2307.09288, 2023.  URL https://api.semanticscholar.org/CorpusID:259950998.
  * Vaswani et al. [2017] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin.  Attention is all you need.  _Advances in neural information processing systems_ , 30, 2017.
  * Wortsman et al. [2022] M. Wortsman, G. Ilharco, S. Y. Gadre, R. Roelofs, R. Gontijo-Lopes, A. S. Morcos, H. Namkoong, A. Farhadi, Y. Carmon, S. Kornblith, and L. Schmidt.  Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.  _ArXiv_ , abs/2203.05482, 2022.  URL https://api.semanticscholar.org/CorpusID:247362886.
  * Yadav et al. [2023] P. Yadav, D. Tam, L. Choshen, C. Raffel, and M. Bansal.  TIES-merging: Resolving interference when merging models.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=xtaX3WyCj1.
  * Yang et al. [2024] S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim.  Gated linear attention transformers with hardware-efficient training, 2024\.
  * Yu et al. [2023] L. Yu, B. Yu, H. Yu, F. Huang, and Y. Li.  Language models are super mario: Absorbing abilities from homologous models as a free lunch.  _CoRR_ , abs/2311.03099, 2023.  URL https://doi.org/10.48550/arXiv.2311.03099.
  * Zhang et al. [2024] M. Zhang, K. Bhatia, H. Kumbong, and C. Ré.  The hedgehog & the porcupine: Expressive linear attentions with softmax mimicry.  _arXiv preprint arXiv:2402.04347_ , 2024.

## Appendix

##  Appendix A Ablations

Choice of search algorithm. By default, we use a form of the single-level
DARTS [23] search algorithm in all of our experiments requiring search. We
optionally evaluate whether or not to take alternating update, that is, we
alternately take gradient steps in the architecture and model parameters—we
treat this choice as a task-dependent hyperparameter. However, there are many
alternative NAS algorithms that we could have used for search. In our ablation
of the choice of search algorithm, we also evaluate DASH [36] on our Penn
Treebank completions synthetic—the results of which are shown in Table 5. In
general, we found that using DASH was unable to recover strong architectures
in our search space. We postulate that this is because DASH simply aims to
solve a different problem, and is not suited to our search space: namely, DASH
is used to search for lower-level operations, rather than LM blocks. We also
found that alternating DARTS updates was somewhat helpful, compared to
simultaneously updating all of the parameters at once—for our experiments, we
treated this choice as a hyperparameter.

Table 5:  Comparison of NAS search methods on our Penn Treebank completions synthetic.  Alternating? | DARTS | DASH
---|---|---
Yes | 1.2854 | 2.5899
No | 1.3635 | 2.5968

Whether or not to discretize after search. We perform an ablation of whether
or not to perform discretization on our MAD task experiments in which we
compare to existing hybrids. We find that while discretization can sometimes
improve performance, the performance differences are often marginal. If final
parameter count is a concern, then discretization is beneficial.

Table 6:  A comparison of non-discretized vs. discretized Manticore.  Task | Manticore | Manticore
---|---|---
(non-discretized) | (discretized)
In-context Recall | 0.0068 | 0.0081
Fuzzy In-context Recall | 4.1764 | 4.1729
Noisy In-context Recall | 4.1628 | 4.1614
Selective Copying | 0.0849 | 0.0006
Memorization | 8.9416 | 8.9402

##  Appendix B Additional MAD results

In the main text of the paper, we presented results comparing Manticore
hybrids trained from scratch to existing hybrids from the literature—namely
Mambaformer and the Striped MH Hyena + MOE architecture from MAD. Notably, the
Striped MH Hyena + MOE architecture was only the second best architecture
presented in the MAD paper. We found that their best architecture, the Striped
Hyena Experts + MOE model, performed slightly worse on the harder versions of
the MAD tasks that we evaluated. We present these results in Table 7.

Table 7:  Trained from scratch on MAD tasks, Manticore beats or matches the
performance of existing hybrids on all but one task. The best test losses are
bolded and the second best are underlined.

Task | Striped Hyena Experts | Striped MH Hyena | Mambaformer | Manticore
---|---|---|---|---
\+ MoE-MLP | \+ MoE-MLP
In-context Recall | 4.0315 | 3.7153 | 0.0020 | 0.0048
Fuzzy In-context Recall | 4.1749 | 4.1714 | 4.1712 | 4.1750
Noisy In-context Recall | 4.1640 | 4.1643 | 4.1646 | 4.1607
Selective Copying | 2.1731 | 1.8021 | 0.0005 | 0.0171
Memorization | 8.8537 | 8.8353 | 5.2179 | 8.9254

##  Appendix C Additional Pathfinder Results

We ran several additional variants of the pathfinder task for which the
required sequence length exceeded the maximum supported sequence length of
GPT-Neo. We report these results in Table 8.

Table 8: Additional Pathfinder results. Note that since these variants of Pathfinder exceed the maximum sequence length of GPT-Neo, we set its mixture weight to be 0 and evaluate using Mamba. Pathfinder task | GPT-Neo | Mamba | Manticore
---|---|---|---
(A) | (B) | [A, B]
$64\times 64$, 6 paddles | N/A | 80.40 | 80.40
$64\times 64$, 9 paddles | N/A | 90.01 | 90.01
$64\times 64$, 14 paddles | N/A | 86.87 | 86.87
$128\times 128$, 6 paddles | N/A | 75.50 | 75.50

##  Appendix D Hyperparameters

###  D.1 Fine-Tuning Pretrained Hybrids

Penn Treebank completions synthetic. For model weights, we use the AdamW [24]
optimizer with a linear learning rate schedule with an initial learning rate
of $5e-5$. For mixture weights, we use the AdamW [24] optimizer with a linear
learning rate schedule with an initial learning rate of $0.005$ and use
alternating updates.

Fine-tuning on language tasks. For model weights, we use the AdamW [24]
optimizer with a linear learning rate schedule with an initial learning rate
of $5e-5$. For mixture weights, we use the AdamW [24] optimizer with a linear
learning rate schedule with an initial learning rate of $0.005$ and use
simultaneous updates.

###  D.2 Training Hybrids from Scratch

Comparison to existing hybrids on MAD.

We provide the hyperparameters and training details for our MAD evaluations
from Section 4.2

Existing hybrids were trained with a hyperparameter grid search over the space
$[1e-4,5e-4,1e-3]$ for learning rate and $[0.0,0.1]$ for weight decay, similar
to the procedure in MAD [34].

Manticore is trained in two stages. In the first stage, we train the model and
architecture weights in the alternating schedule utilized in DARTS [23]. In
this stage, we perform a hyperparameter grid search of the space
$[1e-4,5e-4,1e-3]$ for model weight learning rate, $[1e-4,1e-4]$ for
architecture weight learning rate, and $[0.1]$ for weight decay. In the second
stage, the architecture weights are frozen and we train only the model weights
using the best learning rate found in the first stage.

Evaluation on LRA. We provide the hyperparameters and training details for our
LRA evaluations.

  * •

ListOps. We trained all models with 5000 steps. The hyperparameter for GPT-Neo
is 8 heads, 6 layers, 512 as the embedding dimension, and 2048 as FFN
dimension. The hyperparameter for Mamba is 12 layers, with 512 as the model
dimension. The vocab size is 18.

  * •

IMDb. We trained all models with 25 epochs and batch size 32.The
hyperparameter for GPT-Neo is 8 heads, 6 layers, 512 as the embedding
dimension, and 2048 as FFN dimension. The hyperparameter for Mamba is 12
layers, with 512 as the model dimension. The vocab size is 129.

  * •

CIFAR10. We trained all models with 10 epochs. The hyperparameter for GPT-Neo
is 4 heads, 3 layers, 64 as the embedding dimension, and 128 as FFN dimension.
The hyperparameter for Mamba is 6 layers, with 64 as the model dimension. The
vocab size is 256, which is the pixel value range of the grayscale image.

  * •

Pathfinder32. We trained all models with 10 epochs. The hyperparameter for
GPT-Neo is 8 heads, 4 layers, 128 as the embedding dimension, and 128 as FFN
dimension. The hyperparameter for Mamba is 8 layers, with 128 as the model
dimension. The vocab size is 256, which is the pixel value range of the
grayscale image.

Comparison to non-hybrids on MAD.

We use two blocks each from GPT-Neo and Mamba, each with a model dimension of
128. We train for 200 epochs and select the best performance during training,
as all of the models overfit across the board. We use the AdamW [24] optimizer
with a linear learning rate schedule with an initial learning rate of $5e-5$.

###  D.3 Programming Hybrids

Mamba evaluation on long Pathfinder tasks. Due to our limited computation
resources, we did not conduct a hyperparameter sweep for the result we
presented. We used Mamba with models of a similar size as Pathfinder32, which
has 8 layers, 128 as the hidden dimension size, and 256 as the vocab size. The
$64\times 64$, 6 paddles version is trained by 10 Epoch with default HP. The
result for other versions is trained with 200 epochs with default HP in
Huggingface trainer.

MAD tasks as a search proxy. For model weights, we use the AdamW [24]
optimizer with a linear learning rate schedule with an initial learning rate
of $5e-5$. For mixture weights, we use the AdamW [24] optimizer with a linear
learning rate schedule with an initial learning rate of $0.01$ and use
simultaneous updates. For search on the MAD tasks, we train scaled-down
versions of GPT-Neo and Mamba each with four blocks, model dimensions of 128,
and no projectors.

###  D.4 Pretraining Projectors

For all non-frozen weights (i.e., projectors, mixture weights, embeddings, and
the LM head), we use the AdamW [24] optimizer with a linear learning rate
schedule with an initial learning rate of $5e-5$.

##  Appendix E Data and MAD Task Parameters

We provide a more detailed description of the datasets that we use in our
experiments. We perform our experiments on a range of synthetic and real tasks
that measure various aspects of modern LM capabilities. We discuss the
specific datasets that we use in our experiments below. MAD synthetics. The
MAD synthetic datasets are a set of tasks introduced by Poli et al. [34] to
systematically evaluate the design space of LMs. These tasks are designed to
serve as proxy unit tests for rapidly prototyping of new hybrid LM
architectures. In our experiments, we use harder variants of the MAD tasks, in
which we use a larger vocabulary size of 128 instead of the default 16 for
most of the tasks, along with fewer training examples. For simplicity, we omit
the compression task as it requires the use of encoder-decoder architectures.

  * •

In-context recall. MAD utilizes a multi-query associative recall task,
challenging models to retrieve values linked to keys within input sequences,
testing their in-context learning ability across randomly shuffled mappings.
We use a vocab size of 128 and 800 training examples.

  * •

Fuzzy in-context recall. This is a variant of in-context recall to assess a
model’s ability to semantically group adjacent tokens. Variable-length keys
and values are randomly paired, testing the model’s capacity for fuzzy recall.
We use a vocab size of 128 and 800 training examples.

  * •

Noisy in-context recall. This is an adaptation of in-context recall to
evaluate a model’s capacity to disregard irrelevant information. This involves
inserting tokens from a separate vocabulary randomly among key-value pairs,
enhancing the memorization challenge. We use a vocab size of 128, a noise
vocab size of 16 with 80% noise, and 800 training examples.

  * •

Selective Copying. MAD employs a selective copying task to evaluate a model’s
ability to remember and replicate specific tokens from an input sequence while
disregarding randomly inserted noise tokens, emphasizing the preservation of
token order. We use a vocab size of 128 with 96 tokens to copy, and 800
training examples.

  * •

Memorization. MAD assesses language models’ factual knowledge retention
through a memorization task, where models learn fixed key-value mappings
without in-context computation, testing pure memorization ability. For this
task, we use a vocab size of 8192.

Long Range Arena. Long Range Arena (LRA) [38] is a benchmark consisting of
various tasks of different modalities that evaluate how well models can learn
long-context data. For simplicity, we omit byte-level document retrieval as it
requires two forward passes per example.

  * •

Long ListOps. This task is designed to understand whether the architecture is
able to model hierarchically structured data in a long-context [28].

  * •

Byte-level text classification. This task attempts to test the model’s ability
to deal with compositionality as in the real world, the model needs to compose
characters into words and words into higher-phrases in not so well defined
boundaries making it a challenging task, we use IMDB dataset[25] in the LRA
paper [38].

  * •

Image classification on a sequence of pixels. This task aims to understand
whether a model is able to capture the 2D spatial structure when presented
with a flattened 1D version of an image to classify, we use pixel information
from CIFAR10[20] dataset.

  * •

Pathfinder. This task helps to understand whether a model can reason about
whether the given 2 dots in an image are connected by a path having dashes or
not. The sequence length is 1024 i.e a 32x32 image is flattened and provided
as input to the model [22, 19].

  * •

Pathfinder-X. An extreme version of Pathfinder with a higher resolution, such
as 64x64 and 128*128, which results in a sequence length of up to 16K

Penn Treebank completions. We generate a synthetic dataset of generated text
from pretrained GPT-Neo-125M [5] and pretrained Mamba-130M models [13]. We
prompt both models using the first four words of every example in the Penn
Treebank [27] validation set, which yields two natural slices of our dataset:
sentence completions generated by GPT-Neo and those generated by Mamba.

Natural language tasks. We evaluate the ability to fine-tune Manticore on
natural language datasets. Specifically, we evaluate on Penn Treebank [26],
the Alpaca instruction tuning dataset [37], and an i.i.d. split of the ELI5
training set [10]. Additionally, we use one billion tokens sampled from the
FineWeb dataset [30] to pretrain our projector weights.

##  Appendix F A Call for Action & Community Recommendations

Throughout our research process, we noted a handful of opportunities that help
to democratize LM research. Should these opportunities be taken up by the
research community, we believe they could help to democratize and help to
decentralize community-driven LM research, all which enabling further research
on pretrained hybrids.

##### A search engine for pretrained models.

Surprisingly, we were unable to easily search for pretrained LMs of certain
sizes or with certain properties (using Huggingface or otherwise). Tools like
this should exist: this would not only significantly democratize LMs, but it
would help to reduce monopolies on LM releases and usage, and thereby
decentralize LM research.

##### Standardized, block-structured LM implementations.

We found that standard tools such as Huggingface and PyTorch were insufficient
to cleanly access intermediate activations across several model
implementations. This could be resolved by adopting standard implementations
or structures for LMs that share the common block structure that we describe
in Section 3.1.1. Instead, our solution was to fork implementations of several
Huggingface models, which is time-consuming, error-prone, and non-scalable. A
solution to this problem would enable and encourage further research on
pretrained hybrid models, which in turn helps to democratize LM research.

##### Removing tokenizers from LM pipelines.

We believe that there are too many possible tokenizers, and that tokenizers
have a significant potential to introduce merge conflicts in model
merging/pretrained hybrid pipelines. In response to this challenge, in our
work, we simply chose an arbitrary tokenizer and relearned our embeddings and
LM head from scratch in all of our experiments. Possible solutions to this
problem would be: as a community, we agree on a standard (small) set of
tokenizers, or we eliminate tokenizers altogether by learning character or
byte-level LMs.

##  Appendix G Limitations

At various points in Section 4, we described limitations with using DARTS (the
off the shelf NAS search algorithm that we used) for search, in that it was
not always able to recover the best architecture in the search space. A
potential limitation of Manticore is that it relies on the existence of good
gradient-based NAS search algorithms, potentially tailored to our search
space. However, we postulate that this is possible, and we leave the task of
developing new search techniques to future work.

##  Appendix H Compute Resources

We ran our experiments on the following GPU hardware:

  * •

2x Nvidia RTX A6000 GPUs with 48GB GPU memory hosted locally in a nook in the
lead author’s house and in a friend’s basement.

  * •

2x Nvidia RTX 4090 GPUs with 24GB GPU memory each hosted locally in other
friends’ basements.

  * •

2x Nvidia Tesla V100 GPUs with 16GB GPU memory each hosted on AWS (p3.2xlarge
instances).

In total, we estimate that our total number of GPU hours across all
experiments (those which failed as well as those included in the paper)
amounted to roughly 750 GPU hours. We estimate that less than half of these
hours accounted for experiments that were not ultimately included in the
paper.

##  Appendix I Broader Impacts and Safeguards

We acknowledge the possibility of misuse with respect to any form of LM
research. In our work, among other things, we enable the creation of
pretrained hybrid models from existing pretrained models. This has potentially
positive and negative social impacts for the community. As a positive
potential social impact, we enable the community to much more easily create
their own hybrid models of various sizes without large scale pretraining—this
has as much potential for positive impact in that these models can be used for
good. On the other hand, the ability to create large pretrained hybrids,
potentially with custom sets of skills, has the potential to open the door to
misuse. To safeguard against such things, we will include appropriate licenses
and rules for usage when we ultimately deploy a Python package for the
community to more broadly use our framework.

##  Appendix J Expanded Version of Figure 5 (Right)

To show how the architectures evolve over search on all of the MAD tasks in
our mixture weights programming experiment, we provide a more detailed version
of Figure 5 (Right) – this is shown in Figure 6. Here, we plot the
architecture trajectories throughout training on all of the MAD tasks, and
superimpose them onto the architecture-loss landscape of the Penn Treebank
completions task. The trajectories roughly follow what appears to be a
gradient in the loss landscape, and all of the trajectories are roughly
similar. We derive our final ‘programmed’ alphas by taking the average of the
final alpha values on each of the MAD tasks, after training.

Figure 6:  Mixture weight sweeps on Penn Treebank completions using pretrained
GPT-Neo-125M and Mamba-130M as our component models. There is a region of the
search space where we improve over Mamba when using two Manticore blocks, and
our technique for hybrid programming using MAD discovers this region.