Joint Fine-tuning and Conversion of Pretrained Speech and Language Models towards Linear Complexity

## Abstract

Abstract Architectures such as Linformer and Mamba have recently emerged as competitive linear time replacements for transformers. However, corresponding large pretrained models are often unavailable, especially in non-text domains. To remedy this, we present a Cross-Architecture Layerwise Distillation (CALD) approach that jointly converts a transformer model to a linear time substitute and fine-tunes it to a target task. We also compare several means to guide the fine-tuning to optimally retain the desired inference capability from the original model. The methods differ in their use of the target model and the trajectory of the parameters. In a series of empirical studies on language processing, language modeling, and speech processing, we show that CALD can effectively recover the result of the original model, and that the guiding strategy contributes to the result. Some reasons for the variation are suggested.

## 1 Introduction

Recent progress in language modeling from BERT (Devlin et al., ) to Llama3 (Dubey et al., ) has witnessed the
rise of the attention mechanism and transformer architecture. Such models are able to scale up to very large model capacity and pretraining data for natural language processing (NLP), and have lead to breakthroughs in a broad range of downstream applications under either fine-tuning or zero-shot/in-context scenarios. This is surely not restricted to the text domain: Off-the-shelf speech and audio models like Wav2Vec2 (Baevski et al., ), HuBERT (Hsu et al., ), and WavLM (Chen et al., ) are pretrained on extensive audio data; they can be fine-tuned to obtain good performance on various tasks from automatic speech recognition (ASR) to spoken language understanding (SLU). However, such large-scale models are computationally expensive; the main reason is that in the standard attention mechanism, each token interacts with all other tokens in the input sequence, forming a complete graph. The computation time therefore increases quadratically with the input length, and a linear-growing KV-cache is required. As a result, the computation requirements can become prohibitive, especially for long contexts.
A large cohort of architectures and algorithms have been proposed to mitigate the problem and to reach linear complexity, using techniques such as low-rank projection, hashing, and sparse attention (Tay et al., ). More recently, there has been a revival of recurrent networks such as state space models (Gu et al., ; Gu & Dao, ), RWKV (Peng et al., ), and RetNet (Sun et al., ). By recurrently updating states, those models allow linear time and constant space inference while retaining the transformer’s parallel training capability. It has been demonstrated that the latest recurrent models, notably Mamba2 (Dao & Gu, ), are capable of reaching performance comparative with the most well-trained transformer, even in large-scale models.

In spite of such transformer substitutes, there remain obstacles to putting them into practical use. These new models are not simply plug-and-play upgrades for existing pretrained models, but are generally pretrained from scratch again. The pretrained weights of these models are often not publicly released; the data and computational cost of repeating such pretraining for every new model can be prohibitive for users and researchers of downstream applications. In particular, most models are primarily developed for processing language or text. These models are typically general-purpose substitutes for the standard attention mechanism, also applicable to other domains such as speech and audio. However, the pretrained models on those latter domains simply never exist; this hinders those domains from utilizing the latest models and algorithms created in the larger machine learning (ML) and NLP community.
This is a particular concern for research in end-to-end speech processing involving long-distance dependency in the lower information density audio form: such models are a prerequisite to efficiently process the abundant long-form but poorly segmented audio in the wild.

To address this issue, a possible solution is to make use of off-the-shelf pretrained models and convert them into the target linear-complexity model. Several previous works have explored this method to obtain a target model with much less computation. Many of the standard transformer layers (e.g., embeddings and feed-forward) are still part of the target architecture. Compared to pretraining from scratch, an immediate improvement would be the reuse of those pretrained layers (Kasai et al., ; Choi, ). In addition to such parameter transfer, we can further guide the conversion by transferring model behavior using knowledge distillation (KD). Given model inputs, KD aligns the model outputs or hidden states between the original transformer (teacher) and the target model (student) using an extra loss term. In an earlier work, Tang et al. () explored conversion from BERT to RNN by aligning outputs; more recently, Zhang et al. () and Bick et al. () tried to reproduce the attention weights of the standard attention mechanism in linear attention models. However, it has been pointed out that reproducing attention weights is not a prerequisite to reproduce performance (Mao, ), and even causes training inefficiency and instability (Mercat et al., ). More importantly, this method is restricted to target models for which we can easily determine an attention weight matrix equivalent to that of the standard one. This is not always the case: e.g., the matrix in Linformer has a lower rank (Wang et al., ). Furthermore, distillation only on outputs may not be sufficient to guide the model.
The above works are restricted to NLP tasks without considering other more demanding domains, e.g., speech. Furthermore, most of these works are focused on conversion using an efficient re-pretraining stage, which is still difficult for downstream users with limited access to computational resources and pretraining data. Instead, it would be more attractive to directly carry out the conversion in the fine-tuning phase.

Given the considerations above, we explore the possibility of converting an existing pretrained transformer into a linear-complexity model on the downstream task. Our approach broadly covers different source models, target models, and target tasks. This is made possible by combining parameter transfer and layerwise distillation with hidden states from a target teacher model, i.e., a pretrained standard transformer fine-tuned for the target task. In addition to direct guidance from hidden states of the target teacher, we further investigate two modes of guidance:
Recognizing the importance of the pretrained knowledge of the original transformer that can be lost during fine-tuning, we can guide the converted model by the sequence of models on the teacher’s fine-tuning trajectory. Alternatively, we can use the distilled model as a good initialization for conversion, and allow deviation from the teacher’s behavior in the later training phase.

We demonstrate the advantage of our novel Cross-Architecture Layerwise Distillation (CALD) approach under a broad range of scenarios, including conversion from RoBERTa to Linformer for NLP tasks, from Pythia to Mamba for language modeling, and from Wav2Vec2 to Mamba2 for speech tasks including ASR, SLU, and speaker identification (SID). In all of these scenarios, our converted models demonstrate minor to no performance loss compared to the standard transformer, much better than naive parameter transfer. All our code and models will be publicly available.

## 2 Related work

### 2.1 Linear-complexity and recurrent models

Linear-complexity models. There is a large body of work on reducing the $O(N^{2})$ time complexity of transformer owing to the attention computation $A({\bm{Q}},{\bm{K}},{\bm{V}})=\text{softmax}(\frac{{\bm{Q}}{\bm{K}}^{T}}{\sqrt
{d}}){\bm{V}}$. The $N\times d$ matrices ${\bm{Q}}$, ${\bm{K}}$, and ${\bm{V}}$ represent the sequence of $N$ query, key, and value feature vectors, each with constant $d$ dimensions. A straightforward approach is to limit the attention to several pre-specified positions in the sequence, e.g., neighboring tokens, several special global tokens, or a sparse subset of the sequence (Child et al., ; Beltagy et al., ; Ainslie et al., ; Zaheer et al., ). Alternatively, the attended positions can be determined on-the-fly by hashing and clustering (Kitaev et al., ; Roy et al., ). With Linformer, Wang et al. () further found that the attention matrix ${\bm{A}}$ is inherently low-rank, and projects ${\bm{K}}$ and ${\bm{V}}$ into a lower $k$-dimensional space using $k\times N$ matrices ${\bm{E}}$ and ${\bm{F}}$. The two matrices can be shared, or even shared between all layers. Then, the $O(N^{2})$ matrix multiplication can be avoided, as ${\bm{A}}$ can be provably approximated by

$$ $\displaystyle\tilde{A}({\bm{Q}},{\bm{K}},{\bm{V}})=\text{softmax}(\frac{{\bm{Q }}({\bm{E}}{\bm{K}})^{T}}{\sqrt{d}}){\bm{F}}{\bm{V}}.$ (1) $$

Recurrent models. By applying a kernel trick in a reverse direction,
Katharopoulos et al. () reduce transformers to RNNs of infinite dimension. In this RNN, hidden states ${\bm{z}}_{t}=\sum_{i\leq t}\phi({\bm{K}}_{i,:})$ are the sum of an infinite-dimensional feature $\phi({\bm{K}}_{i,:})$ mapped from each key vector (more specifically, along with their product with values). Using various finite dimension approximations of $\phi$ (Peng et al., ; Choromanski et al., ; Zhai et al., ), the RNN can be implemented as an approximation of the standard attention, which paves the way for the later revival of recurrent models.
Such models are nevertheless different from classical RNNs in that the training can be parallel as long as there is no non-linear transform in the transition from ${\bm{z}}_{t}$ to ${\bm{z}}_{t+1}$. In addition to simply summing up features as in ${\bm{z}}_{t}$ above, RWKV (Peng et al., ) and RetNet (Sun et al., ) use exponential time decay of past features. Gating (Yang et al., ) or linear transform (Orvieto et al., ) have also been introduced. Similarly, a state-space model (SSM) is defined as a transformation of a 1-D input sequence $x(t)$ to output $y(t)$ via an $N$-D hidden state $h(t)$, following

$$ $\displaystyle h^{\prime}(t)$ $\displaystyle={\bm{A}}h(t)+{\bm{B}}x(t)$ (2) $\displaystyle y(t)$ $\displaystyle={\bm{C}}h(t).$ (3) $$

Hence the discretization of an SSM is akin to an RNN with linear state transition (Gu et al., ; ). Mamba then renders the transition matrices as ${\bm{A}}_{t}$, ${\bm{B}}_{t}$, and ${\bm{C}}_{t}$ dependent on the input $x(t)$ (Gu & Dao, ). Mamba2 (Dao & Gu, ) further allows a multi-head formulation in a hardware-efficient way under a generalized framework of various SSMs and linear attention models.

Linear-complexity models for speech. With the lower information density or longer form of spoken language, these efficient models have been also introduced to audio and speech processing. There have been attempts to perform ASR using Longformer (Rekesh et al., ; Koluguri et al., ) and RWKV (An & Zhang, ). Also, Mamba has been applied to speech separation (Jiang et al., ; Li & Chen, ) and enhancement (Li & Chen, ; Zhang et al., ). Zhang et al. () investigated different options of bidirectional Mamba for ASR. Furthermore, there have been attempts to build Mamba-based pretrained speech models (Shams et al., ; Yadav & Tan, ; Bhati et al., ). Such pretraining is nevertheless expensive, especially when taking in to account the need of utilizing the ever-growing set of linear-complexity models.

### 2.2 Distillation and model conversion

Knowledge distillation. Distilling knowledge from teacher models to a student one by matching the output probabilities is a standard approach to compress ML models, including pretrained language models (LMs) (Hinton et al., ; Sanh et al., ). Intermediate hidden states of the teacher and student can be further aligned by layerwise distillation (Sun et al., ; Aguilar et al., ). In addition, Brown et al. () have attempted to distill and compress existing pretrained linear-complexity LMs. In these works, parameters from the teacher are also transferred to the student when possible; an example of the direct parameter transfer approach.

Model conversion. Distinct from model compression, there is also a series of works aimed at converting (a.k.a. uptraining) existing models into new, more efficient ones. Parameter transfer is generally applied, e.g., in the conversion from standard attention to sparse attention (Zaheer et al., ), RNN (Kasai et al., ), linear attention (Mao, ; Mercat et al., ), and multi/grouped-query attention (Ainslie et al., ). Distillation, or behavior transfer, from the standard transformer as the teacher has been also introduced, starting from early attempts to convert transformers to RNN using teacher outputs or generations (Senellart et al., ; Tang et al., ). Zhang et al. () further tried to reproduce the attention matrix of standard attention in linear attention with a loss to align the teacher and student attention weights. Bhati et al. () also used output distillation for pretraining Mamba-based speech models, but without parameter transfer. A concurrent work then combined the distillation on attention weights, outputs, and hidden states (Bick et al., ), and managed to convert large-scale LMs into Mamba with limited computation costs and performance loss. These works are nevertheless focused on conversion during pretraining; by contrast, we focus on conversion during fine-tuning, which is generally cheaper and more accessible. Choi () investigated conversion to RetNet or Mamba during fine-tuning on commonsense reasoning tasks, but using only parameter transfer. Zhang et al. () also considered the case of conversion after fine-tuning by attention weight alignment. A concurrent work converts Llama3 into Mamba using distillation on outputs and generations during instruction fine-tuning (Wang et al., ). Meanwhile, LoSparse jointly fine-tunes and compresses linear layers into LoRA + sparse matrix (Li et al., ). Our work is nevertheless distinct in that we investigate different modes of distillation under a general framework of joint fine-tuning and conversion into efficient architectures based on layerwise distillation, which allows a broader range of models and tasks as discussed in Section 1 above.

## 3 Methods

We aim at building a broadly applicable framework to convert an existing pretrained model into a task-specific linear-complexity one. To achieve this goal, we combine parameter transfer and distillation to best reproduce the performance of the target teacher model, a standard transformer model fine-tuned on the target task from the pretrained source teacher model.
For parameter transfer, we replace only the attention layers in the teacher with the more efficient sequence-mixing modules we intend to use, forming the student model as a result; most other layers including embeddings, feed-forward, and layernorm will be preserved. Following this approach, the teacher parameters can be used to initialize the student model as much as possible.
This constitutes our baseline unguided approach, in which we directly fine-tune the student model with transferred parameters. Instead, by posing loss on the difference of hidden states, we enforce constraints on the deviation of the student’s hidden states from the teacher. In this way, we directly guide each layer in the student to reproduce the behavior of the corresponding teacher layer. Different possible modes of such guidance are illustrated in Figure 1 and introduced below:

- •
Target Guided. We can directly transfer the parameters from the fine-tuned teacher target model and distill from it. More specifically, given the model inputs for training on a target classification task, we denote the hidden states from the student and the teacher target as ${\bm{H}}^{(s)}$ and ${\bm{H}}^{(t)}$, each with $m$ vectors; outputs (class probabilities) of the student and the teacher target as ${\bm{y}}^{(s)}$ and ${\bm{y}}^{(t)}$; and one-hot labels as ${\bm{y}}$. Then the loss terms are written as
$\displaystyle\mathcal{L}_{\text{CE}}({\bm{y}}^{(s)},{\bm{y}})$
$\displaystyle=-\sum_{i}{\bm{y}}_{i}\log({\bm{y}}^{(s)}_{i})$
(4)
$\displaystyle\mathcal{L}_{\text{KD}}({\bm{y}}^{(s)},{\bm{y}}^{(t)})$
$\displaystyle=\sum_{i}\big{(}\frac{{\bm{y}}^{(t)}_{i}}{\beta}\big{)}\log\left(
\frac{{\bm{y}}^{(t)}_{i}/\beta}{{\bm{y}}^{(s)}_{i}/\beta}\right)$
(5)
$\displaystyle\mathcal{L}_{\text{LD}}({\bm{H}}^{(s)},{\bm{H}}^{(t)})$
$\displaystyle=\frac{1}{m}\sum_{i=1}^{m}\big{(}{\bm{H}}^{(s)}_{i}-{\bm{H}}^{(t)
}_{i}\big{)}^{2}$
(6)
$\displaystyle\mathcal{L}$
$\displaystyle=\alpha_{\text{CE}}\mathcal{L}_{\text{CE}}+\alpha_{\text{KD}}
\mathcal{L}_{\text{KD}}+\alpha_{\text{LD}}\mathcal{L}_{\text{LD}}$
(7)
We set the temperature $\beta=2$ following Hinton et al. (). The final loss $\mathcal{L}$ is a combination of the cross entropy $\mathcal{L}_{\text{CE}}$, the KL divergence $\mathcal{L}_{\text{KD}}$ between teacher and student outputs in standard KD, and an MSE loss $\mathcal{L}_{\text{LD}}$ between hidden states for layerwise distillation. The three terms are weighted by a set of hyperparameters $\alpha$. Instead, the unguided mode uses $\mathcal{L}_{\text{CE}}$ only. More specifically, as the target sequence-mixing module (e.g., RNN) is nevertheless distinct from standard attention, it can be more difficult to reproduce the hidden states right after the sequence-mixing layer. Instead, we extract the hidden states after the processing of the following feed-forward layer as ${\bm{H}}$ for layerwise distillation.
- •
Trajectory guided. Models initialized from pretrained parameters perform much better than random initialization when trained on downstream tasks. This can be explained by the rich knowledge internalized into the pretrained model, which may be lost during fine-tuning. It has been found that preserving the knowledge during fine-tuning can be helpful for downstream performance (Li et al., ; He & Garner, ). Therefore, it can be sub-optimal to only reproduce the hidden states produced by the fine-tuned teacher target model, which has already lost certain knowledge. We are further inspired by Li et al. () who preserve certain hidden states from shifting during fine-tuning for better downstream performance, as well as classical homotopy or continuation methods that solve an optimization problem by tracing a series of surrogate problems that gradually lead to the goal (Allgower & Georg, ; Lin et al., ). Hence we first initialize the student with source teacher parameters, and then simultaneously fine-tune the teacher and the student. In each step, the distillation loss is not computed against the hidden states of the final teacher target model, but the current teacher model instead. In this way, the student model will reproduce the trajectory of the hidden states during the whole teacher fine-tuning process. As a result, we can better capture the pretraining knowledge present at the early stage of fine-tuning, and adapt to the target task in a way more similar to the teacher. As the converted architecture is different from the original one and may require multiple steps to approximate the current teacher hidden states, we also consider updating the teacher at a slower pace, e.g., every $T_{u}$ steps.
- •
Waypoint guided. Joint fine-tuning of the teacher and student will essentially double the computation. Instead, we can make use of the checkpoints in teacher fine-tuning. That is to say, we mark the teacher checkpoint of every $T_{w}$ steps as the guiding waypoint. During conversion, we use the nearest waypoint ahead of the current training step as the current teacher. In this way, we approximate the trajectory guided approach in a more coarse-grained way but without extra computation compared to the target guided approach.
- •
Hybrid. The target architecture is after all distinct from the standard transformer, and hidden features in an optimal transformer can be sub-optimal for the target architecture. Therefore, a hybrid scheme between the target guided and unguided modes can be useful. Similar to KD annealing (Jafari et al., ), we apply the target guided mode at the initial phase of fine-tuning. We later switch to the unguided mode, which allows the model to train without the distillation constraint. In this way, the initial distillation provides a good initialization for the following optimization.

Figure: Figure 1: A conceptual illustration of the hidden states shift during training under different modes of distillation. Given the trajectory (green line) of the hidden states during the fine-tuning from the source teacher model (i.e. pretrained transformer) to the target teacher model (i.e. transformer fine-tuned on the target task), we consider: a) Unguided: parameter transfer without any distillation; b) Target Guided: distill from the target teacher; c) Trajectory/Waypoint Guided: gradually distill from a series of models on the trajectory; d) Hybrid: distill from the target teacher until a certain step.
Refer to caption: x1.png

## 4 Empirical studies

### 4.1 Model configuration

We carry out experiments on several scenarios that are representative to demonstrate the advantages of our approach, including

- •
Converting RoBERTa to Linformer, fine-tuned on language processing tasks, and
- •
Converting Wav2Vec2 to the latest Mamba2, fine-tuned on speech tasks.

Although not our focus, we also carry out an extra experiment to convert a widely-used open-source LM of various sizes, namely Pythia (Biderman et al., ), into Mamba for language modeling by retraining on a small subset of the pretraining corpus.

As mentioned above, we build the target architecture that allows parameter transfer as much as possible. In this way, the architecture may not be identical to the standard Linformer, Mamba, or Mamba2, as specifically described below:

- •
Language processing by RoBERTa $\rightarrow$ Linformer: We only add the $E$ and $F$ matrices to project the hidden features, which are randomly initialized by $\mathcal{N}(0,1)$ to preserve the scale of features after projection. All other parts of the model are left unchanged and initialized with pretrained RoBERTa parameters. We follow the reported best Linformer configuration (KV or layer-shared $E$ and $F$, with rank $k=256$).
- •
Language modeling by Pythia $\rightarrow$ Mamba: We only replace each attention layer with a Mamba mixer of the same hidden size, initialized under the Mamba scheme.
- •
Speech processing by Wav2Vec2 $\rightarrow$ Mamba2: Wav2Vec2 is a bidirectional model, while Mamba2 is unidirectional. Hence we build a bidirectional Mamba2 similar to the one by Zhang et al. (). Each attention layer is replaced by two Mamba2 mixers of the same hidden size, but with expand factor = 1, as shown in Figure 2.

Figure: Figure 2: Encoder layers in the student bidirectional Mamba2 model (b), converted from Wav2Vec2 encoder layers in the teacher model (a). In each encoder layer, the attention layer is replaced by two new forward and backward Mamba2 mixers. Inputs and outputs of the backward mixer are timewise inverted. Hidden states after the feed-forward layer are extracted for layerwise distillation.
Refer to caption: x2.png

In this way, all the models have roughly the same number of parameters after conversion. Classification tasks are carried out by adding a linear prediction head after mean pooling. The modeling and training configuration under each scenario is further introduced below, and more details including hyperparameters are provided in Appendix A.

### 4.2 Language processing

We follow the settings in the Linformer paper to fine-tune the models and report the validation accuracy on a set of widely-used benchmarks: QNLI (Rajpurkar et al., ) for natural language inference, QQP (Iyer et al., ) for text similarity, as well as SST2 (Socher et al., ) and IMDB (Maas et al., ) for sentiment classification; the former three being part of the GLUE benchmark (Wang et al., ). The models are converted from and compared with pretrained RoBERTa-base. In this way, we are able to directly compare with the reported results of the fully re-pretrained Linformer, which is not publicly available. We fine-tune all models for 10 epochs, and a waypoint is taken after each epoch.

The results are given in Table 1. Compared with our fine-tuned standard RoBERTa, the re-pretrained Linformer has a slight drop in performance (1.3% on average). This is typical and can be explained by the reduced expressivity compared to the standard attention. While without extensive pretraining, unguided conversion with only parameter transfer leads to much lower accuracy and severe training instability. Despite much lower computational costs, such performance regression is unacceptable. Fortunately, the regression can be alleviated using CALD. By simply using the target guided approach, the performance drop can be reduced to merely 0.7% on average compared to pretrained Linformer. With the trajectory guided approach, the CALD model can even reach better results (+0.2% on average) but without any pretraining. The approximated waypoint guided approach leads to only slightly lower results, but still better than the target guided one. The hybrid approach leads to lower results compared to target guided, indicating that the guidance is helpful even after early training updates, which matches the rather low unguided results.
In addition, we try to use the target guided approach but initialized from the teacher source parameters, which shows slightly lower results.
All these results indicate that the behavior transfer by CALD is effective and critical for the joint conversion and fine-tuning in this scenario to minimize the performance regression. Particularly, the trajectory or waypoint guided approaches are helpful.

**Table 1: Performance of Linformer models converted from RoBERTa, measured by accuracy on benchmarks reported by the original Linformer paper (Wang et al., ), in comparison with the reported results of a fully re-pretrained Linformer. The best results are bolded.**
|  | QNLI | QQP | SST2 | IMDB | Average |
| --- | --- | --- | --- | --- | --- |
| Pretrained Linformer | 91.2% | 90.8% | 93.1% | 94.1% | 92.3% |
| Std. RoBERTa | 92.4% | 91.8% | 95.3% | 95.7% | 93.8% +1.3 |
| Unguided | 53.1% | 73.3% | 82.6% | 82.6% | 72.9% -19.4 |
| CALD |  |  |  |  |  |
| - Target Guided | 89.0% | 91.8% | 93.3% | 92.3% | 91.6% -0.7 |
| - Src. init. | 88.5% | 91.7% | 93.1% | 92.3% | 91.4% -0.9 |
| - Trajectory Guided | 91.2% | 91.9% | 94.0% | 93.1% | 92.5% +0.2 |
| - Waypoint Guided | 89.9% | 91.9% | 93.7% | 92.8% | 92.1% -0.2 |
| - Hybrid | 86.8% | 90.8% | 91.4% | 90.5% | 89.9% -2.4 |

**Table 2: Zero-shot performance of downstream evaluation tasks reported by Pythia (Biderman et al., ) on Mamba-based language models converted from Pythia-1B, using a small subset of Pile.**
| Model | Lambada | PIQA | Winog. | WSC | ArcE | ArcC | SciQ | LogiQA | Avg. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Pythia-1B | 0.562 | 0.707 | 0.535 | 0.670 | 0.570 | 0.244 | 0.839 | 0.221 | 0.544 |
| 0.5% Pile / 1.5B tokens |  |  |  |  |  |  |  |  |  |
| Unguided | 0.394 | 0.671 | 0.493 | 0.542 | 0.502 | 0.211 | 0.778 | 0.246 | 0.479 |
| Tgt. Gd. | 0.410 | 0.686 | 0.504 | 0.608 | 0.538 | 0.210 | 0.775 | 0.230 | 0.495 |
| Hybrid | 0.432 | 0.683 | 0.507 | 0.608 | 0.537 | 0.209 | 0.788 | 0.238 | 0.500 |
| 1% Pile / 3.0B tokens |  |  |  |  |  |  |  |  |  |
| Unguided | 0.453 | 0.673 | 0.514 | 0.571 | 0.525 | 0.217 | 0.794 | 0.217 | 0.495 |
| Tgt. Gd. | 0.449 | 0.689 | 0.518 | 0.626 | 0.535 | 0.220 | 0.798 | 0.218 | 0.507 |
| Hybrid | 0.477 | 0.692 | 0.513 | 0.630 | 0.529 | 0.221 | 0.812 | 0.240 | 0.514 |
| 2% Pile / 6.0B tokens |  |  |  |  |  |  |  |  |  |
| Unguided | 0.475 | 0.687 | 0.535 | 0.612 | 0.543 | 0.222 | 0.802 | 0.237 | 0.514 |
| Tgt. Gd. | 0.459 | 0.697 | 0.527 | 0.656 | 0.547 | 0.218 | 0.817 | 0.235 | 0.520 |
| Hybrid | 0.485 | 0.708 | 0.510 | 0.645 | 0.556 | 0.224 | 0.829 | 0.244 | 0.525 |

### 4.3 Language modeling

We then present the results of experiments on converting Pythia to Mamba for language modeling. This is nevertheless not our focus as recent pretrained LMs are typically directly evaluated by zero or few-shot performance in downstream tasks without fine-tuning. Therefore, we will essentially need to do the re-pretraining under this scenario. Yet we try to convert the model by training only on a tiny (0.5%, 1%, or 2%) subset of the pretraining corpus, i.e. the deduplicated Pile (Gao et al., ; Biderman et al., ). Under this scenario, we can only compare the unguided, target guided, and hybrid approaches. Models are trained for 2 epochs.
Due to limitations in computational resources, we perform experiments only on Pythia-1B.

We evaluate the models using a set of benchmarks based on those reported in the Pythia paper as shown in Table 2. Compared with the standard Pythia-1B, there remains a performance gap. Nevertheless, we can find that the gap gets limited using 2% data. More importantly, the CALD models consistently get better results than the unguided ones, especially with the hybrid method. We also observe that the unguided training is rather unstable with frequent NaNs. All the results support the effectiveness of our CALD approach.

### 4.4 Speech processing

**Table 3: Performance (%) of Mamba2-converted Wav2Vec2 models: word error rate for ASR on TED-LIUM, as well as accuracy for intent classification on SLURP and speaker ID on VoxCeleb1. The best results are bolded.**
|  | ASR WER $\downarrow$ | IC Acc. $\uparrow$ | SID Acc. $\uparrow$ |
| --- | --- | --- | --- |
| Std. Wav2Vec2 | 6.24 | 91.70 | 96.09 |
| Unguided | 11.29 | 79.68 | 84.24 |
| CALD |  |  |  |
| - Target Guided | 6.56 | 90.43 | 96.56 |
| - Waypoint Guided | 6.92 | 90.32 | 96.16 |
| - Hybrid | 6.41 | 91.23 | 96.41 |

Figure: Figure 3: Average cosine distance between the hidden states produced by the initial model and the checkpoints during fine-tuning after each epoch (for NLP tasks) or every 10,000 steps (for speech tasks). Unlike NLP models, features produced by the fine-tuned speech models are far from the initial ones since the early phase of fine-tuning.
Refer to caption: x3.png

We then convert the Wav2Vec2-large models into the latest Mamba2, using a bidirectional architecture specified above. We evaluate them on three common speech tasks with corresponding common benchmarks: ASR with TED-LIUMv3 (Hernandez et al., ), intent classification, a standard SLU task, with SLURP (Bastianelli et al., ), along with speaker ID with VoxCeleb1 (Nagrani et al., ), following the SUPERB configuration (Yang et al., ). Specifically, ASR training is performed by CTC loss instead of cross entropy as $L_{CE}$. The converted models are compared with standard fine-tuned Wav2Vec2 models. We take a waypoint every 10,000 steps during fine-tuning. Test word error rate (WER) and accuracy are reported respectively.

Results are given in Table 3, which are consistent with previous findings. There is a significant gap between the performance of standard Wav2Vec2 and the unguided conversion. This can be largely alleviated by CALD using the target guided approach, reaching results close to or even better than standard Wav2Vec2. The hybrid approach further brings slight improvements on ASR and IC.

However, the waypoint guided approach is found to be not helpful or even harmful; we choose not to proceed with the more precise trajectory guided approach. Our assumption is that the trajectory guided approach helps retain the knowledge in the hidden states of the pretrained model. Hence We hypothesize that the hidden states of the Wav2Vec2 model undergo significant shift during fine-tuning. As a result, the hidden states of the original pretrained model or those in the early phase of fine-tuning are less useful for the target task. To examine the hypothesis, for each model, we compute the average cosine distance of the hidden states produced by the source teacher and the waypoints, using 100 random data samples. As shown in Figure 3, as for fine-tuning RoBERTa on NLP tasks, the cosine distance is relatively small and increases gradually. While the distance is much larger during the fine-tuning of Wav2Vec2 on speech tasks, even at the early phase of fine-tuning. In fact, the distance has already reached 0.372 after only 2,000 steps of fine-tuning on TED-LIUM. Therefore, our hypothesis is supported that the trajectory guided approach works better in the cases when there is limited hidden state shift during fine-tuning, and when the pretrained model features remain relevant for the target task. While the condition is not satisfied for speech tasks.

Figure: Figure 4: The trajectory of hidden states during training under different modes of guidance, visualized using t-SNE.
Refer to caption: x4.png

### 4.5 Hidden state trajectory

We further visualize the trajectory of hidden states during fine-tuning on the QNLI dataset using different guidance modes by applying t-SNE to concatenated hidden states of 20 random samples in the training set. Hidden states produced by every half-epoch checkpoint are included, hence a total of 20 points during training are plotted for each model. As shown in Figure 4, the actual trajectory matches our expectation conceptualized in Figure 1 that the target guided model approaches the fine-tuned target teacher model (the end of the green trajectory), while the waypoint or trajectory guided models further mimic the trajectory, which is found to be beneficial to the final results. In sharp contrast, the unguided model will completely deviate from the teacher, which leads to training instability and failure to reproduce the original performance.

## 5 Conclusion

We investigate the task of converting various existing pretrained transformer models into efficient linear-complexity architectures without the need to redo the whole pretraining. For this goal, we propose a novel and broadly compatible Cross Architecture Layerwise Distillation (CALD) approach, and further enhance CALD with a trajectory-based guidance or a hybrid approach. Through our empirical experiments, we confirm that CALD can effectively convert the transformer into an efficient linear-complexity model with performance close to or on par with the standard transformer, much better than directly converted models. We also verify the versatility of CALD by experiments on multiple tasks with various speech and language models. We further identify the tasks where the enhancement approaches are effective in improving the results.

## Appendix A Implementation details

**Table 4: Hyperparameters for language processing tasks.**
|  |  | QNLI | QQP | SST2 | IMDB |
| --- | --- | --- | --- | --- | --- |
| Std. RoBERTa | LR | 1e-5 | 3e-5 | 1e-5 | 1e-5 |
| Unguided | LR | 1e-4 | 1e-4 | 1e-4 | 3e-4 |
| Target Guided | LR | 1e-4 | 1e-4 | 1e-4 | 3e-4 |
| $\alpha_{LD}$ | 25 | 20 | 10 | 15 |  |
| Share mode | Layer | Layer | KV | KV |  |
| Traj. Guided | LR | 3e-4 | 1e-4 | 1e-4 | 3e-4 |
| $\alpha_{LD}$ | 30 | 25 | 30 | 40 |  |
| Share mode | Layer | KV | KV | KV |  |
| $T_{u}$ | 5 | 1 | 3 | 3 |  |
| Waypoint Guided | LR | 1e-4 | 1e-4 | 1e-4 | 3e-4 |
| $\alpha_{LD}$ | 10 | 20 | 5 | 25 |  |
| Share mode | KV | Layer | KV | KV |  |

We build the models as mentioned above, and train the models using the AdamW optimizer. On Mamba/Mamba2 models, we also find it important to apply layer normalization on hidden states prior to computing the layerwise distillation loss, i.e. we take the hidden states after the layernorm operation at the following layer. Otherwise the hidden states will have a rather wide range and bring training instability.
We only use output distillation by setting $\alpha_{KD}=1$ in ASR experiments. On other experiments, we find that the distillation on output probabilities actually leads to suboptimal results, and thus set $\alpha_{KD}=0$. We always use $\alpha_{CE}=1$. As for the hybrid mode of distillation, we start with the optimal target-guided configuration, but switch to unguided mode when the loss is close to converge, generally at around 30% of the total training steps.

For each target task, we train the models based on its common training recipe for standard transformers, and further search and determine hyperparameters and the model checkpoint to use according to the validation performance. Details for each task are listed below, and more specific details can be accessed from our source code.

- •
Language processing: We train the model using a learning rate (LR) schedule of linear decay with 6% warmup steps. A total batch size of 32 is used. Other hyperparameters are listed in Table 4.
- •
Language modeling: We train the model using a learning rate of 3e-4 and a schedule of cosine decay with 3% warmup steps and a 3e-5 minimum learning rate. We use a total batch size of 128, each with 2048 tokens. We also apply the weight decay of scale 0.1 on Mamba parameters following the Mamba training scheme, along with gradient norm clipped to 1. We set $\alpha_{LD}=15$ in the experiments. FP16 mixed precision training is adopted.
- •
Speech processing: We train the model using a learning rate schedule of exponential decay with 5000 warmup steps and 1.5e-6 minimum learning rate. We use a dynamic batch strategy that leads to a batch size of approximately 64 in average. We also apply the weight decay of scale 5e-3, along with gradient norm clipped to 1. For ASR, IC, and SID, we set $\alpha_{LD}$ to 15, 15, and 30, respectively. BF16 mixed precision training is adopted. We find that these hyperparameters work well across different modes of distillation.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see the above "weakness" section.
soundness: 2
strengths: The paper is mostly well-written and the problem it tries to solve is an important one. It does a good job reference related works. Novelty-wise, although the individual components of the proposed method aren't new, it has a fair amount of novelty for combining them and designing different fine-tuning strategies for training. The figure showing the actual trajectory of hidden states is interesting and informative to show how parameters change in those different methods.
summary: This paper focuses on the issue of the quadratic computational complexity of standard self-attention in transformer models and proposes a way to convert pre-trained transformer models with quadratic attention to more efficient linear attention, while trying to minimize performance degradation. They call this method Cross-Architecture Layerwise Distillation (CALD), and the core of this approach is the combination of layer-wise distillation with parameter transfer, making the fine-tuning process more efficient, and a few specially designed fine-tuning strategies. 

The authors propose a few strategies to compare with a simple "unguided" approach. They report on language and speech tasks, and show that their proposed methods consistently outperform the unguided approach, and in some cases, some methods can even outperform linear-attention models that are repretrained.
weaknesses: I'm not fully convinced of the claimed merits of the method. While it seems to be convincing that those proposed methods are better than unguided approach, there are unanswered questions regarding those approaches, because it seems to have different behaviors in different tasks. E.g. in experiments reported in Table 1, unguided's performance significantly lags behind other methods, while in Table 2, the relative difference is much smaller. This might indicate that the unguided hyper-params might not be well-tuned for certain tasks. Also, in Table 1, the "hybrid" approach gives the worst performance among all CALD variants, but it seems to give the best performance in many of the cases in Table 2 and 3. This to me suggests that there might be some other fundamental factors that caused those different behaviors shown in those Tables. 

There are a couple of misreferences in the writing, where the authors wanted to reference Figure 1 but said Figure 3 instead. There's a reference of Figure 4.1 and Figure 4.4 which I believe are typos that point to other Figures in the paper. In section 2.2, there's an incomplete sentence "an example of the direct parameter transfer approach." 

It took me a while to fully understand Figure 1, and I think the source of the confusion comes from that the diagram involves both shifts in hidden states and also relapse of time, like the "trajectory/waypoint guided" blue arrow actually "follows" the green arrow as time goes by. I understand the authors might want to show this diagram here to be consistent with Figure 4 shown later in the analysis, but my feeling is this diagram isn't the clearest way to demonstrate the difference between the approaches. I would feel that providing simple pseudo-code instead of this diagram + lots of text in method bullet points might be a better way to accurately represent the methods. 

For ASR tasks, the authors use CTC loss instead of CE. How that does change Equation 4? E.g. do you still compute the CE per-frame or some other computation? Also, since CTC adopts a blank token in its output and we have research works that show a CTC model would predict blanks for most frames, I feel there should be some interactions with this aspect instead of treating blanks and non-blanks equally.

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Q1 - How does CALD compare with other recent state-space models in terms of accuracy, inference time, and training stability, particularly for large-scale speech tasks?

Q2 - Have the authors tested CALD in real-world deployment scenarios for speech and language tasks? If so, how does its performance hold up compared to controlled experiments?
soundness: 3
strengths: 1.) The introduction of the Cross-Architecture Layerwise Distillation (CALD) is noteworthy, as it aims to convert pre-trained transformer models into efficient linear-complexity architectures. This bridges a critical gap by enabling the reuse of existing pre-trained models without the need for resource-intensive pretraining, especially for non-text domains like speech.

2.) The empirical studies are well-structured and cover various conversion tasks from RoBERTa to Linformer for NLP tasks, Wav2Vec2 to Mamba2 for speech tasks, and Pythia to Mamba for language modeling. The detailed experiments provide a strong basis for assessing the effectiveness of CALD, showing that guided approaches outperform unguided ones significantly.

3.) The paper does a thorough job of comparing the proposed method with existing approaches, including both guided and unguided methods. The results highlight that CALD, especially with trajectory-guided or waypoint-guided distillation, can effectively maintain or improve performance close to standard transformer models.

4.) The inclusion of diverse benchmarks (e.g., QNLI, QQP, TED-LIUM, SLURP, VoxCeleb1) and the report of performance improvements in areas like word error rate (WER) and accuracy provide strong evidence for the robustness of the proposed methodology.

5.) The paper provides a solid theoretical explanation for why trajectory-guided distillation can retain pre-training knowledge and ensure better downstream task performance.
summary: This paper introduces a method called Cross-Architecture Layerwise Distillation (CALD). The goal of CALD is to convert existing pre-trained transformer models into linear-complexity models and also fine-tune tje,, making them more computationally efficient without the need for extensive re-pretraining. This approach enables the efficient adaptation of models to different architectures, such as transforming RoBERTa into Linformer for NLP tasks and Wav2Vec2 into Mamba2 for speech processing tasks. CALD combines parameter transfer, where attention layers are replaced with efficient sequence-mixing modules, with knowledge distillation, where the student model learns from the teacher model's behavior. Four distillation modes are presented: Target Guided, Trajectory Guided, Waypoint Guided, and Hybrid.
The paper highlights that CALD effectively minimizes performance loss compared to directly converted models, demonstrating its efficacy across various language and speech processing tasks.
weaknesses: Limited Discussion on Limitations: While the results are strong, the paper lacks an in-depth discussion on potential limitations or failure cases of the CALD approach, especially in more complex or resource-constrained real-world scenarios.

Computational Cost Analysis: Although the proposed method is presented as a more efficient alternative, there is insufficient analysis of the actual computational cost savings compared to traditional pre-training or conversion processes. A more detailed breakdown of memory and time requirements would enhance the practical relevance.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - How do you choose hyper parameters in equ. 7?
- In Table 1, Src. init results are worse, why?
- Hybrid approach is better in Table 2 and 3, but it is worse in Table 1. Any explanation?  The trajectory or waypoint guided approaches are worse than the target Guided method in Table 3 but better in Table 1. Any explanation? Has the model learnt enough from the waypoint model?
soundness: 3
strengths: - CALD provides an effective and cost-effective approach to build linear complexity based language model by leveraging pre-trained transformer based language model.
- Multiple knowledge distillation methods are proposed and examined.
- The experiments are conducted with different modalities including both speech and text, and different linear complexity transformers. The results show CALD could achieve good results.
summary: In this work, a cross architecture layer wise distillation (CALD) approach is proposed, which includes converting a transformer model to a linear time substitute by replacing attention modules and fine-tuning the new model towards the target task. The fine-tuning is enhanced by knowledge distillation at different levels and stages. The proposed method is examined in both language modeling and speech processing tasks. The results show CALD can effectively narrow the gaps between linear time based substitute and original transformer model on the downstream tasks.
weaknesses: - Experimental descriptions could be improved by providing more information. For example,  in section 4.2, "the models are converted from and compared with retrained RoBERTa-base."  But in the following description about Table 1, "In addition, we try to use the target guided approach but initialized from the teacher source parameters, which shows slightly lower results". Are other models are not initialized from RoBERTa-base? 
- More detailed analysis about different distillation methods would be helpful.
- The proposed method could be further improved by including the discussion about linear complexity attention module initialization, which could be very important for the final results

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: Q: the conversion/distillation is made for a particular downstream task, would it be possible to make it for the pre-trained - general purpose - model ?

Q: can we further speed-up inference through quantization (with added benefits)?
soundness: 4
strengths: -new approach to address model conversion/distillation from transformer (quadratic) architectures to mamba 2 (linear) architectures

-not only applied to written language but also to speech with convincing results (3 types of tasks overall: NPL, LM and Speech)
summary: Latests long-context recurrent models such as Mamba2 have demonstrated performance similar to transformer models, even in large scale scenarios, while being more efficient (no quadratic complexity of attention).
However, they need to be retrained from scratch, which could slow down the adoption of their architecture—especially with so many pre-trained transformer models already available.
Moreover, for non-text modalities such as speech, large scale Mamba 2 models are simply not available yet.
Hence authors propose leveraging off-the-shelf pre-trained (transformer) models and converting them into the target (mamba 2) model with linear complexity.
More specifically, they explore the possibility of converting an existing pre-trained transformer into a linear-complexity model for a specific downstream task; approach proposed is called Cross-Architecture Layerwise Distillation (CALD) and combines parameter transfer and distillation.
Multiple scenarios are explored, including converting RoBERTa to Linformer for NLP tasks, Pythia to Mamba for language modeling, and Wav2Vec2 to Mamba2 for speech tasks; converted models show minimal to no performance loss compared to standard transformers while being of linear (instead of quadratic) complexity.
weaknesses: -the obtained model is no longer a general purpose model (converted for a specific downstream task)

-some inference speed evaluations would have been a plus (we know that obtained models with linear complexity should be faster though)

Meta-review: This paper proposes a approach based on distillation for efficient fine-tuning, with an advantage of linear complexity.

All reviewers praise the idea interesting, the paper easy to follow, and the experiments solid across multiple modalities and benchmarks.

Despite the recommendation of acceptance, the paper is not without weaknesses. The paper is conceptually not a big departure from prior work; hence the recommendation of a poster. There are valuable suggestions provided by the reviewers, and I encourage the authors to incorporate the feedback to improve the paper.

Decision: Accept (Poster)