# OTCE: Hybrid SSM and Attention with Cross Domain Mixture of Experts to
construct Observer-Thinker-Conceiver-Expresser

Jingze Shi  Ting Xie  Bingheng Wu  Chunjun Zheng  Kai Wang

###### Abstract

Recent research has shown that combining Mamba with Transformer architecture,
which has selective state space and quadratic self-attention mechanism,
outperforms using Mamba or Transformer architecture alone in language modeling
tasks. The quadratic self-attention mechanism effectively alleviates the
shortcomings of selective state space in handling long-term dependencies of
any element in the sequence. We propose a position information injection
method that connects the selective state space model with the quadratic
attention, and integrates these two architectures with hybrid experts with
cross-sharing domains, so that we can enjoy the advantages of both. We design
a new architecture with a more biomimetic idea: Observer-Thinker-Conceiver-
Expresser (OTCE), which can compete with well-known medium-scale open-source
language models on a small scale in language modeling tasks.

##  1 Introduction

The Transformers (Attention is All You Need [17]) architecture is popular in
modern deep learning language modeling, which can directly capture the
relationship between any two elements in a sequence, effectively handle long-
distance dependencies, however, the architecture has two main drawbacks.
First, when processing long sequences, its self-attention mechanism’s
quadratic complexity and cache size limit the ability to handle long contexts.
Second, Transformer lacks a single summary state, which means that each
generated token must compute over the entire context.

Meanwhile, the Selective State Model (Mamba [5]) has emerged. Mamba achieves
linear scaling of sequence length during training and maintains a constant
state size during generation through its selective state update mechanism.
Moreover, due to its linear recursive state update mechanism, Mamba has a
single summary state. However, Mamba also has a major drawback, that is, its
positional information depends on the implicit local positional information
provided by the causal convolution, while long-distance dependencies depend on
the matrix D that skips the connection between input and output. This makes
Mamba perform poorly in capturing long-distance dependencies, such as
correctly capturing input-output formats in context learning (ICL).

An efficient model must have a small state, and an effective model must have a
state that contains all the necessary information from the context. To build a
model that is both efficient and effective, the key is to design a state that
is both compact and comprehensive in capturing the necessary context
information. Our main goal is to combine self-attention and the Selective
State Model to overcome their respective limitations, further combine them
with a mixed expert with extensive general and cross-domain knowledge to build
a better basic model architecture than Transformers or Mamba. The model has
the ability to learn long context dependencies, aggregate states, and
efficient reasoning. This paper proposes a bionic perspective, aiming to
explore new model architectures by cleverly combining the Selective State
Model with self-attention mechanism. This approach can fully utilize the
advantages of the two mechanisms and promote language modeling in a more
efficient and effective direction.

#### Positional Information.

We first identified a key challenge in combining the Selective State Model
with self-attention: the effective integration of positional information. In
Mamba, positional information is provided by the implicit local positional
information from the causal convolution, while self-attention itself cannot
provide positional information, it relies on positional encoding to provide
global positional context. To address this issue, we designed a relative
positional information injection method that connects the inner product state
of the Selective State Space with the inner product state of self-attention,
allowing the input gate-state-output gate of the Selective State Space to make
filtered relevant information have discrete relative positional information,
and in self-attention, the discrete relative positional information is re-
continuous to build long-term dependency relationships of relevant
information. This method not only enables our model to have the ability to
selectively process input sequences but also effectively capture long-distance
dependency relationships. In complex multi-query associative recall tasks, our
model trained on the same dataset outperforms larger-scale Mamba, Transformer
models, and models that mix Mamba and Transformer without our proposed
relative positional information injection method.

#### Cross-Domain Mixture of Experts.

In human society, knowledge is widely distributed across different domains,
and these domains are interconnected through common foundational knowledge and
cross-domain connections. To simulate this phenomenon in the model, we
designed two types of cross-domain mixed experts: Cohesive Cross-Domain Expert
and Expansive Cross-Domain Expert. These experts (multi-layer perceptrons)
store and transfer common foundational knowledge and cross-domain knowledge
between different domains by sharing parameters. The Cohesive Cross-Domain
Expert achieves close integration between domains by sharing linear layer
parameters within all experts, which is more suitable for small-scale models
with fewer experts because of its faster computation speed. The Expansive
Cross-Domain Expert shares a complete MLP parameter, adding a common domain
knowledge gate in each expert to control the flow of common domain knowledge
into private MLP parameters, which is more suitable for large-scale models
with more experts because it allows for more flexible adjustment and
utilization of common knowledge. Experimental results show that the
performance of these two cross-domain mixed experts on the same dataset is
better than shared expert isolated mixed experts, confirming the effectiveness
of our design in promoting cross-domain knowledge transfer and improving model
generalization.

#### Architecture Design.

From a biological perspective, the relationship between input and output can
be described by observing, thinking, conceiving, and expressing four stages.
Inspired by this idea, we designed a new architecture with a bionic
perspective: Observer-Thinker-Conceiver-Expresser (OTCE). The OTCE
architecture mimics the natural process of information processing in biology,
aiming to optimize information processing and transmission in a modular way.
In the Observer module, we use the selective state space’s selection ability
to filter out irrelevant information in the sequence to retain relevant
information. In the Thinker module, we use the self-attention’s ability to
capture dependencies between any two elements in a sequence of any length,
regardless of their position in the sequence, to build long-term dependency
relationships. In the Conceiver module, we use the linear recursive state
update mechanism of the state space to build a single summary state
information. In the Expresser module, we combine the context-aware state
information produced by self-attention considering all elements with the
summary state information to build a context-aware summary state. We also
studied the combination of the Selective State Space and self-attention with
ordinary multi-layer perceptrons and cross-domain mixed experts, and finally
determined a model combination with the lowest perplexity at the same
parameter scale.

We empirically validated OTCE on multiple tasks, including semantic similarity
evaluation, long-short text classification, natural language inference,
keyword recognition, different domain selection tasks, context learning, and
multi-query associative recall tasks. These experiments demonstrate the
effectiveness of the OTCE architecture in handling complex language tasks.

Figure 1: (OTCE Architecture.) OTCE demonstrates the overall combined
architecture and process of using Observer, Thinker, Conceiver, and Expresser
modules in language modeling tasks. Observer, Thinker, Conceiver, and
Expresser show their internal combination of selective state space, self-
attention, multi-layer perceptron, and cross-domain mixed experts.

##  2 Background

###  2.1 Selective State Space Models

Selective state space models [5] consider that a fundamental problem in
sequence modeling is to compress the context into a smaller state. From this
perspective, the attention mechanism of Transformers explicitly stores the
entire context information, as if reviewing all previous inputs and generated
tokens before writing each token. In contrast, RNNs only refer to a fixed
number of previous tokens each time, which allows for faster writing but may
forget key tokens.

In the precursor state space models of selective state space models, there are
4 parameters $(\Delta,A,B,C)$, and they do not change with the input. These
parameters control the following two stages:

$\displaystyle h^{\prime}(t)$ $\displaystyle=\bm{A}h(t)+\bm{B}x(t)$ (1a)
$\displaystyle y(t)$ $\displaystyle=\bm{C}h(t)$ (1b)

$\displaystyle h_{t}$
$\displaystyle=\overline{\bm{A}}h_{t-1}+\overline{\bm{B}}x_{t}$ (2a)
$\displaystyle y_{t}$ $\displaystyle=\bm{C}h_{t}$ (2b)

$\displaystyle\bm{\overline{K}}$
$\displaystyle=(\bm{C}\bm{\overline{B}},\bm{C}\bm{\overline{A}}\bm{\overline{B}},\dots,\bm{C}\bm{\overline{A}}^{k}\bm{\overline{B}},\dots)$
(3a) $\displaystyle y$ $\displaystyle=x\ast\bm{\overline{K}}$ (3b)

  * •

In the first stage (1a 1b), the continuous parameters $(\Delta,\bm{A},\bm{B})$
are transformed into discrete parameters
$(\overline{\bm{A}},\overline{\bm{B}})$ through a fixed formula
$\overline{\bm{A}}=f_{A}(\Delta,\bm{A})$ and
$\overline{\bm{B}}=f_{B}(\Delta,\bm{A},\bm{B})$, where $(f_{A},f_{B})$ is
called the discretization rule. The most common is the zero-order hold (ZOH)
rule, defined as $\overline{\bm{A}}=\exp(\Delta\bm{A})$ and
$\overline{\bm{B}}=(\Delta\bm{A})^{-1}(\exp(\Delta\bm{A})-\bm{I})\cdot\Delta\bm{B}$.

  * •

In the second stage (2a 2b and 3a 3b), after the parameters are transformed
from $(\Delta,\bm{A},\bm{B},\bm{C})$ to
$(\overline{\bm{A}},\overline{\bm{B}},\bm{C})$, linear recursion or global
convolution can be used for computation.

State space models cannot update states selectively based on different input
information.

The solution of selective state space models (Mamba [5]) is that, compared to
state space models that compress all historical information, they design a
simple selection mechanism that parameterizes the input of the state space
model, making $(\Delta,\bm{B},\bm{C})$ a function of the input.

Algorithm 1 SSM (S4)

1:$x:\mathtt{(B,L,D)}$

2:$y:\mathtt{(B,L,D)}$

3:$\bm{A}:\mathtt{(D,N)}\leftarrow\mathsf{Parameter}$ $\triangleright$
Represents structured $N\times N$ matrix

4:$\bm{B}:\mathtt{(D,N)}\leftarrow\mathsf{Parameter}$

5:$\bm{C}:\mathtt{(D,N)}\leftarrow\mathsf{Parameter}$

6:$\Delta:\mathtt{(D)}\leftarrow\tau_{\Delta}(\mathsf{Parameter})$

7:$\overline{\bm{A}},\overline{\bm{B}}:\mathtt{(D,N)}\leftarrow\mathsf{discretize}(\Delta,\bm{A},\bm{B})$

8:$y\leftarrow\mathsf{SSM}(\overline{\bm{A}},\overline{\bm{B}},\bm{C})(x)$
$\triangleright$ Time-invariant: recurrence or convolution

9:return $y$

Algorithm 2 SSM + Selection (S6)

1:$x:\mathtt{(B,L,D)}$

2:$y:\mathtt{(B,L,D)}$

3:$\bm{A}:\mathtt{(D,N)}\leftarrow\mathsf{Parameter}$ $\triangleright$
Represents structured $N\times N$ matrix

4:$\bm{B}:{\color[rgb]{0.72,0,0}\mathtt{(B,L,N)}}\leftarrow{\color[rgb]{0.72,0,0}s_{B}(x)}$

5:$\bm{C}:{\color[rgb]{0.72,0,0}\mathtt{(B,L,N)}}\leftarrow{\color[rgb]{0.72,0,0}s_{C}(x)}$

6:$\Delta:{\color[rgb]{0.72,0,0}\mathtt{(B,L,D)}}\leftarrow\tau_{\Delta}(\mathsf{Parameter}{\color[rgb]{0.72,0,0}+s_{\Delta}(x)})$

7:$\overline{\bm{A}},\overline{\bm{B}}:{\color[rgb]{0.72,0,0}\mathtt{(B,L,D,N)}}\leftarrow\mathsf{discretize}(\Delta,\bm{A},\bm{B})$

8:$y\leftarrow\mathsf{SSM}(\overline{\bm{A}},\overline{\bm{B}},\bm{C})(x)$
$\triangleright$ Time-varying: recurrence (_scan_) only

9:return $y$

  * •

The sizes of the matrices $\bm{B}$ that affect the input and $\bm{C}$ that
affect the state change from $(D,N)$ to $(B,L,N)$, i.e., from the fixed input
vector dimension $D$ and the state dimension $N$ of the SSM to the variable
input batch size $B$, sequence length $L$, and state dimension $N$. The size
of $\Delta$ changes from $D$ to $(B,L,D)$, which means that each token in a
batch has a unique $\Delta$. The matrices $\bm{B}$, $\bm{C}$, and $\Delta$ at
each position are different, which means that each input token has different
$\bm{B}$ and $\bm{C}$, solving the content-aware problem.

  * •

The dimension changes are specifically implemented by
$s_{B}(x)=Linear_{N}(x)$, $s_{C}(x)=Linear_{N}(x)$, and
$s_{\Delta}(x)=Linear_{D}(x)$, which sequentially make $(\Delta,B,C)$ data-
dependent.

  * •

Although $A$ has not become data-dependent, after the discretization operation
of the SSM, $(\overline{\bm{A}},\overline{\bm{B}})$ will be transformed into a
data-dependent tensor of $(B,L,D,N)$ through vector products, which is a
parameter-efficient way to achieve data dependency. That is, after the
discretization of $A$, $\overline{\bm{A}}=\exp(\Delta A)$, the data dependency
of $\Delta$ can make the overall $\overline{\bm{A}}$ related to the input.

This allows the model to selectively process information to focus on or ignore
specific information and adaptively adjust its behavior.

###  2.2 Attention and SSM

Attention is a mechanism that computes the relevance scores between each
element in a sequence and all other elements, allowing each element to
"attend" to other elements. The most important variant of attention is softmax
self-attention.

$\displaystyle Y$ $\displaystyle=\operatorname*{softmax}({QK^{\top}})\cdot V$
(4)

Algorithm 3 Self-Attention

$x:\mathtt{(B,L,D)}$

$y:\mathtt{(B,L,D)}$

$Q:\mathtt{(B,L,D)}\leftarrow xW_{Q}$

$K:\mathtt{(B,L,D)}\leftarrow xW_{K}$

$V:\mathtt{(B,L,D)}\leftarrow xW_{V}$

$A:\mathtt{(B,L,L)}\leftarrow\operatorname*{softmax}({QK^{\top}})$

$y:\mathtt{(B,L,D)}\leftarrow AV$

return $y$

A notable feature of self-attention is that it can capture dependencies
between any positions in the input sequence without being limited by distance,
giving it an advantage in handling arbitrary element dependencies in long
sequences.

Many variants of attention have been proposed, all of which are based on the
attention score as the core. Linear attention [10] discards softmax by folding
it into the kernel feature map and rewrites it using the kernel property of
matrix multiplication as $(QK^{\top})\cdot V=Q\cdot(K^{\top}V)$. In the case
of causal (autoregressive) attention, they show that when the causal mask is
merged on the left side as $(L\circ QK^{\top})\cdot V$, where $L$ is a lower
triangular matrix, the right side can be expanded into a recursive form. In
Transformers are SSMs [3], the simple computation of scalar structured SSM is
proven to be equivalent to quadratic masked attention by materializing the
half-separable matrix $M=L\circ CB^{\top}=L\circ QK^{\top}$ and performing
quadratic matrix-vector multiplication.

$\displaystyle Y=(L\circ QK^{\top})\cdot V=(L\circ CB^{\top})\cdot X$ (5)

###  2.3 Mixture of Experts

The mixture of experts architecture [7] aims to train a larger model with
fewer training steps under limited computational resources, often yielding
better results than training a smaller model with more steps. The mixture of
experts architecture consists of two key components:

  * •

The MOE layer, which replaces the feedforward network layer in traditional
Transformer models. Each expert is an independent neural network, typically an
MLP, but it can also be a more complex network structure or even the MOE layer
itself.

  * •

The gating network or router, which determines which tokens are sent to which
expert. A token can be sent to multiple experts. The router consists of
learnable parameters and is trained together with other parts of the network.

To ensure that experts receive non-overlapping concentrated knowledge,
DeepSeekMoE [2] isolates K experts for sharing, aiming to capture common
knowledge and reduce redundancy.

Figure 2: (Shared Expert Isolation.) The hidden state of the entire sequence
calculated by the shared expert is added to the hidden state of each token
calculated by the most relevant routing expert. Each token is composed of the
shared expert state and the routing expert state.

##  3 Methods

###  3.1 Positional Information

#### Rotational Positional Encoding.

The so-called rotational positional encoding, in which the absolute positional
embedding is removed, and the rotational positional encoding is added at each
layer of the network [9], is a form of relative positional encoding
implemented in the form of absolute positional encoding using the idea of
complex numbers.

$\displaystyle Q_{m}$ $\displaystyle=f_{Q}(x_{m},m)$ (6a) $\displaystyle
K_{n}$ $\displaystyle=f_{K}(x_{n},n)$ (6b)

$Q_{m}$ represents the query vector integrated with the positional information
$m$ of the $m$-th token corresponding to the word embedding $x_{m}$, and
$K_{n}$ represents the key vector integrated with the positional information
$n$ of the $n$-th token corresponding to the word embedding $x_{n}$.

#### Rotational Positional Encoding for SSM.

In the matrix transformation form of the state space model mentioned in
Transformers are SSMs [3], by definition, $h_{0}=B_{0}x_{0}$. By induction,

$\displaystyle h_{t}$ $\displaystyle=A_{t}\dots A_{1}B_{0}x_{0}+A_{t}\dots
A_{2}B_{1}x_{1}+\dots+A_{t}A_{t-1}B_{t-2}x_{t-2}+A_{t}B_{t-1}x_{t-1}+B_{t}x_{t}$
$\displaystyle=\sum_{s=0}^{t}A_{t:s}^{\times}B_{s}x_{s}$

Multiplying by $C_{t}$ to produce $y_{t}$, and vectorizing the equation to
$t\in[\mathtt{T}]$ ($\mathtt{T}$ is the sequence length), we derive the matrix
transformation form of SSM.

$\displaystyle y_{t}$
$\displaystyle=\sum_{s=0}^{t}C_{t}^{\top}A_{t:s}^{\times}B_{s}x_{s}$ (7)
$\displaystyle y$ $\displaystyle=\mathsf{SSM}(A,B,C)(x)=Mx$ $\displaystyle
M_{ji}$ $\displaystyle\coloneqq C_{j}^{\top}A_{j}\cdots A_{i+1}B_{i}$

Then the matrix form of SSM is represented using SSS (Sequentially
Semiseparable) as $M=\mathsf{SSS}(A,B,C)$, where
$M_{ji}=C_{j}^{\top}A_{j:i}B_{i}$ 7, and then considering $A$ is just a
scalar, rearranged as

$\displaystyle M_{ji}=A_{j:i}\cdot(C_{j}^{\top}B_{i})$

Vectorized as

$\displaystyle L$ $\displaystyle\coloneqq\mathsf{1SS}(a)$ (8) $\displaystyle
M$ $\displaystyle=L\circ(CB^{\top})$

Finally, it is proved that the matrix transformation form of SSM is equivalent
to Attention $(L\circ QK^{\top})\cdot V=(L\circ CB^{\top})\cdot X$ 5.

Now we have enough theoretical support to give rotational positional encoding
to the $C$ and $B$ matrices in SSM.

$\displaystyle C_{j}$ $\displaystyle=f_{C}(x_{j},j)$ (9a) $\displaystyle
B_{i}$ $\displaystyle=f_{B}(x_{i},i)$ (9b)

$C_{j}$ represents the output weight matrix integrated with the positional
information $j$ of the $j$-th token corresponding to the word embedding
$x_{j}$, and $B_{i}$ represents the input weight matrix integrated with the
positional information $i$ of the $i$-th token corresponding to the word
embedding $x_{i}$.

To utilize the relative positional information between tokens, we assume that
the inner product operation between $C_{j}$ and $B_{i}$ vectors can be
represented by a function $g$, the input of this function $g$ is the word
embedding vectors $x_{j}$ and $x_{i}$, and their relative positional
information $j-i$, the relationship between $C_{j}$ and $B_{i}$ inner product
and their relative positional information $j-i$ is defined as

$<f_{C}(x_{j},j),f_{B}(x_{i},i)>=g(x_{j},x_{i},j-i)$ (10)

We can use the geometric properties of vectors on a 2-dimensional plane,
combined with the properties of complex numbers, to prove

$\displaystyle<f_{C}(x_{j},j),f_{B}(x_{i},i)>$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}&C_{j}^{(2)}\end{bmatrix}\begin{bmatrix}\cos((j-i)\theta)&-\sin((j-i)\theta)\\\
\sin((j-i)\theta)&\cos((j-i)\theta)\end{bmatrix}\begin{bmatrix}B_{i}^{(1)}\\\
B_{i}^{(2)}\end{bmatrix}$ (11)

The proof process is detailed in Appendix A.

In this way, we can apply rotational positional encoding to SSM as shown in
Figure 3 and Algorithm 4.

Recalling the algorithm of the Selective State Space Model 2, the matrix $B$
controls whether to let the input $x_{t}$ affect the state $h_{t}$, after
adding rotational positional encoding to the matrix $B$, the matrix $B$
selectively filters the input $x_{t}$ to make $h_{t}$ have discrete positional
information, the matrix $A$ controls whether the state $h_{t-1}$ affects the
state $h_{t}$, that is, the matrix $A$ also allows the position information of
important states stored in the past state to affect the current state, the
matrix $C$ controls whether the state $h_{t}$ affects the output $y_{t}$,
after the inner product operation between the matrix $C$ and the $h_{t}$ with
discrete positional information, the output $y_{t}$ has discrete relative
positional information, which is undoubtedly an important supplement to the
Selective State Space Model, a recursive model that is difficult to perform
comprehensive context association and memory tasks (including state
interference, no global position as a reference factor in sequence
transformation, etc.).

Figure 3:  (RoPE for SSM.) The structure diagram of applying rotational
positional encoding to the Selective State Space Model.  Algorithm 4 S6 with
RoPE

$x:\mathtt{(B,L,D)},pos_{id}:\mathtt{(B,L)}$

$y:\mathtt{(B,L,D)}$

$\bm{A}:\mathtt{(D,N)}\leftarrow\mathsf{Parameter}$ $\triangleright$
Represents structured $N\times N$ matrix

$cos,sin:{\color[rgb]{0.72,0,0}\mathtt{(L,N)}}\leftarrow{\color[rgb]{0.72,0,0}Ro_{emb}(\mathtt{N})}$

$\bm{B}:{\color[rgb]{0.72,0,0}\mathtt{(B,L,N)}}\leftarrow{\color[rgb]{0.72,0,0}RoPE(cos,sin,pos_{id},s_{B}(x))}$

$\bm{C}:{\color[rgb]{0.72,0,0}\mathtt{(B,L,N)}}\leftarrow{\color[rgb]{0.72,0,0}RoPE(cos,sin,pos_{id},s_{C}(x))}$

$\Delta:{\mathtt{(B,L,D)}}\leftarrow\tau_{\Delta}(\mathsf{Parameter}{+s_{\Delta}(x)})$

$\overline{\bm{A}},\overline{\bm{B}}:{\mathtt{(B,L,D,N)}}\leftarrow\mathsf{discretize}(\Delta,\bm{A},\bm{B})$

$y\leftarrow\mathsf{SSM}(\overline{\bm{A}},\overline{\bm{B}},\bm{C})(x)$
$\triangleright$ Time-varying: recurrence (_scan_) only

return $y$

#### Converting Discrete Position Information to Continuous Position
Information for Attention Weighting.

Although the $y$ output by this kind of selective state space model with
positional encoding is all effective information, the positional information
is discrete. If we don’t provide continuous relative positional information
like Mamba [5] as shown in Figure 4, by convolution to provide continuous
relative positional information for the matrix $D$ that skips the connection
between the input and output of SSM before the SSM matrix transformation, we
want to directly take this $y$ with only discrete relative positional
information for self-attention weighting to build long-term dependency
relationships is completely insufficient.

Figure 4:  (Mamba’s Positional Information.) Mamba provides continuous
relative positional information for the matrix $D$ that skips the connection
between the input and output of SSM by convolution.

If a position has invalid positional information, after the inner product
operation between $Q$ and $K$, this kind of vacancy interference will be
amplified. Therefore, if we want to remove convolution operations or even
remove the matrix $D$, we need to reintroduce the continuous positional
encoding information as shown in Equation 12 before performing the quadratic
self-attention inner product operation again on $y$ that has discrete relative
positional information, as shown in Figure 5.

$\displaystyle<f_{Q}(x_{m},m),f_{K}(x_{n},n)>$
$\displaystyle=\begin{bmatrix}Q_{m}^{(1)}&Q_{m}^{(2)}\end{bmatrix}\begin{bmatrix}\cos((m-n)\theta)&-\sin((m-n)\theta)\\\
\sin((m-n)\theta)&\cos((m-n)\theta)\end{bmatrix}\begin{bmatrix}K_{n}^{(1)}\\\
K_{n}^{(2)}\end{bmatrix}$ (12)

Figure 5:  (RoPE for Quadratic Self-Attention.) Applying rotational positional
encoding to the quadratic self-attention mechanism, reintroducing continuous
positional encoding information before the $QK$ inner product operation.

###  3.2 Cross Domain Mixture of Experts

#### Structure of a Single Expert.

GLU [14] tried gated linear units and variants in the feedforward network
layer of the Transformer, and found that some variants are more effective than
common activation functions such as ReLU and GELU. It consists of the product
of two linear parameters, one of which is passed through an activation
function. Taking the $Swish$ function as an example, it is defined as
$SwiGLU(x,W,V,b,c)=Swish(xW+b)\cdot(xV+c)$, and its feedforward network layer
can be defined as $FFN_{SwiGLU}(x,W,V,W_{2})=(Swish(xW)\cdot(xV))W_{2}$.

Figure 6:  (Loss and Accuracy of MLP with Different Activation Functions.)
$FFN_{GLU}$ and $FFN_{DGLU}$ layers both have three weight matrices, while the
original $FFN$ layer has only two. To keep the number of parameters and
computational complexity unchanged, we reduced the second dimension of $W$ and
$V$ of $FFN_{GLU}$ and $FFN_{DGLU}$ layers and the first dimension of $W_{2}$
to be consistent with the original $FFN$ layer during validation.

We tried double gated linear units, that is, both linear transformations are
passed through activation functions, that is, double gate linear unit (DGLU),
and we used two activation functions in the other branch, which are $Sigmod$
and $Tanh$.

$\displaystyle FFN_{SwiSigDGLU}(x,W,V,W_{2})$
$\displaystyle=(Swish(xW)\cdot\sigma(xV))W_{2}$ (13a) $\displaystyle
FFN_{SwiTanhDGLU}(x,W,V,W_{2})$ $\displaystyle=(Swish(xW)\cdot\tanh(xV))W_{2}$
(13b)

The results are shown in Figure 6. In the two-category classification task,
the loss curve of the $FFN_{DGLU}$ layer using the $\sigma$ function is the
lowest, and the accuracy of the $FFN_{DGLU}$ layer using the $\sigma$ function
is much higher than other layers. In the three, fifteen, and one hundred
category classification tasks, the loss curve of the $FFN_{DGLU}$ layer using
the $\tanh$ function is the lowest. In the fifteen category classification
task, the accuracy of the $FFN_{DGLU}$ layer using the $\tanh$ function
gradually surpasses other activation functions, and the accuracy of the
$FFN_{DGLU}$ layer using the $\tanh$ function is the highest in the one
hundred category classification task.

In the autoregressive language reasoning task, it can be regarded as a
continuous classification task with a category size of the vocabulary size. As
the number of categories continues to increase, the loss and accuracy of the
$FFN_{DGLU}$ layer using the $\tanh$ function gradually surpass other
activation functions. Theoretically, using the $FFN_{SwiTanhDGLU}$ layer will
achieve better results in pre-training perplexity and downstream language
understanding. Therefore, we use the $FFN_{SwiTanhDGLU}$ layer as the
structure of a single expert in the mixture of experts.

#### Mixture of Experts for Transformers.

A standard Transformer language model is built by stacking $L$ Transformer
blocks, each of which consists of a self-attention block and a feedforward
network block, and its output hidden state can be defined as follows:

$\displaystyle a^{l}_{1:T}$ $\displaystyle=\textbf{Self-
Attention}(h^{l-1}_{1:T})+h^{l-1}_{1:T}$ (14a) $\displaystyle ah^{l}_{t}$
$\displaystyle=\textbf{FFN}(a^{l}_{t})+a^{l}_{t}$ (14b)

where $T$ is the sequence length, $Self-Attention$ is self-attention, $FFN$ is
the feedforward network, $a^{l}_{1:T}\in\mathbb{R}^{T\times d}$ is the hidden
state of all tokens after the $l$-th self-attention block, and
$ah^{l}_{t}\in\mathbb{R}^{d}$ is the hidden state of the $t$-th token after
the $l$-th feedforward network block. For simplicity, we omit normalization in
the above formula.

The classic way to construct a mixture of experts language model is to replace
the $FFN$ block in the Transformer at specified intervals, where each expert
is structurally the same as the standard $FFN$ block, and then assign each
token to the topk experts. Replace the $l$-th $FFN$ block with the $MOE$
block, and its output hidden state can be defined as follows:

$\displaystyle ah^{l}_{t}$
$\displaystyle=\sum_{i=1}^{N}(g_{i,t}\cdot\textbf{FFN}_{i}(a^{l}_{t}))+a^{l}_{t}$
(15a) $\displaystyle g_{i,t}$
$\displaystyle=\begin{cases}aff_{i,t}&aff_{i,t}\in\text{topK}({aff_{j,t}|1\leq
j\leq N},K)\\\ 0&\text{otherwise}\end{cases}$ (15b) $\displaystyle aff_{i,t}$
$\displaystyle=\textbf{Softmax}_{i}(a^{l^{T}}_{t}e^{l}_{i})$ (15c)

where $N$ is the number of experts, $\textbf{FFN}_{i}$ is the feedforward
network of the $i$-th expert, $g_{i,t}$ is the gate value assigned to the
$i$-th expert for the $t$-th token, $aff_{i,t}$ is the affinity between the
token and the expert, $\text{topK}(\cdot,K)$ is the topk selection function of
the affinity between the $t$-th token and all experts, and $e^{l}_{i}$ is the
centroid of the $i$-th expert in the $l$-th layer. For simplicity, we omit
normalization in the above formula.

#### Mixture of Experts for SSM.

Following the idea of Mixtures of Experts for Transformers, we can apply it to
the state space model.

$\displaystyle s^{l}_{1:T}$
$\displaystyle=\textbf{SSM}(h^{l-1}_{1:T})+h^{l-1}_{1:T}$ (16a) $\displaystyle
sh^{l}_{t}$ $\displaystyle=\textbf{FFN}(s^{l}_{t})+s^{l}_{t}$ (16b)

where $T$ is the sequence length, $SSM$ is the state space model, $FFN$ is the
feedforward network, $s^{l}_{1:T}\in\mathbb{R}^{T\times d}$ is the hidden
state of all tokens after the $l$-th state space model block, and
$sh^{l}_{t}\in\mathbb{R}^{d}$ is the hidden state of the $t$-th token after
the $l$-th feedforward network block.

Imitating the conventional routing strategy of the mixture of experts language
model, we assign each token to the topk experts, replace the $l$-th $FFN$
block with the $MOE$ block, and its output hidden state can be defined as
follows:

$\displaystyle sh^{l}_{t}$
$\displaystyle=\sum_{i=1}^{N}(g_{i,t}\cdot\textbf{FFN}_{i}(s^{l}_{t}))+s^{l}_{t}$
(17a) $\displaystyle g_{i,t}$
$\displaystyle=\begin{cases}aff_{i,t}&aff_{i,t}\in\text{topK}({aff_{j,t}|1\leq
j\leq N},K)\\\ 0&\text{otherwise}\end{cases}$ (17b) $\displaystyle aff_{i,t}$
$\displaystyle=\textbf{Softmax}_{i}(s^{l^{T}}_{t}e^{l}_{i})$ (17c)

where $N$ is the number of experts, $\textbf{FFN}_{i}$ is the feedforward
network of the $i$-th expert, $g_{i,t}$ is the gate value assigned to the
$i$-th expert for the $t$-th token, $aff_{i,t}$ is the affinity between the
token and the expert, $\text{topK}(\cdot,K)$ is the topk selection function of
the affinity between the $t$-th token and all experts, and $e^{l}_{i}$ is the
centroid of the $i$-th expert in the $l$-th layer.

Figure 7: (Illustration of Cross Domain Mixture of Experts.) These three
figures respectively show a single cohesive cross domain expert, a single
expansive cross domain expert, and the routing expert calculation method,
which together form a complete cross domain mixture of experts. It is worth
noting that in the specific implementation, Shared Linear can be defined
externally and then passed into each private expert, or the shared hidden
states can be calculated externally and then passed into each private expert.

#### Shared Strategy.

Under the conventional routing strategy of the mixture of experts, tokens
assigned to different experts need common knowledge or information, so
multiple routing experts will have redundant parameters that store common
information when obtaining their respective parameters. This leads to expert
parameter redundancy. However, if these experts have partially shared
parameters or the parameters that store common knowledge have already been
passed through the parameters, then parameter redundancy can be reduced. The
reduction of redundancy will help build a more efficient and specialized
mixture of experts.

#### Cohesive Cross Domain Experts.

In Attention-MOE 15a and SSM-MOE 17a, expand $FFN$ into $FFN_{SwiTanhDGLU}$
13b, which has three weight matrices $W$, $V$, $W_{2}$. We take matrix $V$ as
the shared parameter of different $FFN$ in the mixture of experts, and take
$W$ and $W_{2}$ as the private parameters of each $FFN$. Re-represent the
parameters as $W^{p}$, $V^{s}$, $W_{2}^{p}$, where $p$ is the private
parameter and $s$ is the shared parameter, and keep other parameters and
formulas unchanged.

$\displaystyle FFN^{Cohesive}(x,W^{p},V^{s},W_{2}^{p})$
$\displaystyle=(Swish(xW^{p})\cdot\tanh(xV^{s}))W_{2}^{p}$ (18)

This way of sharing parameters inside $FFN$ can reduce parameter redundancy
and the total number of parameters. We call it cohesive cross domain experts.
However, this method is only suitable for cases with a small number of
experts, because the shared parameter $V^{s}$ cannot adjust the number of
parameters independently and must be consistent with $W^{p}$, which will cause
the proportion of shared parameters to decrease as the number of experts
increases.

#### Expansive Cross Domain Experts.

To cope with a large number of experts while maintaining the proportion of
shared parameters unchanged, still taking $FFN_{SwiTanhDGLU}$ 13b as an
example, use $p$ to represent private parameters, $s$ to represent shared
parameters, define $W^{p}$, $V^{p}$, $W_{2}^{p}$ as private $FFN$ parameters,
$W^{s}$, $V^{s}$, $W_{2}^{s}$ as shared $FFN$ parameters, and $W_{3}^{p}$ as
the gate parameter that receives shared hidden states to private parameters in
each private $FFN$. Re-represent the Attention-MOE 15a and SSM-MOE 17a as
expansive cross domain experts, keeping other parameters and formulas
unchanged.

$\displaystyle FFN^{p}(x,W^{p},V^{p},W_{2}^{p})$
$\displaystyle=(Swish(xW^{p})\cdot\tanh(xV^{p}))W_{2}^{p}$ (19a)
$\displaystyle FFN^{s}(x,W^{s},V^{s},W_{2}^{s})$
$\displaystyle=(Swish(xW^{s})\cdot\tanh(xV^{s}))W_{2}^{s}$ (19b)
$\displaystyle FFN^{Expansive}(x,W_{3}^{p})$
$\displaystyle=FFN^{p}(FFN^{s}(x)\cdot FFN^{s}(x)W_{3}^{p})$ (19c)
$\displaystyle ah^{l}_{t}$ $\displaystyle=\sum_{i=1}^{N}(g_{i,t}\cdot
FFN^{Expansive}_{i}(s^{l}_{t}))+s^{l}_{t}$ (19d) $\displaystyle sh^{l}_{t}$
$\displaystyle=\sum_{i=1}^{N}(g_{i,t}\cdot
FFN^{Expansive}_{i}(s^{l}_{t}))+s^{l}_{t}$ (19e)

In the expansive cross domain experts, tokens routed to various experts first
pass through the shared $FFN$ parameters that store common knowledge, and then
pass through the gate of the private $FFN$ to let the private experts decide
which information from the shared hidden states they should receive. Compared
with the shared expert isolation 2, the expansive cross domain experts achieve
more accurate and targeted knowledge acquisition capabilities, and retain the
advantage of being able to adjust the proportion of shared parameters
arbitrarily.

###  3.3 Architecture Design

#### Self-Attention’s Characteristics and Drawbacks.

Let’s review the algorithm of standard self-attention 3, Self-attention’s
parameters include the weight matrices $W_{Q}$, $W_{K}$, $W_{V}$ and bias of
$Q$, $K$, $V$, and the complete self-attention layer also includes an output
weight matrix $W_{O}$ and bias.

First, calculate $Q$, $K$, $V$ that is $Q=xW_{Q}$, $K=xW_{K}$, $V=xW_{V}$, the
input and output shapes of this matrix multiplication are
$[B,L,D]\times[D,D]=[B,L,D]$, the computational complexity is $3\times
2BLD^{2}=6BLD^{2}$.

Then calculate the matrix multiplication of $QK^{T}$, the input and output
shapes are $[B,L,D]\times[B,D,L]=[B,L,L]$, the computational complexity is
$2BL^{2}D$.

Then calculate the weighted sum on $V$, the input and output shapes of this
matrix multiplication are $[B,L,L]\times[B,L,D]=[B,L,D]$, the computational
complexity is $2BL^{2}D$.

Finally, it is the linear mapping after self-attention, the input and output
shapes of this matrix multiplication are $[B,L,D]\times[D,D]=[B,L,D]$, the
computational complexity is $2BLD^{2}$.

In summary, the computational complexity of the self-attention layer is
$8BLD^{2}+4BL^{2}D$, the total computational complexity of the multi-head
self-attention variant is the same, but it can be calculated in parallel when
calculating $(QK^{T})V$, which is faster.

Next, we can calculate the intermediate activation memory usage of the self-
attention layer $x_{out}=softmax(\frac{QK^{T}}{\sqrt{d}})V\cdot W_{O}+x$.

For $Q$, $K$, $V$, we need to save their common input $x$, the size of $x$ is
$[B,L,D]$, the number of elements is $BLD$, the memory usage is $2BLD$.

For the matrix multiplication of $QK^{T}$, we need to save $Q$, $K$, the
shapes of these two tensors are $[B,L,D]$, the memory usage is $4BLD$.

For $softmax(\frac{QK^{T}}{\sqrt{d}})$, we need to save the input of the
function $QK^{T}$, the memory usage is $2BL^{2}h$, where $h$ is the number of
heads of multi-head attention, the shape of $Q$ is
$[B,head\\_num,L,per\\_head\\_dim]$, the shape of $K$ is
$[B,head\\_num,per\\_head\\_dim,L]$, the shape of $QK^{T}$ is
$[B,head\\_num,L,L]$, after calculating $softmax$, a $dropout$ operation will
be performed, we need to save a $mask$ matrix, the shape of the $mask$ matrix
is the same as $QK^{T}$, the memory usage is $BL^{2}h$.

Calculate the attention on $V$, that is $AV$, we need to save $A$, the memory
usage is $2BL^{2}h$, and $V$, the memory usage is $2BLD$, the memory usage is
$2BL^{2}h+2BLD$.

Calculate the final output mapping and a $dropout$ operation, the output
mapping needs to save its input $y$, the memory usage is $2BLD$, the $dropout$
operation needs to save a $mask$ matrix, the memory usage is $BLD$, the memory
usage is $2BLD+BLD=3BLD$.

Therefore, by adding the memory usage of all the intermediate activations
above, the memory usage of the intermediate activations of the self-attention
layer is $11BLD+5BL^{2}h$.

Through the above calculation, we can see that the computational complexity
and storage complexity of self-attention increase quadratically with the
sequence length $L$, which limits the maximum sequence length $L$, but we can
easily see two characteristics of self-attention from the formula:

  * •

Self-attention does not compress context information, that is, it can capture
the dependency relationship between any position elements in the sequence,
regardless of distance.

  * •

Self-attention weights each element, each element considers the information of
all other elements in the sequence, and produces a context-aware state.

#### SSM’s Characteristics and Drawbacks.

The computational complexity of SSM increases linearly with the sequence
length $L$, and we will not go into details here. Let’s review the Selective
State Model 2.1, which has two important characteristics:

  * •

The data-dependent $\overline{\bm{B}}$ matrix more finely controls whether to
let the input $x$ affect the state $h$, the data-dependent $\overline{\bm{A}}$
matrix controls the influence from the previous state $h_{t-1}$ to the current
state $h_{t}$, and the data-dependent $\bm{C}$ matrix controls whether the
state $h$ can affect the output $y$, the data dependency of these three
matrices gives SSM the ability to selectively process information.

  * •

Because the number of parameters in the word embedding matrix is also
relatively large, the linear recursive aggregation state assigned to the last
token allows us to calculate the matrix multiplication of
$[B,D]\times[D,V]=[B,V]$ to obtain $logits$.

Of course, the drawbacks of SSM are also obvious, by approximating the optimal
solution of the state matrix $A$ through functions to remember history, it can
only accurately capture the state change from $h_{t-1}$ to $h_{t}$, gradually
decay the old state information, which leads to SSM performing poorly in tasks
that require strong copying or context learning ability or long context
reasoning.

#### Combination of Attention and SSM.

If a human who loves Shiba Inu reads an article, we can divide it into four
stages: observation, thinking, conception, and expression.

Figure 8:  Cheems. The prototype of the Shiba Inu Cheems that became popular
on the Internet. If you are a human who has been stamped with the idea of
loving Shiba Inu (uncompressed memory dependency of any length), you should
have a sudden realization — the process of observation, thinking, conception,
and expression corresponds to the combination of SSM and self-attention in the
Observer, Thinker, Conceiver, and Expresser modules of the OTCE architecture.

In the observation stage, we are often attracted by the pictures first, and
this interesting picture priority corresponds to the selective processing
ability of SSM for information.

In the thinking stage, we carefully observe the decadent expression and
sitting posture of the creature in the picture to think about the relevance to
the love of Shiba Inu, this taking the current information and past memory to
build long-term dependencies corresponds to the global dependency of self-
attention.

In the conception stage, we combine the current information with relevant past
memories to deduce the logical relationship step by step and finally draw a
conclusion, which corresponds to the linear recursive aggregation state of
SSM.

In the expression stage, we accurately say the name found from the distant
memory: "This is (surprised) Cheems (love)!", which corresponds to the
context-aware state of self-attention considering all elements.

This example is not very appropriate, it is more appropriate to compare it
with looking up in the phone book, but under our strong desire, we finally
chose to compare it with the Internet celebrity dog Cheems 8. These four
stages correspond to the combination of SSM and self-attention in the
Observer, Thinker, Conceiver, and Expresser modules of the OTCE architecture
1.

It is worth noting that we design the Expresser module 9 to consider the
information of all elements of the linear recursive aggregation state of the
previous layer SSM after weighting with self-attention and then add it to it,
and finally obtain the final output through a combination of a linear layer,
an activation function, and another linear layer. The purpose of doing this is
to prevent the bias caused by too many SSM layers after the Thinker module,
and at the same time, it can retain the information of the aggregation state
for efficient reasoning.

Figure 9:  Expresser. The Expresser module reweights the information with
gradually decaying long-term dependency relationships through re-attention to
all elements, and learns how SSM recursively aggregates information to correct
knowledge routing bias. Maybe we can do something on the branch that skips the
connection attention.

#### Combination of MLP and MOE.

After ablation experiments, we found that using MLP in the first layer and the
last layer of the model performs best on the pre-training perplexity. Since
the self-attention part of the Expresser considers the information of all
elements of the linear recursive aggregation state of the previous layer SSM
after weighting and then adds it to it, we also consider it as the last layer,
that is, the first layer of the Observer module and the last layer of the
Conceiver and Expresser modules use MLP, and other layers use the combination
of MOE is the best. Ablation experiments will be detailed in the following
sections.

##  4 Empirical Validation

We empirically validated OTCE on semantic similarity, long-short text
classification, natural language inference, keyword recognition, different
domain selection questions, context learning, and multi-query associative
recall tasks. The importance of re-attention weighting in the Expresser module
before output is verified in Section 4.2, and the use of positional encoding
in SSM and Attention to connect the token positions in the hybrid architecture
is verified in Section 4.3. The ablation of different combinations of MOE and
MLP in the OTCE architecture, and the comparison of MOE with different shared
parameter strategies in pre-training, fine-tuning, and multi-domain selection
question tasks are verified in Section 4.4, and the verification on context
learning and multi-query associative recall tasks is verified in Section 4.5.

###  4.1 Training Dataset and Hyperparameters

The pre-training dataset is a mixture of open-source datasets including Book,
Wikipedia, UNCorpus, translation2019zh, WikiMatrix, news-commentry, ParaCrawl
v9, ClueCorpusSmall, CSL, news-crawl, etc., and the SFT fine-tuning dataset is
2 million instructions from the Belle dataset.

The learning rate formula comes from Attention is All You Need [17], to pay
tribute to the first paper we read as a sophomore ordinary student starting
from scratch to study machine learning.

$lrate=d_{model}^{-0.5}\times\min(step\\_num^{-0.5},step\\_num\times
warmup\\_steps^{-1.5})$

In addition, we did not use linear bias terms, the warm-up steps are $10\%$ of
the total steps, $RMSNorm$ is used instead of $LayerNorm$, $AdamW$
hyperparameters are set to $\beta_{1}=0.9,\beta_{2}=0.98,\epsilon=10^{-6}$,
and GLM [4] is used as the architecture backbone, that is, bidirectional GLM
universal language model: MASK spans in the training text sequence, with
autoregressive filling as the training target, requiring the model to predict
the masked token. The validation environment is the open-source PyTorch image
provided by Nvidia [13].

###  4.2 Attention before Output is Important

In this section, we will not adopt most of the improvements mentioned earlier,
but focus on verifying the importance of attention weighting the recursively
aggregated state of SSM at the output stage. As shown in Table 1, when the
number of layers is 8, one of the combinations of OTCE is very similar to a
single Jamba block of another hybrid architecture Jamba [12]. The main
difference between the two is only in the distribution strategy of MLP and
MOE, and the last layer of OTCE is the Expresser module.

Table 1:  (Jamba and OTCE Parameters.) The parameters of Jamba and OTCE with 8 layers, the positional information comes from the convolution before each SSM, without RoPE, and without MOE, the difference is only in the last layer module, Jamba is SSM + MLP, OTCE is Attention + MLP.  Model | Params | Layers | Model Dim | FFN Dim | Attn Heads | Group | State Dim | Hybrid (S=SSM, A=Attention, M=MLP)
---|---|---|---|---|---|---|---|---
Jamba | 283M | 8 | 1024 | 4096 | 16 | 8 | 16 | SMSMSMSMAMSMSMSM
OTCE | 281M | 8 | 1024 | 4096 | 16 | 8 | 16 | SMSMSMSMAMSMSMAM
Table 2:  (Jamba and OTCE Results.) The perplexity performance of Jamba and OTCE in the pre-training stage is very close, but in the SFT fine-tuning stage, the perplexity of OTCE is significantly lower than Jamba. We select the evaluations for the CLUEbenchmark [19] that are more suitable for the current parameter volume: AFQMC: semantic similarity, judge whether two sentences have the same meaning. TNEWS: short text classification, from the news section of today’s headlines, judge the category of the news. IFLYTEK: long text classification, a long text dataset about software descriptions, judge the application theme of the software description. CMNLI: language reasoning, judge the logical relationship between two sentences. WSC: pronoun disambiguation, judge which noun the pronoun in the sentence refers to. CSL: keyword recognition, taken from the abstract and keywords of the paper, judge whether the keyword is in the abstract. OTCE maintains a slight lead in AFQMC, TNEWS, IFLYTEK, WSC, and CSL, and slightly lags behind in CMNLI, with an average accuracy improvement of $4.25\%$ over Jamba.  Model | PreTrain | FineTune | AFQMC | TNEWS | IFLYTEK | CMNLI | WSC | CSL | Average
---|---|---|---|---|---|---|---|---|---
| ppl $\downarrow$ | ppl $\downarrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$
Jamba | 2.53 | 1.98 | 78.65 | 63.15 | 72.40 | 79.66 | 76.64 | 84.83 | 75.89
OTCE | 2.51 | 1.69 | 79.48 | 77.98 | 72.71 | 79.59 | 79.93 | 85.06 | 79.12

As shown in Table 2, the difference in perplexity between Jamba and OTCE in
the pre-training stage is negligible, which indicates that both have learned
almost the same knowledge from pre-training. However, in the SFT fine-tuning
stage, the model usually needs to remember specific formats based on the
context and reason based on these formats in the subsequent context. This has
been proven to be a deficiency in the ability of SSM to aggregate state
information recursively. The Attention part in a single Jamba block is
embedded between the SSM layers, which leads to the ability of attention to
directly route the knowledge of each answer to a single answer token being
attenuated by the subsequent three SSM layers, resulting in a bias in
knowledge routing.

The Expresser module of OTCE re-weights the state information recursively
aggregated by SSM with attention, which is equivalent to learning and
weighting how the information from the first token to the last token is
recursively aggregated. Then, the two types of state information are mixed
together, effectively avoiding the bias in knowledge routing. Therefore, we
can now explain why OTCE has a significant advantage in most validation tasks.

For tasks such as AFQMC, which judge whether the meanings of sentences are the
same, CSL, which judge whether keywords are in the abstract rather than
extracting keywords from the abstract, and CMNLI, which judge whether the
relationship between two sentences is entailment, neutral, or contradiction,
they mainly rely on selectively extracting key information on the sequence and
recursively aggregating state information. Therefore, the performance gap
between Jamba and OTCE on these tasks is not significant, and the reason why
OTCE slightly outperforms Jamba in AFQMC and CSL may be that it corrects some
knowledge routing biases.

In tasks such as TNEWS, which predict the category of short sentences, since
the sentences do not contain or contain very little context information, they
are very dependent on the quality of pre-training and fine-tuning. This
requires that the knowledge learned in the pre-training stage can be correctly
and accurately guided in the fine-tuning stage, which leads to Jamba’s
performance on TNEWS being far worse than OTCE.

In contrast to TNEWS, the text in the IFLYTEK task is all long context, and
the vocabulary of the application theme appears multiple times. This
constitutes a relatively simple single-query associative recall task, where
the vocabulary of the single main theme that appears multiple times is allowed
to be stored in the state by the selective state space model and becomes the
main factor affecting the output. Therefore, the performance gap between the
two on IFLYTEK is not significant.

In the WSC task, which is a pronoun disambiguation task, the model needs to be
able to trace back to the noun that appeared before the pronoun after the
pronoun appears and point the pronoun to the correct noun. This belongs to the
additional requirement of understanding the current language in the single-
query associative recall task. Therefore, the reason why the performance gap
between the two on WSC is second in these tasks is obvious, with the two main
factors being the difference in fine-tuning quality and the re-attention of
the recursive aggregation information in the Expresser module of OTCE on the
long-term strong dependency relationship.

###  4.3 Positional Encoding that connected SSM to Attention

In this section, we evaluate the positional encoding strategy that connects
SSM to Attention. The parameter settings used are consistent with the 281M
OTCE model in Table 1.

We conducted four pre-training experiments, namely without using RoPE, only
using Attention RoPE, only using SSM RoPE, and using both Attention and SSM
RoPE. The experiments used a sequence length of 8192 and a RoPE base frequency
of 10000. The experimental results are detailed in Table 3.

Table 3:  (Evaluation of Different Position Information Sources.) Pre-training perplexity performance without RoPE, only with Attention RoPE, only with SSM RoPE, and with both Attention and SSM RoPE, with Attention accounting for $12.5\%$, SSM accounting for $37.5\%$, and MLP accounting for $50\%$.  Non RoPE | Only Attention RoPE | Only SSM RoPE | Both Attention and SSM RoPE
---|---|---|---
ppl $\downarrow$ | ppl $\downarrow$ | ppl $\downarrow$ | ppl $\downarrow$
2.51 | 2.49 | 2.26 | 2.10

From the experimental results, we can observe that there is little difference
in perplexity performance between not using RoPE and only using Attention
RoPE. This finding is consistent with recent large-scale studies on hybrid
architectures [18].

However, when only using SSM RoPE, the perplexity decreases by nearly $10\%$
compared to not using RoPE. Moreover, when using both Attention and SSM RoPE,
the perplexity decreases by nearly $7\%$ compared to only using SSM RoPE.

This indicates that the one-dimensional convolution layer before SSM provides
local positional information in the model’s hidden dimension, promoting the
consideration of adjacent positional information during linear recursive state
updates, which is also applicable to the matrix $D$ that skips the connection
between input and output in SSM. After selective filtering by SSM, matrix $D$
reassigns discrete effective information with continuous positional context.
However, the local positional information provided by the one-dimensional
convolution has limited impact on building dependencies between any two
elements in a sequence of any length in Attention. Therefore, when only using
Attention RoPE, although the improvement in perplexity is not significant
compared to not using RoPE, there is still an improvement. We believe that the
slight negative reason for the improvement is that Attention accounts for a
small proportion of the entire model, while the positive reason is that the
final Expresser module combines Attention with MLP. The negative contrast of
this observation is also reflected in the recent large-scale study on hybrid
architectures [18], which caused a negative impact on a sequence length of
16K.

When only using SSM RoPE, we are actually injecting continuous global relative
positional information into the SSM state. Even after selective filtering,
this information is still stored in the SSM state in discrete positional form,
affecting the linear recursive $h_{t}$ to consider the discrete relative
position of the previous effective information state. In the subsequent
Attention calculation, it still relies on the local positional information
provided by the one-dimensional convolution in the hidden dimension.
Nevertheless, when only using SSM RoPE, the perplexity performance decreases
by nearly $10\%$ compared to not using RoPE, significantly better than not
using RoPE. This result indicates that performing position calculation and
update in the SSM state is far superior to the method of providing local
positional information by convolution first and then reassigning continuous
positional information by matrix $D$.

When using both Attention and SSM RoPE, we can utilize the global relative
positional information in the SSM state. After selective filtering by SSM,
discrete effective information regains continuous positional information in
Attention RoPE. Through Attention calculation, the model can not only consider
the relative position between effective information states and ineffective
information states but also build dependencies between any two elements in a
sequence of any length and any element. This is why when using both Attention
and SSM RoPE, the perplexity continues to decrease by nearly $7\%$ compared to
only using SSM RoPE. Moreover, it is worth noting that in this case, we can
completely remove the one-dimensional convolution layer in the Mamba structure
and the matrix $D$ in SSM, saving much more computation than the additional
computation required for RoPE.

###  4.4 Cross-Domain can Improve the Efficiency of Fine-tuning

In this section, we will conduct a combination ablation experiment of MOE and
MLP in the OTCE architecture, and evaluate the zero-shot and five-shot
accuracy of shared expert isolation mixed expert models, cohesive cross domain
mixed expert models, and expansive cross domain mixed expert models in pre-
training, fine-tuning, and cross-domain multiple-choice questions.

The parameters of the MOE and MLP combination ablation experiment are set as
shown in Table 4, and the results are shown in Table 5.

Table 4:  (The Parameters of MOE and MLP Combination Ablation Experiment.) Because the ablation experiment is very time-consuming, we set a small parameter size to facilitate the ablation experiment with limited computing resources.  Layers | Model Dim | Attn Heads | Group | State Dim | Experts | Experts topK
---|---|---|---|---|---|---
8 | 1024 | 16 | 2 | 16 | 4 | 2
Table 5:  (The Results of MOE and MLP Combination Ablation Experiment.) We temporarily use ordinary MOE 15 17, adjust the dimension size of $FFN$ to ensure that any combination has the same total parameters, test several representative MOE and MLP combinations, iterate the same number of steps on the pre-training subset, and evaluate the performance of different combinations through perplexity.  Hybrid (S=SSM, A=Attention, M=MLP, E=MOE) | MOE Ratio | ppl $\downarrow$
---|---|---
SMSMSMSMAMSMSMAM | 0% | 4.86
SESMSMSMAMSMSMAM | 6.25% | 4.83
SMSMSMSMAESMSMAM | 6.25% | 4.84
SMSMSMSMAMSESMAM | 6.25% | 4.82
SMSMSMSMAMSMSEAM | 6.25% | 4.88
SMSMSMSMAMSMSMAE | 6.25% | 4.92
SESMSMSMAMSMSMAE | 12.5% | 4.68
SMSMSMSEAMSESMAM | 12.5% | 4.68
SMSMSESEAMSMSMAM | 12.5% | 4.67
SMSMSMSMAMSESEAM | 12.5% | 4.68
SMSESESEAMSMSMAM | 18.75% | 4.77
SESMSMSMAMSMSEAE | 18.75% | 4.69
SMSMSMSEAESESMAM | 18.75% | 4.63
SMSESMSEAMSESMSE | 25% | 4.61
SMSESESEAESESMAM | 31.25% | 4.48
SMSESESEAESESEAM | 37.5% | 4.56
SESESESEAESESEAE | 50% | 4.54

In most cases, the perplexity of different combinations decreases with the
increase of the MOE ratio, and the training cost of MOE is lower than that of
MLP with the same total parameters, so we choose the combination with a higher
MOE ratio as much as possible. We found that the perplexity is the lowest when
the first layer and the last layer (because of the special nature of the
Expresser module, it should be regarded as the last two layers) of the model
use MLP, and we speculate that using dense activation parameters in the two
parts of the model that process the input and prepare the output can help
improve stability. So we choose the first layer of the Observer module as MLP,
the last layer of the Conceiver module as MLP, and the last combination of the
Expresser module as MLP, as shown in Figure 1.

Table 6:  (The Parameters of Three Mixed Expert Models with Shared
Parameters.) We mainly increase the number of model layers to improve the data
saturation limit of the model, and abbreviate Shared Expert Isolation Mixture
of Experts as SEIMOE, Cohesive Cross Domain Mixture of Experts as CCDMOE,
Expansive Cross Domain Mixture of Experts as ECDMOE. In addition, because the
shared parameter ratio of the cohesive cross domain expert is only related to
the number of experts and cannot be adjusted independently, we first determine
the total number of parameters and the shared parameter ratio between experts
of the cohesive cross domain mixed expert model, and then adjust the size of
the $FFN$ of the shared expert isolation mixed expert model and the expansive
cross domain mixed expert model to ensure that the total number of parameters
and the shared parameter ratio between experts of the three models are the
same.

Layers | Model Dim | Attn Heads | Group | State Dim | Hybrid
---|---|---|---|---|---
|  |  |  |  | s=ssm a=attention m=mlp e=moe
24 | 1024 | 16 | 2 | 64 | $\makebox[16.00003pt]{{SM}}+\makebox[16.00003pt]{{SE}}\times 11+\makebox[16.00003pt]{{AE}}\times 2+\makebox[16.00003pt]{{SE}}\times 8+\makebox[16.00003pt]{{SM}}+\makebox[16.00003pt]{{AM}}$

Model | Parameters | Experts | Experts topK | Shared parameter ratio | PreTrain | FineTune
---|---|---|---|---|---|---
|  |  |  |  | ppl $\downarrow$ | ppl $\downarrow$
SEIMOE | 1B | 4 | 2 | 1:8 | 1.99 | 0.89
CCDMOE | 1B | 4 | 2 | 1:8 | 1.93 | 0.82
ECDMOE | 1B | 4 | 2 | 1:8 | 1.88 | 0.68

We list the basic parameter settings of the shared expert isolation mixed
expert model, cohesive cross domain mixed expert model, and expansive cross
domain mixed expert model in Table 6, the hybrid of SSM, Attention, MLP, and
MOE, the total number of parameters, the number of experts, the number of
expert activation, and the shared parameter ratio between experts are the
same.

For sparse activation model architectures, an important criterion for
evaluating whether they are excellent is the data efficiency of training, as
shown in Table 6, the comparison of pre-training perplexity and fine-tuning
perplexity of SEIMOE, CCDMOE, and ECDMOE, when we know that our training data
volume cannot train the model to the data saturation limit, we set the total
number of parameters and the shared parameter ratio to be the same, so we can
see the training data efficiency of the model from the perplexity. The pre-
training perplexity of CCDMOE is slightly better than SEIMOE, and the pre-
training perplexity of ECDMOE is slightly better than CCDMOE, but in the fine-
tuning perplexity performance, ECDMOE is far better than SEIMOE and CCDMOE,
with an increase of $23.59\%$ and $17.07\%$, respectively.

In Table 7, we detail the zero-shot and five-shot scores of SEIMOE, CCDMOE,
and ECDMOE in each domain. The average scores of CCDMOE and ECDMOE on zero-
shot and five-shot are better than SEIMOE, and the average scores of ECDMOE on
zero-shot and five-shot are better than CCDMOE. It is worth noting that the
average scores of SEIMOE and CCDMOE after Five-shot Fine-tune are lower than
the zero-shot scores, while ECDMOE performs well, with an increase of $2\%$,
which is consistent with the performance of the three models in fine-tuning
perplexity, indicating that ECDMOE is more efficient in training data than
SEIMOE and CCDMOE.

Table 7:  (The Zero-shot and Five-shot Evaluations of Three Mixed Expert Models with Shared Parameters.) These tasks come from CEvalBenchmark [6], we detail the zero-shot and five-shot accuracy of three mixed expert models with shared parameters in each subtask.  Task | SEIMOE | CCDMOE | ECDMOE | SEIMOE | CCDMOE | ECDMOE
---|---|---|---|---|---|---
| zero-shot $\uparrow$ | zero-shot $\uparrow$ | zero-shot $\uparrow$ | five-shot $\uparrow$ | five-shot $\uparrow$ | five-shot $\uparrow$
computer network | 42.10 | 52.63 | 53.51 | 47.37 | 52.63 | 54.87
operating system | 63.16 | 52.63 | 60.11 | 47.37 | 47.36 | 51.41
computer architecture | 61.90 | 52.38 | 54.21 | 61.90 | 52.38 | 52.09
college programming | 27.03 | 40.54 | 36.36 | 35.14 | 32.43 | 36.36
college physics | 36.84 | 36.84 | 37.65 | 42.11 | 42.11 | 43.15
college chemistry | 33.33 | 45.83 | 44.37 | 37.50 | 54.17 | 61.34
advanced mathematics | 36.84 | 47.37 | 43.76 | 42.11 | 57.89 | 47.76
probability and statistics | 44.44 | 44.44 | 45.66 | 44.44 | 38.89 | 41.21
discrete mathematics | 56.25 | 43.75 | 44.44 | 50.00 | 37.50 | 39.57
electrical engineer | 48.65 | 43.24 | 42.55 | 40.54 | 48.65 | 51.34
metrology engineer | 41.67 | 37.50 | 36.45 | 41.67 | 41.67 | 39.98
high school mathematics | 50.00 | 38.89 | 44.44 | 38.89 | 33.33 | 44.44
high school physics | 47.37 | 57.89 | 55.88 | 42.11 | 47.37 | 50.31
high school chemistry | 47.37 | 52.63 | 54.54 | 47.37 | 52.63 | 54.54
high school biology | 63.16 | 47.37 | 48.76 | 52.63 | 57.89 | 56.54
middle school mathematics | 63.16 | 47.37 | 65.76 | 63.16 | 42.11 | 65.76
middle school biology | 38.10 | 38.10 | 38.10 | 38.10 | 38.10 | 38.10
middle school physics | 36.84 | 57.89 | 59.87 | 31.58 | 47.37 | 48.76
middle school chemistry | 60.00 | 45.00 | 45.00 | 60.00 | 35.00 | 43.32
veterinary medicine | 52.17 | 43.48 | 51.11 | 47.83 | 39.13 | 51.78
college economics | 29.09 | 36.36 | 34.87 | 27.27 | 29.09 | 31.23
business administration | 27.27 | 36.36 | 34.44 | 30.30 | 36.36 | 41.21
marxism | 42.10 | 52.63 | 54.67 | 42.11 | 57.89 | 57.43
mao zedong thought | 29.17 | 25.00 | 25.00 | 33.33 | 25.00 | 24.43
education science | 31.03 | 20.69 | 33.43 | 31.03 | 27.59 | 34.32
teacher qualification | 22.73 | 27.27 | 29.98 | 27.27 | 31.82 | 32.77
high school politics | 47.37 | 52.63 | 65.47 | 63.16 | 57.89 | 68.98
high school geography | 47.37 | 52.63 | 53.21 | 42.11 | 52.63 | 53.21
middle school politics | 38.10 | 38.10 | 38.10 | 38.10 | 42.86 | 37.81
middle school geography | 50.00 | 75.00 | 71.10 | 50.00 | 58.33 | 45.78
modern chinese history | 34.78 | 47.83 | 46.43 | 39.13 | 56.52 | 62.12
ideological and moral cultivation | 57.89 | 42.11 | 45.55 | 52.63 | 31.58 | 42.21
logic | 54.55 | 40.91 | 54.55 | 54.55 | 36.36 | 54.55
law | 37.50 | 45.83 | 47.76 | 45.83 | 54.17 | 52.32
chinese language and literature | 43.48 | 56.52 | 55.55 | 43.48 | 56.52 | 55.55
art studies | 39.39 | 27.27 | 31.21 | 30.30 | 30.30 | 34.44
professional tour guide | 41.38 | 51.72 | 49.87 | 37.93 | 41.38 | 45.78
legal professional | 47.83 | 52.17 | 53.43 | 47.83 | 56.52 | 56.66
high school chinese | 42.10 | 47.37 | 45.56 | 36.84 | 47.37 | 48.64
high school history | 50.00 | 30.00 | 45.00 | 50.00 | 45.00 | 55.00
middle school history | 40.91 | 50.00 | 49.94 | 40.91 | 36.36 | 45.77
civil servant | 25.53 | 31.91 | 30.76 | 25.53 | 36.17 | 34.77
sports science | 47.37 | 42.11 | 45.54 | 57.89 | 63.16 | 70.12
plant protection | 40.91 | 50.00 | 45.88 | 45.45 | 45.45 | 45.88
basic medicine | 57.89 | 63.16 | 65.67 | 52.63 | 57.89 | 61.32
clinical medicine | 50.00 | 50.00 | 50.00 | 40.91 | 40.91 | 40.91
urban and rural planner | 41.30 | 45.65 | 41.30 | 41.30 | 32.61 | 41.30
accountant | 28.57 | 36.73 | 32.44 | 22.45 | 38.78 | 34.34
fire engineer | 35.48 | 32.26 | 30.32 | 29.03 | 45.16 | 37.75
environmental impact assessment | 45.16 | 41.94 | 46.56 | 45.16 | 32.26 | 46.56
tax accountant | 34.69 | 30.61 | 34.69 | 32.65 | 32.65 | 32.65
physician | 22.45 | 34.69 | 35.32 | 28.57 | 38.78 | 38.78
Average | 42.96 | 44.10 | 45.88 | 42.26 | 43.73 | 46.86

###  4.5 Associative Recall

In this section, we trained a 1B parameter Jamba model and a 1B parameter OTCE
model on 80% of the pre-training data set with a sequence length of 4096, and
then continued to pre-train the remaining 20% of the pre-training data set as
a 128K context length data set to continue training the model. The fine-tuning
data set is also packaged as a 128K context length data set to fine-tune the
model. The two models are extended to 128K context length for evaluation. The
parameter settings of the two models are shown in Table 8.

Table 8:  (Model Parameters for Associative Recall.) The specific
implementation of Jamba [12] is exactly the same as the original paper,
without using the improved modules in this paper. Also, because the proportion
of mixed experts is different, we adjust the dimension of $FFN$ in Jamba so
that the total number of parameters is the same as OTCE.

Model | Layers | Model Dim | Attn Heads | Group | State Dim | Experts | Experts topK | Experts Type
---|---|---|---|---|---|---|---|---
Jamba | 24 | 1024 | 16 | 2 | 64 | 4 | 1 | -
OTCE | 24 | 1024 | 16 | 2 | 64 | 4 | 1 | Expansive

Model | Hybrid
---|---
| s=ssm a=attention m=mlp e=moe
Jamba | $(\makebox[16.00003pt]{{SM}}+\makebox[16.00003pt]{{SE}}+\makebox[16.00003pt]{{SM}}+\makebox[16.00003pt]{{SE}}+\makebox[16.00003pt]{{AM}}+\makebox[16.00003pt]{{SE}}+\makebox[16.00003pt]{{SM}}+\makebox[16.00003pt]{{SE}})\times 3$
OTCE | $\makebox[16.00003pt]{{SM}}+\makebox[16.00003pt]{{SE}}\times 11+\makebox[16.00003pt]{{AE}}\times 2+\makebox[16.00003pt]{{SE}}\times 8+\makebox[16.00003pt]{{SM}}+\makebox[16.00003pt]{{AM}}$

Table 9:  (Evaluations for Associative Recall.) We select the evaluations for associative recall including: standard short context task: PIQA [1], natural long context task: NarrativeQA [11], synthetic long context task: HotpotQA [20], to evaluate the performance of Jamba and OTCE on these three tasks.  Model | PIQA | NarrativeQA | HotpotQA | Average
---|---|---|---|---
| acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$ | acc $\uparrow$
Jamba | 78.65 | 23.15 | 34.40 | 45.4
OTCE | 79.48 | 27.98 | 42.71 | 50.0

In Table 9, we show the performance of Jamba and OTCE on standard short
context tasks, natural long context tasks, and synthetic long context tasks.
It is known that Jamba [12] can be successfully trained on tasks with a
context length of 1M, not only performing well on tasks with a context length
of 256K, but also surpassing many well-known open-source models such as Llama
[16], Gemma [15], Mixtral [8] in common tasks. Therefore, we believe that
Jamba is an excellent baseline model, and comparing with it is equivalent to
comparing with many well-known open-source models.

In standard short context tasks, it is obvious that the model does not need to
have the ability to handle long-distance dependencies. Our goal is to evaluate
whether the accuracy of short-distance associative recall will be affected
when the OTCE model is extended to handle long context environments in common
tasks. The experimental results show that OTCE outperforms Jamba in standard
short context tasks. Therefore, while adapting to long context environments,
OTCE does not sacrifice accuracy in common tasks.

In natural long context tasks, OTCE achieved a performance improvement of up
to $20.86\%$ compared to Jamba. We analyze that the main reason for this
improvement is that SSM has a learnable matrix $(\Delta,B,C)$ during training.
However, in the inference stage, although these matrices can be dynamically
calculated according to different input data, the parameters used for these
calculations (i.e., the function or mapping that determines how to calculate
these matrices and step sizes) are fixed, and these parameters are reused in
the inference stage after being determined in the training stage. This means
that the state of SSM may still be affected by irrelevant documents in the
inference stage. In contrast, the last layer of Jamba consists of SSM and MOE.
In the inference stage, not only the output state may be confused when
approaching the output, but also as the ablation experiment of MOE and MLP in
the previous section 5 shows, the output state may also be unstable. These
factors together lead to Jamba’s performance in handling long context tasks is
not as good as OTCE.

We all know that the Transformer model tends to copy the answers in the
context examples rather than predict the answers to the actual questions. In
contrast, the SSM model tends to predict answers rather than directly copy
answers from the context. In synthetic long context tasks, the model needs to
recall information from lengthy irrelevant text, requiring the ability to
track across contexts and aggregate information. The reason why OTCE
outperforms Jamba by $27.06\%$ in synthetic long context tasks is similar to
the attention decay and knowledge routing bias issues discussed in the
previous section 4.2. The Expresser module of OTCE reweights the recursively
aggregated information with gradually decaying long-term dependency
relationships through re-attention to all elements, and learns how SSM
recursively aggregates information to correct knowledge routing bias. This
allows OTCE to outperform Jamba in synthetic long context tasks.

In the Needle in a haystack problem, the model is required to recall
information from lengthy irrelevant text input, as shown in Figure 10. OTCE
has excellent performance on the Needle in a haystack evaluation.

Figure 10:  (Needle in a haystack.) OTCE can recall information across 128K
context length (our maximum extended sequence length), and has excellent
performance on the Needle in a haystack evaluation.

##  5 Conclusion

We propose a position information encoding method that connects SSM and
Attention, a sparse cross-domain mixture of experts with higher learning
efficiency, and design a hybrid SSM and Attention sparse model architecture:
OTCE with a biomimetic idea, achieving state-of-the-art performance in
language modeling.

## References

  * [1] Yonatan Bisk, Rowan Zellers, Jianfeng Gao and Yejin Choi  “PIQA: Reasoning about Physical Commonsense in Natural Language”  In _Proceedings of the AAAI conference on Artificial Intelligence_ 34, 2020
  * [2] Damai Dai, Chengqi Deng, Chenggang Zhao, R.. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y.. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui and Wenfeng Liang  “DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models”  In _CoRR_ abs/2401.06066, 2024  URL: https://arxiv.org/abs/2401.06066
  * [3] Tri Dao and Albert Gu  “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality”  In _International Conference on Machine Learning (ICML)_ , 2024
  * [4] Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang and Jie Tang  “GLM: General Language Model Pretraining with Autoregressive Blank Infilling”  In _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , 2022, pp. 320–335
  * [5] Albert Gu and Tri Dao  “Mamba: Linear-Time Sequence Modeling with Selective State Spaces”  In _arXiv preprint arXiv:2312.00752_ , 2023
  * [6] Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Jiayi Lei, Yao Fu, Maosong Sun and Junxian He  “C-Eval: A Multi-Level Multi-Discipline Chinese Evaluation Suite for Foundation Models”  In _Advances in Neural Information Processing Systems_ , 2023
  * [7] Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample and Lucile Saulnier  “Mistral 7B”  In _arXiv preprint arXiv:2310.06825_ , 2023
  * [8] Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna and Florian Bressand  “Mixtral of Experts”  In _arXiv preprint arXiv:2401.04088_ , 2024
  * [9] Su Jianlin  “The Road to Transformer Upgrade: 2. Rotational Position Encoding”, 2021  URL: https://kexue.fm/archives/8265
  * [10] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas and François Fleuret  “Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention”  In _International Conference on Machine Learning_ , 2020, pp. 5156–5165  PMLR
  * [11] Tomas Kovcisky, Bryan McCann, James Bradbury, Caiming Xiong and Richard Socher  “The NarrativeQA Reading Comprehension Challenge”  In _Transactions of the Association for Computational Linguistics_ 6, 2018, pp. 317–328
  * [12] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov and Shai Shalev-Shwartz  “Jamba: A Hybrid Transformer-Mamba Language Model”  In _arXiv preprint arXiv:2403.19887_ , 2024
  * [13] Meta NVIDIA  “PyTorch Container Image”, https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch, 2022
  * [14] Noam Shazeer  “GLU Variants Improve Transformer”  In _arXiv preprint arXiv:2002.05202_ , 2020
  * [15] Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale and Juliette Love  “Gemma: Open models based on gemini research and technology”  In _arXiv preprint arXiv:2403.08295_ , 2024
  * [16] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava and Shruti Bhosale  “Llama 2: Open foundation and fine-tuned chat models”  In _arXiv preprint arXiv:2307.09288_ , 2023
  * [17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin  “Attention is All You Need”  In _arXiv preprint arXiv:1706.03762_ , 2017
  * [18] Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi and Bryan Catanzaro  “An Empirical Study of Mamba-based Language Models”  In _arXiv preprint arXiv:2406.07887_ , 2024
  * [19] Liang Xu, Hai Hu, Xuanwei Zhang, Lu Li, Chenjie Cao, Yudong Li, Yechen Xu, Kai Sun, Dian Yu, Cong Yu, Yin Tian, Qianqian Dong, Weitang Liu, Bo Shi, Yiming Cui, Junyi Li, Jun Zeng, Rongzhao Wang, Weijian Xie, Yanting Li, Yina Patterson, Zuoyu Tian, Yiwen Zhang, He Zhou, Shaoweihua Liu, Zhe Zhao, Qipeng Zhao, Cong Yue, Xinrui Zhang, Zhengliang Yang, Kyle Richardson and Zhenzhong Lan  “CLUE: A Chinese Language Understanding Evaluation Benchmark”  In _Proceedings of the 28th International Conference on Computational Linguistics_ Barcelona, Spain (Online): International Committee on Computational Linguistics, 2020, pp. 4762–4772  DOI: 10.18653/v1/2020.coling-main.419
  * [20] Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W Cohen, Ruslan Salakhutdinov and Christopher D Manning  “HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering”  In _arXiv preprint arXiv:1809.09600_ , 2018

##  Appendix A RoPE for SSM

###### Proof of Equation 11.

The rotation position encoding for $C$ and $B$ matrices in SSM.

$\displaystyle C_{j}$ $\displaystyle=f_{C}(x_{j},j)$ $\displaystyle B_{i}$
$\displaystyle=f_{B}(x_{i},i)$

$C_{j}$ represents the output weight matrix of the $j$-th token corresponding
to the word vector $x_{j}$ integrated with the position information $j$,
$B_{i}$ represents the input weight matrix of the $i$-th token corresponding
to the word vector $x_{i}$ integrated with the position information $i$.

To utilize the relative positional information between tokens, we assume that
the inner product operation between the $C_{j}$ vector and the $B_{i}$ vector
can be represented by a function $g$, where the input of the function $g$ is
the word embedding vectors $x_{j}$ and $x_{i}$, and their relative positional
information $j-i$, the inner product of $C_{j}$ and $B_{i}$ and their relative
positional information $j-i$ is defined as

$<f_{C}(x_{j},j),f_{B}(x_{i},i)>=g(x_{j},x_{i},j-i)$

Now, assuming the word embedding vector dimension is $d=2$, we have
$f_{C}(x_{j},j)=(W_{C}x_{j})e^{\imath j\theta}$, for the first half of the
formula $W_{C}x_{j}$, we know that $W_{C}$ is a two-dimensional matrix,
$x_{j}$ is a two-dimensional vector, the result of the multiplication is
naturally a two-dimensional vector, represented by $C_{j}$

$\displaystyle C_{j}$ $\displaystyle=\begin{bmatrix}C_{j}^{(1)}\\\
C_{j}^{(2)}\end{bmatrix}=W_{C}x_{j}=\begin{bmatrix}W_{C}^{(11)}&W_{C}^{(12)}\\\
W_{C}^{(21)}&W_{C}^{(22)}\end{bmatrix}\begin{bmatrix}x_{j}^{(1)}\\\
x_{j}^{(2)}\end{bmatrix}$

For the second half $e^{\imath j\theta}$, according to Euler’s formula
$e^{\imath x}=\cos(x)+\imath\sin(x)$, we have

$\displaystyle e^{\imath j\theta}$
$\displaystyle=\cos(j\theta)+\imath\sin(j\theta)$

We know

$\displaystyle f_{C}(x_{j},j)$ $\displaystyle=(W_{C}x_{j})e^{\imath
j\theta}=C_{j}e^{\imath j\theta}$

$C_{j}$ is represented in complex form,

$\displaystyle C_{j}$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)},C_{j}^{(2)}\end{bmatrix}=\begin{bmatrix}C_{j}^{(1)}+\imath
C_{j}^{(2)}\end{bmatrix}$

Thus,

$\displaystyle f_{C}(x_{j},j)$ $\displaystyle=C_{j}e^{\imath
j\theta}=\begin{bmatrix}C_{j}^{(1)}+\imath C_{j}^{(2)}\end{bmatrix}e^{\imath
j\theta}$

According to the above derivation, we know that $f_{C}(x_{j},j)$ is the
product of two complex numbers,

$\displaystyle f_{C}(x_{j},j)$ $\displaystyle=C_{j}e^{\imath
j\theta}=\begin{bmatrix}C_{j}^{(1)}+\imath
C_{j}^{(2)}\end{bmatrix}\times(\cos(j\theta)+\imath\sin(j\theta))$

Considering the following two formulas about complex numbers

$\displaystyle(a+\imath b)\times(c+\imath d)$ $\displaystyle=ac+\imath
bc+\imath ad+\imath^{2}bd=(ac-bd)+\imath(bc+ad)$ $\displaystyle\imath^{2}$
$\displaystyle=-1$

We have

$\displaystyle C_{j}e^{\imath j\theta}$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}+\imath
C_{j}^{(2)}\end{bmatrix}\times(\cos(j\theta)+\imath\sin(j\theta))=\begin{bmatrix}C_{j}^{(1)}\cos(j\theta)-C_{j}^{(2)}\sin(j\theta)\end{bmatrix}+\imath\begin{bmatrix}C_{j}^{(2)}\cos(j\theta)+C_{j}^{(1)}\sin(j\theta)\end{bmatrix}$

Expressing this result as a real vector,

$\displaystyle C_{j}e^{\imath j\theta}$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}\cos(j\theta)-C_{j}^{(2)}\sin(j\theta),C_{j}^{(2)}\cos(j\theta)+C_{j}^{(1)}\sin(j\theta)\end{bmatrix}$

Therefore, $C_{j}$ multiplied by a rotation matrix is obtained.

$\displaystyle f_{C}(x_{j},j)$ $\displaystyle=(W_{C}x_{j})e^{\imath
j\theta}=C_{j}e^{\imath j\theta}$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}\cos(j\theta)-C_{j}^{(2)}\sin(j\theta),C_{j}^{(2)}\cos(j\theta)+C_{j}^{(1)}\sin(j\theta)\end{bmatrix}$
$\displaystyle=\begin{bmatrix}\cos(j\theta)&-\sin(j\theta)\\\
\sin(j\theta)&\cos(j\theta)\end{bmatrix}\begin{bmatrix}C_{j}^{(1)}\\\
C_{j}^{(2)}\end{bmatrix}$

Similarly, $B_{i}$ vector can be obtained

$\displaystyle f_{B}(x_{i},i)$ $\displaystyle=(W_{B}x_{i})e^{\imath
i\theta}=B_{i}e^{\imath i\theta}$
$\displaystyle=\begin{bmatrix}B_{i}^{(1)}\cos(i\theta)-B_{i}^{(2)}\sin(i\theta),B_{i}^{(2)}\cos(i\theta)+B_{i}^{(1)}\sin(i\theta)\end{bmatrix}$
$\displaystyle=\begin{bmatrix}\cos(i\theta)&-\sin(i\theta)\\\
\sin(i\theta)&\cos(i\theta)\end{bmatrix}\begin{bmatrix}B_{i}^{(1)}\\\
B_{i}^{(2)}\end{bmatrix}$

The function $g$ can be represented as

$\displaystyle g(x_{j},x_{i},j-i)$
$\displaystyle=\Re\begin{bmatrix}(W_{C}x_{j})(W_{B}x_{i})^{*}e^{\imath(j-i)\theta}\end{bmatrix}$

where $\Re$ represents the real part of the complex number $x$,
$(W_{C}x_{j})(W_{B}x_{i})^{*}$ represents the conjugate of the product of two
complex numbers. Considering

$\displaystyle z$ $\displaystyle=a+\imath b$ $\displaystyle z^{*}$
$\displaystyle=a-\imath b$

we have

$\displaystyle W_{C}x_{j}$ $\displaystyle=C_{j}=C_{j}^{(1)}+\imath
C_{j}^{(2)}$ $\displaystyle W_{B}x_{i}$
$\displaystyle=B_{i}=B_{i}^{(1)}+\imath B_{i}^{(2)}$
$\displaystyle(W_{B}x_{i})^{*}$ $\displaystyle=B_{i}^{*}=B_{i}^{(1)}-\imath
B_{i}^{(2)}$ $\displaystyle e^{\imath(j-i)\theta}$
$\displaystyle=\cos((j-i)\theta)+\imath\sin((j-i)\theta)$

We now want to prove that

$\displaystyle g(x_{j},x_{i},j-i)$
$\displaystyle=\Re\begin{bmatrix}(W_{C}x_{j})(W_{B}x_{i})^{*}e^{\imath(j-i)\theta}\end{bmatrix}$
$\displaystyle=\Re\begin{bmatrix}(C_{j}^{(1)}+\imath
C_{j}^{(2)})(B_{i}^{(1)}-\imath
B_{i}^{(2)})(\cos((j-i)\theta)+\imath\sin((j-i)\theta))\end{bmatrix}$
$\displaystyle=\Re\begin{bmatrix}((C_{j}^{(1)}B_{i}^{(1)}+C_{j}^{(2)}B_{i}^{(2)})+\imath(C_{j}^{(2)}B_{i}^{(1)}-C_{j}^{(1)}B_{i}^{(2)}))(\cos((j-i)\theta)+\imath\sin((j-i)\theta))\end{bmatrix}$
$\displaystyle=(C_{j}^{(1)}B_{i}^{(1)}+C_{j}^{(2)}B_{i}^{(2)})\cos((j-i)\theta)-(C_{j}^{(2)}B_{i}^{(1)}-C_{j}^{(1)}B_{i}^{(2)})\sin((j-i)\theta)$

Recalling the vectorized form of SSM 8, the $C$ vector at position $j$ and the
$B$ vector at position $i$ will perform an inner product operation, that is,

$\displaystyle f_{C}(x_{j},j)$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}\cos(j\theta)-C_{j}^{(2)}\sin(j\theta),C_{j}^{(2)}\cos(j\theta)+C_{j}^{(1)}\sin(j\theta)\end{bmatrix}$
$\displaystyle f_{B}(x_{i},i)$
$\displaystyle=\begin{bmatrix}B_{i}^{(1)}\cos(i\theta)-B_{i}^{(2)}\sin(i\theta),B_{i}^{(2)}\cos(i\theta)+B_{i}^{(1)}\sin(i\theta)\end{bmatrix}$

We have

$\displaystyle<f_{C}(x_{j},j),f_{B}(x_{i},i)>$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}\cos(j\theta)-C_{j}^{(2)}\sin(j\theta)\end{bmatrix}\begin{bmatrix}B_{i}^{(1)}\cos(i\theta)-B_{i}^{(2)}\sin(i\theta)\end{bmatrix}$
$\displaystyle+\begin{bmatrix}C_{j}^{(2)}\cos(j\theta)+C_{j}^{(1)}\sin(j\theta)\end{bmatrix}\begin{bmatrix}B_{i}^{(2)}\cos(i\theta)+B_{i}^{(1)}\sin(i\theta)\end{bmatrix}$
$\displaystyle=C_{j}^{(1)}\cos(j\theta)B_{i}^{(1)}\cos(i\theta)-C_{j}^{(1)}\cos(j\theta)B_{i}^{(2)}\sin(i\theta)$
$\displaystyle-
C_{j}^{(2)}\sin(j\theta)B_{i}^{(1)}\cos(i\theta)+C_{j}^{(2)}\sin(j\theta)B_{i}^{(2)}\sin(i\theta)$
$\displaystyle+C_{j}^{(2)}\cos(j\theta)B_{i}^{(2)}\cos(i\theta)+C_{j}^{(2)}\cos(j\theta)B_{i}^{(1)}\sin(i\theta)$
$\displaystyle+C_{j}^{(1)}\sin(j\theta)B_{i}^{(2)}\cos(i\theta)+C_{j}^{(1)}\sin(j\theta)B_{i}^{(1)}\sin(i\theta)$

Considering

$\displaystyle\sin(a+b)$ $\displaystyle=\sin(a)\cos(b)+\cos(a)\sin(b)$
$\displaystyle\sin(a-b)$ $\displaystyle=\sin(a)\cos(b)-\cos(a)\sin(b)$
$\displaystyle\cos(a+b)$ $\displaystyle=\cos(a)\cos(b)-\sin(a)\sin(b)$
$\displaystyle\cos(a-b)$ $\displaystyle=\cos(a)\cos(b)+\sin(a)\sin(b)$

We have

$\displaystyle<f_{C}(x_{j},j),f_{B}(x_{i},i)>$
$\displaystyle=C_{j}^{(1)}B_{i}^{(1)}(\cos(j\theta)\cos(i\theta)+\sin(j\theta)\sin(i\theta))$
$\displaystyle+C_{j}^{(1)}B_{i}^{(2)}(-\cos(j\theta)\sin(i\theta)+\sin(j\theta)\cos(i\theta))$
$\displaystyle+C_{j}^{(2)}B_{i}^{(1)}(-\sin(j\theta)\cos(i\theta)+\cos(j\theta)\sin(i\theta))$
$\displaystyle+C_{j}^{(2)}B_{i}^{(2)}(\sin(j\theta)\sin(i\theta)+\cos(j\theta)\cos(i\theta))$
$\displaystyle=C_{j}^{(1)}B_{i}^{(1)}\cos((j-i)\theta)+C_{j}^{(1)}B_{i}^{(2)}\sin((j-i)\theta)$
$\displaystyle-
C_{j}^{(2)}B_{i}^{(1)}\sin((j-i)\theta)+C_{j}^{(2)}B_{i}^{(2)}\cos((j-i)\theta)$
$\displaystyle=(C_{j}^{(1)}B_{i}^{(1)}+C_{j}^{(2)}B_{i}^{(2)})\cos((j-i)\theta)+(C_{j}^{(1)}B_{i}^{(2)}-C_{j}^{(2)}B_{i}^{(1)})\sin((j-i)\theta)$
$\displaystyle=(C_{j}^{(1)}B_{i}^{(1)}+C_{j}^{(2)}B_{i}^{(2)})\cos((j-i)\theta)-(C_{j}^{(2)}B_{i}^{(1)}-C_{j}^{(1)}B_{i}^{(2)})\sin((j-i)\theta)$
$\displaystyle=g(x_{j},x_{i},j-i)$

It is proved that the inner product of the $C$ vector at position $j$ and the
$B$ vector at position $i$ is the function $g$.

Finally, using the matrix-vector multiplication form

$\displaystyle<f_{C}(x_{j},j),f_{B}(x_{i},i)>$
$\displaystyle=\begin{bmatrix}\begin{bmatrix}\cos(j\theta)&-\sin(j\theta)\\\
\sin(j\theta)&\cos(j\theta)\end{bmatrix}\begin{bmatrix}C_{j}^{(1)}\\\
C_{j}^{(2)}\end{bmatrix}\end{bmatrix}^{T}\begin{bmatrix}\begin{bmatrix}\cos(i\theta)&-\sin(i\theta)\\\
\sin(i\theta)&\cos(i\theta)\end{bmatrix}\begin{bmatrix}B_{i}^{(1)}\\\
B_{i}^{(2)}\end{bmatrix}\end{bmatrix}$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}&C_{j}^{(2)}\end{bmatrix}\begin{bmatrix}\cos(j\theta)&\sin(j\theta)\\\
-\sin(j\theta)&\cos(j\theta)\end{bmatrix}\begin{bmatrix}\cos(i\theta)&-\sin(i\theta)\\\
\sin(i\theta)&\cos(i\theta)\end{bmatrix}\begin{bmatrix}B_{i}^{(1)}\\\
B_{i}^{(2)}\end{bmatrix}$

Expanding the product of the two rotation matrices, we have

$\displaystyle\begin{bmatrix}\cos(j\theta)\cos(i\theta)+\sin(j\theta)\sin(i\theta)&-\cos(j\theta)\sin(i\theta)+\sin(j\theta)\cos(i\theta)\\\
-\sin(j\theta)\cos(i\theta)+\cos(j\theta)\sin(i\theta)&\sin(j\theta)\sin(i\theta)+\cos(j\theta)\cos(i\theta)\end{bmatrix}$

Finally, we get

$\displaystyle<f_{C}(x_{j},j),f_{B}(x_{i},i)>$
$\displaystyle=\begin{bmatrix}C_{j}^{(1)}&C_{j}^{(2)}\end{bmatrix}\begin{bmatrix}\cos((j-i)\theta)&-\sin((j-i)\theta)\\\
\sin((j-i)\theta)&\cos((j-i)\theta)\end{bmatrix}\begin{bmatrix}B_{i}^{(1)}\\\
B_{i}^{(2)}\end{bmatrix}$

The above derivation is only for the case of word embedding dimension $d=2$,
when $d>2$, the two-dimensional case can be extended to any dimension as
follows

$\displaystyle f_{\\{C,B\\}}(x_{j},j)$
$\displaystyle=\mathbb{R}_{\Theta,j}^{d}W_{\\{C,B\\}}x_{j}$

The inner product satisfies linearity, so for any even-dimensional RoPE, we
can represent it as a concatenation of the two-dimensional case, that is,
grouping the elements of the word embedding vector in pairs

$\displaystyle\mathbb{R}_{\Theta,j}^{d}=\begin{bmatrix}\cos
j\theta_{0}&-sinj\theta_{0}&0&0&\dots&0&0\\\ \sin j\theta_{0}&\cos
j\theta_{0}&0&0&\dots&0&0\\\ 0&0&\cos j\theta_{1}&-sinj\theta_{1}&\dots&0&0\\\
0&0&\sin j\theta_{1}&\cos j\theta_{1}&\dots&0&0\\\
\vdots&\vdots&\vdots&\vdots&\ddots&\vdots&\vdots\\\ 0&0&0&0&\dots&\cos
j\theta_{d/2}&-sinj\theta_{d/2-1}\\\ 0&0&0&0&\dots&\sin j\theta_{d/2}&\cos
j\theta_{d/2-1}\\\ \end{bmatrix}$

Each group applies the same rotation operation and the rotation angle of each
group is calculated as follows:

$\displaystyle\Theta$
$\displaystyle=\\{\theta_{i}=10000^{-2(i-1)/d},i\in[1,2,\dots,d/2]\\}$

RoPE is a kind of relative positional encoding, and relative positional
encoding is a special form of Toeplitz matrix

$\displaystyle\begin{bmatrix}0&&&&&&\\\ 1&0&&&&\\\ 2&1&0&&\\\ 3&2&1&0\\\
\vdots&\ddots&\ddots&\ddots&\ddots\\\ n-1&n-2&n-3&\dots&1&0\end{bmatrix}$

We can know that the position distribution after RoPE is unbalanced, 0 appears
most frequently as the low bit, and n-1 appears least frequently as the high
bit, which also leads to a problem of RoPE, its high bit is not sufficiently
trained, and the generalization ability is not as good as the low bit. We take
the average value of the effective sequence length of the training data as the
base, for the sequence length greater than the base, we have

$\displaystyle C\times\max(1,\log_{base}n)$

The part of the sequence length within the base is not affected, and the part
of the sequence length greater than the base is expanded according to the
ratio of $\log_{base}n$, so that the problem of insufficient generalization
ability of the high bit can be solved.

∎