Joint Fine-tuning and Conversion of Pretrained Speech and Language Models towards Linear Complexity
Authors: Mutian He
Philip N. Garner
Venue: ICLR 2025 Poster (2025)
Keywords: pretrained models
efficient attention
uptraining
speech processing
Primary area: applications to computer vision, audio, language, and other modalities
TL;DR: Fine-tuning and converting pretrained speech and language transformer models into task-specific efficient/linear-complexity models without re-pretraining

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.

Published as a conference paper at ICLR 2025
JOINT FINE-TUNING AND CONVERSION OF PRE- TRAINED SPEECH AND LANGUAGE MODELS TOWARDS LINEAR COMPLEXITY
Mutian Het-2, Philip N. Garner
1 Idiap Research Institute, Martigny, Switzerland
2 Ecole Polytechnique Fédérale de Lausanne (EPFL., Switzerland (mutian. he, phil aet.aaa....ii....ar.r.......... garner)eidiap.co
ABSTRACT
Architectures such Linformer and Mamba have recently emerged as competi- tive linear time replacements for transformers. However, corresponding large pre- trained 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 linear time substitute and fine-tunes it to 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 parame- ters. In a series of empirical studies on language processing, language modeling. and speech processing,
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.
INTRODUCTION
Recent progress in language modeling from BERT (Devlin al., 2019) Llama3 (Dubey al., 2024) has witnessed the rise of the attention mechanism and transformer architecture. Such models are able scale up to very large model capacity and pretraining data for natural language pro- cessing (NLP), and have lead to breakthroughs in broad range of downstream applications under either fine-tuning or cero--soooin-nonteeee scenarios. This surely not restricted to the text domain: Off-the-shelf speech and audio models like Wav2Vec2 (Baevski et al.) (222)) HuBERT (Hsu et al.) 202T), and WavLM (Chen et al., (022) are pretrained on extensive audio data; they can be fine- tuned to obtain good performance
various tasks from automatic speech recognition (ASR) to spoken language understanding (SLU). However, such large-scale models are computationally ex- pensive; the main reason is that in the standard attention mechanism, each token interacts with all other tokens the input sequence, forming a complete graph. The computation time therefore in- creases quadratically with the input length, and linear-growing K V-cache is required. As result, the computation requirements can become prohibitive, especially for long contexts. A large co- hort of architectures and algorithms have been proposed to mitigate the problem and reach linear complexity, using techniques such as low-rank projection, hashing, and sparse attention (Tay et al.) 2023 More recently, there has been a revival of recurrent networks such state space models (Gu et al. 2022 Gu & Dao 2024), RWKV (Peng et al. 2023), and RetNet (Sun et al. 2023). By recur- rently 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 mod- els, notably Mamba2 (Dao & Gui 2024), 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 pub- licly released; the data and computational cost of repeating such pretraining for every model can be prohibitive for users and researchers of downstream applications. In particular, most models primarily developed for processing language text. These models are typically general-purpose

---

Published as a conference paper at ICLR 2025
substitutes for the standard attention mechanism, also applicable to other domains such 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 (MI.) and NL. community. This is a particular concern for research in end-to-end speech processing involving long-distance dependeney in the lower information density audio form: such models are prerequisite to efficiently process the abundant long-form but poorly segmented audio in the wild.
To address this issue, possible solution is to make use of off-the-shelff pretrained models and con- vert them into the target linear-complexii 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 (Kasa et al., 2021 Choi 2024). In addition to such parameter transfer, we can further guide the conversion by transfer- ring 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. (2019) explored conversion from BERT to RNN by aligning outputs; more recently, Zhang al. (0224)) and Bick et al. (024) tried to repro- duce the attention weights of the standard attention mechanism in linear attention models. However, it has been pointed out that reproducing attention weights is not prerequisite reproduce perfor- mance (Mao) 2022), and even causes training inefficiency and instability (Mercat et al.) 2024). More importantly, this method restricted to target models for which we can casily 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 lower rank (Wang al. 2020). 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 trans- former into cimeal-cmplyexity model on the downstream task. Our approach broadly covers differ- ent source models, target models, and target tasks. This is made possible by combining parameter transfer and layerwise distillation with hidden states from target teacher model, i.e., a pretrained standard transformer fine-tume -uned 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. Alter- natively, we can use the distilled model as 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) ap- proach under 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. Code and resources to reproduce our work are available at https: 7/gthhub. cow wwwaa//iinaaaaae www.inaarrieddisstlileee distill celsstil--reeraiine pretrained transformers
2 RELATED WORK
2.1 LINEARCCOPPLEXITY AND RECURRENT MODELS
Linear-complexity models. There is a large body of work on reducing the O(N2) time complexity of transformer owing to the attention computation A(Q, K, V) = softmax QK" matrices Q, K, and 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. 2019 [Bllagy et al.) 20200] Ainslic et al. 2020) Zaheer et al.) 2020). Alternatively, the attended positions can be determined oh- by hashing and clustering (iitaev
Va V. The N x d
2

---

Published as a conference paper at ICLR 2025
With Linformer, Wang et al. (2020) further found that the attention matrix A is inherently low-rank, and projects K and V into a lower k-dimensional space using k x N matrices E and F. The two matrices can be shared, or even shared between all layers. Then,
al. 2020; Roy al. 2021).
the O(N2) matrix multiplication can be avoided, as A can be provably approximated by
Q Vd
A(Q, K, V) softmax(
FV.
Recurrent models. By applying kernel trick in a reverse direction, Katharopoulos et al. (2020) reduce transformers to RNNs of infinite dimension. In this RNN, hidden states 2, = Lise @ K.:)
are the sum of an infiniteddimensionall feature @( ..:) mapped from each key vector (more specif- ically, along with their product with values). Using various finite dimension approximations of (Peng et al. 2021; Choromanski et al., 2021; Zhai et al.) (021), 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 long as there is no non-linear transform in the transition from 2, to žt+1- In addition to simply summing up features as in 2 above, RWKV (Peng et al. (023) and RetNet (Sun et al.) 2023) use exponential time decay of past features. Gating (Yang et al. 2024) or linear transform (Orvieto et al. (023) have also been introduced. Similarly, a state-space model (SSM) is defined as a transformation of 1-D imput sequence a=t) to output (t) via an N-D hidden state h(t). following
h' (t) Ah(t) + Bx(()
yt) Ch(t).
(3)
Hence the discretization of an SSM is akin to an RNN with linear state transition (Gu et al., 2021 2022). Mamba then renders the transition matrices as Ar. Br, and C, dependent on the input att) (Gu & Dao 2024). Mamba2 (Dao & Gu 2024) further allows multi-head formulation in hardware-efficient way under generalized framework of various SSMs and linear attention models.
Linear-complexity models for speech. With the lower information density or longer form of spo- ken 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. 2023 Koluguri et al., 2024) and RWKV (An & Zhang, 2023). Also, Mamba has been applied to speech separation (Jiang 2024b). Zhang (2024b) investigated different options of bidirectional Mamba for ASR. Furthermore, there
et al.
2024 Li & Chen 2024) and enhancement (Li & Chen 2024 Zhang et al.
cet al.
have been attempts to build Mamba-bian based pretrained speech models (Shams et al. 2024: Yadav & Tan [2024] [Bhati et al. 2024). Such pretraining is nevertheless expensive, especially when taking in 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 student one by matching the output probabilities is a standard approach to compress ML models, including pretrained language models (LMs) (Hinton al. 2015; [aanh et al. 2019). Intermediate hidden states of the teacher and student can be further a aligned by layerwise distillation (Sun et al.) 2019 AAgulaar et al.) 2020). In addition, Brown al. (2023) have attempted distill and compress existing pretrained linear- complexity LMs. In these works, parameters from the teacher are also transferred to the student when possible, hence also represent the direct parameter transfer approach.
Distinct from model compression, there is also a series of works aimed at con- verting (a.k.a. uptraining) existing models into new, more efficient ones. Parameter transfer gener- ally applied, c.g., in the conversion from standard attention to sparse attention (Zaheer et al. 2020).
Model conversion.
RNN (Kasai et 2021), linear attention (Mao) 2022 Mercat et al. (002)) and attention (Ainslie et al. 2023). 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 al. 2018 Tang et al. 2019). Zhang et al. (2024a) further tried to reproduce the attention matrix of standard attention in linear attention with loss to align the teacher and student attention weights. [Bhati et al. (2024) also used output distillation for pretraining Mamba- based speech models, but without parameter transfer. A concurrent work then combined the distillation attention weights, outputs, and hidden states (Bick et al. 2024). and
3

---

Published as a conference paper at ICLR 2025
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 fo- conversion during fine-tuning which is generally cheaper and more accessible. Choi (2024) investigated conversion to RetNet or Mamba during fine-tuning on commonsense reasoning tasks, but using only parameter transfer. [Zhang al. (2024a) 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., 2024). Mean- while, LoSparse jointly fine-tunes and compresses linear layers into LoRA + sparse matrix (Li et al., 2023). Our work is nevertheless distinct in that we investigate different modes of distillation under 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 Section T aboov..
METHODS
We aim at building a broadly applicable framework to convert an existing pretrained model into task-specific iomer-complexity one. To achieve this goal, we combine parameter transfer and distil- lation to best reproduce the performance of the target teacher model. 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 aequence-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 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. In addition, we consider different ways or modes of such guidance, introduced below. These modes are also illustrated in Figure [ and pseudo-codes for the different modes of training algorithms are available in Appendix B
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 target classification task, we denote the hidden states from the student and the teacher target as (+)) and H (0, each with m vectors; outputs (class probabilities) of the student and the teacher target as (8)) and (t); and one-hot labels as y. Then the loss terms are written as
(c .y) Σ yb log(y!") (x)
(4)
18
( (,) = Σ
log
(5)
y
i
m H
LLD(H()) (()))
H11))
(6)
i=1
oCECC + CKD + OLDLLD
We set the temperature / 2 following HHinton et al. (2015). The final loss L is a com- bination of the cross entropy LCE. the KI. divergence LKD between teacher and student outputs in standard KDD, and MSE loss LLD between hidden states for layerwise distil- lation. The three terms are weighted by a set of hyperparameters a. Instead, the unguided mode uses LcE 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 H for layerwise distillation. As for tasks other than classification, we will replace the Lee term with the corresponding loss function, e.g. CTC loss for ASR tasks.
Trajectory guided. Models initialized from pretrained parameters perform much better than random initialization when trained 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
4

---

Published as a conference paper at ICLR 2025
Trajectory or Waypoint Guided
Source Teacher
Target
Teacher
Convert
Comvet
Target
Guided
Hybrid
Unguided
Student
Student (Target initialized)
(Source initialized)
Figure I: 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/Wayponn Guided: gradually distill from series of models on the trajectory; d) Hybrid: distill from the target teacher until certain step.
for downstream performance (Li et al. 2018 He & Garner 2023) Therefore, it can be sub-optimal 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. (2019) 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 opti- mization problem by tracing series of surrogate problems that gradually lead to the goal (Allgower & Georg 1990 Lin et al.) (023)) Hence we first initialize the student with source teacher parameters, and then simultaneously fine-tune
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 carly stage of fine-tuning, and adapt to the target task in way more similar the teacher. As the converted archi- tecture is different from the original one and may require multiple steps to approximate the current teacher hidden states, also consider updating the teacher slower pace, e.g..
every T. 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., 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 grained way but without extra computation compared to the target guided approach.
coarse-
The target architecture is after all distinct from the standard transformer, and hidden features 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 KD annealing (Jafari al,, (222), 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.
Hybrid.
5

---

Published as a conference paper at ICLR 2025
LayerNorm
Attention
Droprua
Nomm
Feed-Forward
Layoroorm
Mambal Micer
(F Formwads)
Dippost
LoyeWorm
Feod-Forwar
Irkear (Rackwart)
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.
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 language processing tasks, and
Converting Wav2Vec2 to the latest Mamba2, fine-tuned speech tasks.
Although not our focus, also carry out extra experiment to convert widely-used open-source LM various sizes, namely Pythia (Biderman et al. 2023), into Mamba for language modeling by retraining 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 Linformer: We only add the and F matrices to project the hidden features, which are randomly initialized by A/0,, 1) to preserve the scale features after projection. All other parts 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 256).
Language modeling by Pythia - Mamba: We only replace each attention layer with Mamba mixer of the same hidden size. The mixer is initialized using the specific Mamba state-space scheme to best preserve the past memory (Gu & Dao 2024).
bidirectional model, while Mamba2 is unidirectional. Hence we build a bidirectional Mamba2 similar to the one by Zhang et al. (2024b). Each attention layer is replaced by two Mamha2 mixers of the same hidden size, but with expand factor 1, shown in Figure 2 The mixer is similarly initialized with the specific Mamba2 state-space scheme (Dao & Gu 2024).
Speech processing by Wav2Vec2 + Mamba2: Wav2Vec2
In this way, all the models have roughly the same number of parameters after conversion. Classi- fication tasks carried out by adding linear prediction head after mean pooling. The modeling and training configuration under each scenario is further introduced below, and more details includ- ing hyperparameters provided in Appendix A Additional details the computational costs given in Appendix C
4.2 LANGUAGE PROCESSING
We follow the settings in the Linformer paper to fine-tune the models and report the validation accuracy set of widely-used benchmarks: QNLI (Rajpurkar al. 2016) for natural language

---

Published as a conference paper at ICLR 2025
Table I: Performance of Linformer models converted from RoBERTa, measured by aceuracy on benchmarks reported by the original Linformer paper (Wang et al., 2020). in comparison with the reported results of 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
69.4%
84.3%
83.6%
82.6%
80.0% -12.3
CALD
Target Guided
89.0%
91.8%
93.3%
92.3%
91.6% -07
Src. init.
88.5%
91.7%
93.1%
92.3%
91.4%
-Trajectory Guided
91.2%
91.9%
94.0%
93.1%
92.5% 40.2
Waypoint Guided
89.9%
91.9%
93.7%
92.8%
92. 1%
Hybrid
86.8%
90.8%
91.4%
90.5%
89.9% -24
inference, QQP (Iyer et 2017) for text similarity, well SST2 (Socher et al. (013) and IMDB MMas et al. 201 for sentiment classification; the former three being part of the GLUE benchmark Wang et al. 2019). 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 Compared with our fine-tuned standard RoBERTa, the re-pretrained Linformer has slight drop in performance (1.3% on average). This is typical and can be explained by the reduced expressivity compared 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 be alleviated using CALD. By simply using the target guided ap- proach, the performance drop can reduced to merely 0.7% on average compared to pretrained Linformer. With the trajectory guided approach, the CALD model can (+0.2% on average) but without any pretraining. The approximated waypoint guided approach leads 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; this is further discussed below in Section 4.4 In addition, we try to use the target guided approach but initialized from the teacher source parameters. Since the distillation should be easier if the student model is initialized with pa- rameters exactly from the teacher model (i.e. the teacher target), these Src. init. models show lower results as we expected. All these results indicate that the behavior transfer by CALD is effective and critical for the joint conversion and fine-tuning in this scenario minimize the performance regression. Particularly, the trajectory or waypoint guided approaches are helpful.
reach better results
4.3 LANGUAGE MODELING
We then present the results of experiments on converting Pythia 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 tiny (0.5%, %, or 2%) subset of the pretraining corpus, i.e. the deduplicated Pile (Gao al. 202T] Biderman et al.) 2023). Under this scenario, can only compare the unguided, target guided, and hybrid approaches. Models are trained for epochs. Due to limitations in computational resources, perform experiments only Pythia-
We evaluate the models using a set benchmarks based on those reported in the Pythia paper shown in Table 2 Compared with the standard Pythia-IB, 1B, there remains 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.

---

Published as a conference paper at ICLR 2025
Table 2:
Zero-shot performance of downstream evaluation tasks reported by Pythia (Biderman al.
2023 on Mamba-based language models converted from Pythia- using a small subset of Pile.
Model
Lambada
PIQA
Winog.
WSC
ArcE
ArcC
SciQ
LogiQA
Avg.
Pythia-IB
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 Tgt. Gd. Hybrid
0.394
0.671
0.493
0.542
0.502
0.211
0.778
0.246
0.479
0.410
0.686
0.504
0.608
0.538
0.210
0.775
0,230
0,495
0.432
0.683
0.507
0.608
0.537
0.209
0.788
0.238
0.500
1% Pile / 3.0B tokens
Unguided Tgt. Gd. Hybrid
0.453
0.673
0.514
0.571
0.525
0.217
0.794
0.217
0.495
0.449
0.689
0.518
0.626
0.535
0.220
0.798
0.218
0.507
0.479
0.693
0.520
0.648
0.531
0.224
0.808
0.247
0.519
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. Hybrid 0.485
0.459
0.697
0.527
0.656
0,547
0.218
0.817
0,235
0.520
0.708
0.510
0.645
0.556
0.224
0.829
0.244
0.525
Table 3: Performance (%) of Mamba2-converted Way2Vec2 models: word error rate for ASR on TED-LIUM, as well as accuracy for intent classification SLURP and speaker ID on VoxCelebl. The best results are bolded.
ASR WER
IC Acc. t
SID Acc.
Std. Wav2Vec2 Unguided
6.24
91.70
96.09
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
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
We then convert the Wav2Vec2-large models into the latest Mamba2, using bidirectional architec- ture specified above. We evaluate them on three common speech tasks with corresponding common benchmarks: ASR with TED-LIUMv3 (Herannez et al. 2018), intent classification, standard SLU task, with SLURP (Bastianelli et al. [220)), along with speaker with VoxCeleb1 (Nagrani et al. 2020), following the SUPERB configuration (Yang et al. [202]). Specifically, ASR training is performed by CTC loss instead of cross entropy as LCF. The converted models are compared with standard fine-tuned Wav2Vec2 models. We take waypoint every 10,000 steps during fine-tuning. Test word error rate (WER) and accuracy are reported respectively.
Results given in Table 3 which consistent with previous findings. There is a significant gap between the performance standard Wav2Vec2 and the unguided conversion. This can be largely alleviated by CALD using the target guided approach, reaching results close to standard Wav2Vec2. The hybrid approach further brings slight improvements on ASR and IC.
better than
However, the waypoint guided approach is found to be not helpful even harmful; choose not to proceed with the more precise trajectory guided approach. Our assumption that the trajectory guided approach helps retain the knowledge in the hidden states of the pretrained model. Hence
8

---

Published as a conference paper at ICLR 2025
QCP
0.60
0.55
0.50
0.45
0.40
0.30
0.25
0.20
0.15
0.10
0.05
0.00
3
6
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.
We hypothesize that the hidden states of the Wav2Vec2 model undergo significant shift during fine- tuning. As result, the hidden states of the original pretrained model 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 for fine-tuning RoBERTa on NLP tasks, the cosine distance is relatively small and increases gradually. While the distance 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, in which case allowing the model to deviate more from the initial behavior using the hybrid approach will be beneficial.
In addition, we have also attempted to train the models that are unguided and without any parameter transfer. Such large Mamba2-based speech model trained the target task from scratch is rather unstable and fails to converge well. This highlights the necessity of pretraining in such large models.
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 the actual trajectory matches our expectation conceptualized in Figure [ that the target guided model approaches the fine-tuned target teacher model (the end of the green trajectory), while the waypoint 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.
trajectory
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, pro-

---

Published as a conference paper at ICLR 2025
15
Std. Transformer
Unguided
Hybrid
Tar get Guided
Trajectory ded
10
Waypoint Guided
-15
-00
6.2 2.5
Figure 4: The trajectory of hidden states during training under different modes of guidance, visual- ized using I-SNE.
pose a novel and broadly compatible Cross Architecture Layerwise Distillation (CALD) approach, and further enhance CALD with a trajectory- based guidance or hybrid approach. Through our em- pirical experiments, we confirm that CALD can effectively convert the transformer into efficient linear-complexity model with performance close to 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.
par with the standard transformer, much
ACKNOWLEDGEMENTS
This work received funding under project SteADI, Swiss National Science Foundation grant 197479.
REFERENCES
Gustavo Aguilar, Yuan Ling, Yu Zhang, Benjamin Yao, Xing Fan, and Chenlei Guo. Knowledge distillation from internal representations. In The Thirty-Fourth AAAI Conference on Artificial Intelligence, AAAI 2020, pp. 7350-7357. AAAI Press, 2020. URL https:://oi..ogg 16699/aai. v34105. 6229
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. ETC: Encoding long and structured inputs in transformers. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, pp. 268-284- 2020. URL org/ 20tp/2220..o maammlp-mmin....
Joshua Ainslie, James Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: Training generalized multi-query transformer models from multi-head check- points. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Pro- cessing, EMNLP 2023, pp. 4895-4901, 2023. URL httts::: http://aaaaatttoooooooo oogg:/oolol..////222..... emn p-mma ma 298
10

---

Published as a conference paper at ICLR 2025
Eugene L Allgower and Kurt Georg. Springer Science & Business Media, 1990.
Numerical continuation methods: An introduction, volume 13.
Keyu An and Shiliang Zhang. Exploring RWKV for memory efficient and low latency stream- ing ASR, arxiv preprint arxiv:2309. 2023.
14758
Alexei Baevski, Yuhao Zhou, Abdelrahman Mohamed, and Michael Auli. wav2vec 2.0: A frame- work for self-supervised learning of speech representations. In Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, 2020.
Emanuele Bastianelli, Andrea Vanzo, Pawel Swietojanski, and Verena Rieser. SLURP: A spoken
language understanding resource package. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing, EMNLP 2020, pp- 7252-7262, 2020. URL https:
//aclanthology www.oogyy.o.g.ooooooooooo 2020. manip-main..1
Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arxiv preprint arxiv:2004.055/3: 2020. URL /tttps/://aaiiv.orrg////ss//
051 50
Saurabhchand Bhati, Yuan Gong, Leonid Karlinsky, Hilde Kuehne, Rogerio Feris, and James Glass. DASS: Distilled audio state space models are stronger and more duration-scalable learners. axxiv preprint arxiv:2407.0444 00.0000020..... 2024. URL
04082
Aviv Bick, Kevin Y. Li, Eric P. Xing, J. Zico Kolter, and Albert Gu. Transformers to SSMs: Distill- ing quadratic knowledge to subquadratic models. Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, 2024.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, Aviya Skowron, Lintang Sutawika, and Oskar van der Wal. Pythia: suite for analyzing large language models across training and scaling. In International Conference on Machine Learning (ICML) 2023, pp. 2330-2430. PMLR, 2023. URL http::ttp//eetttt #202/biderman223 html
/ /ttpe//eccccccegssoe/oe/ mlr. .ress//
Nathan Brown, Ashton Williamson, Tahj Anderson, and Logan Lawrence. Efficient transformer knowledge distillation: A performance review. In Proceedings of the 2023 Conference on Empir- ical Methods Natural Language Processing: EMNLP 2023 Industry Track, pp. 54-65, 2023.
URL
emnlp incustry.
Sanyuan Chen, Chengyi Wang. Zhengyang Chen, Yu Wu, Shujie Liu, Zhuo Chen, Jinyu Li, Naoyuki Kanda, Takuya Yoshioka, Xiong Xiao, Jian Wu, Long Zhou, Shuo Ren, Yanmin Qian, Yao Qian, Jian Wu, Michael Zeng, Xiangzhan Yu, and Furu Wei. WavLM: Large-scale self-supervised pre- training for full stack speech processing. IEEE J. Sel. Top. Signal Process., 1665):500--1518 2022. URL /ttttts:///oot www.eorrggl.0... 1119/..110////////.....1 JSTSP. 2022. 3188113
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arxiv:1904.10500 2019. URL http://arxiv..org
Sehyun Choi. Cross-architecture transfer learning for linear-cost inference transformers. arxiv htsp//b///240/ /240/ 02684
preprint arxiv:2404.026644 2024.
Krzysztof Marcin Choromanski, Valerïi Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Benjamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with per- formers. In 9/h International Conference on Learning Representations, ICLR 2021. OpenRe- view.net, 2021.
wwt.e/wwweeeeeeeeeeeeeeeeeeeeeee wow.foodduuuaaaaa 6zuk OWRH
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference Machine Learn- ing. ICML 2024. OpenReview.net, 2024. URL https:/ http:///ppeeeeviiie net/ wott/forustt forummic ztr td
URLhattss::/ppeeeeiiee

---

Published as a conference paper at ICLR 2025
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language ogy.
Technologies, NAACI-HLI 2019, pp. 4171-4186, 2019. URL. BTg/N19-1422 1423
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, and et al. The Llama 3 herd of models. arxiv preprint arxiv:2407.21188, 2024. URL BRLL:ttp:///rxii..ooo http:////wxxxvvoooo///////////// wtt../rrrrr//a////.1111.55 21783
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800GB dataset of diverse text for language modeling. arxiv preprini urxiv:2 1et:22210.0002... 2021. URL / http://axxxxv.ooaa///////// httpg/bb/221111 60027
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In Conference on Language Modeling (COLM), 2024.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-time models with linear state space layers. In Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, pp. 572- 585,20 2021.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The Tenth International Conference on Learning Representations, ICLR 2022. OpenReview.net, 2022.
Mutian He and Philip N. Garner. The interpreter understands your meaning: End-to-end spoken language understanding aided by speech translation. Findings of the Association for Compu- tational Linguistics: EMNLP 2023, pp. 4408 4423-442,,222 2023. URL htttpp::/aacaataatatt org/2023. findings emm.e.mmlppp...
François Hernandez, Vincent Nguyen, Sahar Ghannay, Natalia A. Tomashenko, and Yannick Estève. TED-LIUM 3: Twice much data and corpus repartition for experiments speaker adaptation. In International Conference on Speech and Computer (SPECOM) 2018, volume 1096 of Lecture Notes in Computer Science, pp. 198-208. Springer, 2018. URL https: http:://ddi..rrg////
10C7/978 3- 119--19-999 9957999999999 21
Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in neural network. arxiv preprint arxiv: (oxiv:/503.00 0/5500333311 2015. /tlL:ttt///xxx..o/////ba//
02531
Wei-Ning Hsu, Benjamin Bolte, Yao-Hung Hubert Tsai, Kushal Lakhotia, Ruslan Salakhutdi- nov, and Abdelrahman Mohamed. HuBERT: Self-supervised speech representation learning by masked prediction of hidden units. IEEE ACM Trans. Audio Speech Lang. Process., 29:3451-
3460, 2021. URL httpt:/// org/ 10. S119//TALL... 2021 3122291
Shankar Iyer, Nikhil Dandekar, and Kornel Csemmi. First quora dataset lease: Question pairs. 2017.
URL bttpp::/quoradatt.t quora. com/
rrat-quora-Datsset leale Quest. on-Pairs
Aref Jafari, Mehdi Rezagholizadeh, Pranav Sharma, and Ali Ghodsi. Annealing knowledge distilla-
tion. In Proceedings of the 16th Conference of the European Chapter of the Association for Com- putational Linguistics, EACL 2021, pp. 2493-2500 2021. URL pettps://aclanteen
erg/2021.ese eacl in. 212
Xilin Jiang, Cong Han, and Nima Mesgarani. Dual-path Mamba: Short and long-term bidirectional selective structured state space models for speech separation. arxiv preprint arxiv:2403. 18257, 2024. http::tts: //ar:iv///xxv wrrgor///aaooo abs ab/////aa//233... 2403.44. 18257
Jungo Kasai, Hao Peng, Yizhe Zhang, Dani Yogatama, Gabriel Ilharco, Nikolaos Pappas, Yi Mao, Weizhu Chen, and Noah A. Smith. Finetuning pretrained transformers into RNNs. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. EMNLP 2021, pp. 10630-10643, 2021. 830
URL https: http:///aaaathoootooooo wttp//looyooooo/oooooooooo /2211 manlp-main..
12

---

Published as a conference paper at ICLR 2025
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers RNNs: Fast autoregressive transformers with linear attention. In Proceedings of the 37th In- ternational Conference Machine Learning. ICML 2020, pp. 5165. PMLR, 2020. URL
ngs www.mmlm.....m.......mmm..m....m tts.s
wrwssressss... /vlss//vvii 9tttt//attaapaaae os20a .html
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th
International Conference on Learning Representations, ICLR 2020. OpenReview.net, 2020. URL
Nithin Rao Koluguri, Samuel Kriman, Georgy Zelenfroind, Somshubra Majumdar, Dima Rekesh, Vahid Noroozi, Jagadeesh Balam, and Boris Ginsburg. Investigating end-to-end ASR archi- tectures for long form audio transcription. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2024, pp. 13366-13370. IEEE, 2024. URL https: //doi. 10 1109/ Ct/.//CSSSSPPP4B 202.84455.2222.... 2024 10448309 http:///........................
Kai Li and Guo Chen. SPMamba: State-space model is all you need in speech separation. arxiv preprint arxiv:2404.02063 2024. |tttpss:: orq/ abs/ 2404 02063
Xingjian Li, Haoyi Xiong, Hanchao Wang, Yuxuan Rao, Liping Liu, and Jun Huan. Delta: Deep learning transfer using feature map with attention for convolutional networks. In 7th International Conference on Learning Representations, ICLR 2019. OpenReview.net, 2019. URL https:
Xuhong Li, Yves Grandvalet, and Franck Davoine. Explicit inductive bias for transfer learning with convolutional networks. In Proceedings of the 35th International Conference on Machine Learning. ICML 2018, pp. 2830-282800000 2839. PMLR, 2018. URL httpp://proceedings .mlr. press/v80 www.lall......
Yixiao Li, Yifan Yu, Qingru Zhang, Chen Liang, Pengcheng He, Weizhu Chen, and Tuo Zhao. LoSparse: Structured compression of large language models based on low-rank and sparse ap- proximation. In International Conference on Machine Learning, ICML 2023, pp. 20336-20350. mlr www. /tteess/v20/// wttp//i22........../../.. html
PMLR, 2023. URL htttps::/proceedingggg
Xi Lin, Zhiyuan Yang, Xiaoyuan Zhang, and Qingfu Zhang. Continuation path learning for homo- topy optimization. In International Conference on Machine Learning, ICML 2023, pp. 21288- 21311. PMLR, 2023. URL httpt::://proceediii ngs m.rr. prwss/v 2202 //in2nnnnnnnnmmmmm html
Andrew L. Maas, Raymond E. Daly, Peter T. Pham, Dan Huang, Andrew Y. Ng. and Christopher Potts. Learning word vectors for sentiment analysis. In The 49th Anmual Meeting of the As- sociation for Computational Linguistics: Human Language Technologies, ACL-HLT 2011, pp- 150, 2011.
Huanru Henry Mao. Fine-tuning pre-trained transformers into decaying fast weights. Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, EMNLP 2022, pp. 10236-10242, 2022. URL
wwtph/owooooo.../2222222//// org/2022 1a:eee2mmmmlpmmmammm mommm.in.....
697
Jean Mercat, Igor Vasiljevic, Sedrick Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearizing large language models. In Conference on Language Modeling (COLM), 2024.
Arsha Nagrani, Joon Son Chung, Weidi Xie, and Andrew Zisserman. VoxCeleb: Large-scale speaker verification in the wild. Comput. Speech Lang.. 60, 2020. URL ltttps://oi..cg//1
/016/j.csl.. 2019.101027
Antonio Orvieto, Samuel L. Smith, Albert Gu, Anushan Fernando, Çaglar Gülgehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. Interna- lional Conference on Machine Learning (ICML) 2023, pp. 26670-26699 PMLR, 2023. URL mir. press /222//2222//// 1ttpeess//202//vviee2222222 http://vveeeoaaaaa.....
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, and et al. RWKV: Reinventing RNNs for the transformer
In Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 14048-14077, 2023. URL https: httsss//aaaaaattttoaaooeo orr/ 2023. find emn emn 1p. 936
13

---

Published as a conference paper at ICLR 2025
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention. In 9th International Conference on Learning Representa- tions, ICLR 2021. OpenReview.net, 2021. httpt::://pppreeeeeeee net.eeew.eeeeeeeeeeeeeeeeeeeeeee wwr.moormmmm
Qt.KTdVFFRR.
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100, 000+ ques- tions for machine comprehension of text. In Proceedings of the 2016 Conference Empir- ical Methods in Natural Language Prucessing, EMNLP 2016, pp. 2383-2392, 2016. URL
https: / http://acaantthoooo www.://y...///1-1115550500
Dima Rekesh, Nithin Rao Koluguri, Samuel Kriman, Somshubra Majumdar, Vahid Noroozi, He Huang, Oleksii Hrinchuk, Krishna C. Puvvada, Ankur Kumar, Jagadeesh Balam, and Boris Ginsburg. Fast conformer with linearly scalable attention for efficient speech recognition. In IEEE Automatic Speech Recogmition and Understanding Workshop (ASRU) 2023, pp. 1-8. IEEE, 2023. URL
org/10..//...o.... 10. 1011//1000/9/66-0 2023. 10389701
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Trans. Assoc. Comput. Linguistics, 9:53-68, 2021. URL
fttpp:://do wtt.oo.orr//..111//aaaa...
1162/tacl a 00353
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arxiv preprint arxiv: 1910.01 108, 2019. URL
http:://axtt wwwg//aaa....ooooo.... abs 1910 01108
Jean Senellart, Dakun Zhang, Bo Wang, Guillaume Klein, lean-Pierre Ramatchandirin, Josep Maria Crego, and Alexander M. Rush. OpenNMT system description for WNMT 2018: 800 words/sec on single-core CPU. In Proceedings of the 2nd Workshop on Neural Machine Translation and Generation, IMT@ACL. 2018, pp. 122-128, 2018. URL https: /aass://aaaaattt lgy. www.ooooo...ooooooooo oorg/
218-27157
Siavash Shams, Sukru Samet Dindar, Xilin Jiang, and Nima Mesgarani. SSAMBA: Self-supervised audio representation learning with Mamba state space model. In IEEE Spoken Language Tech- nology Workshop (SLT) 2024, pp. 1053-1059, IEEE, 2024.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Y. Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, EMNLP 2013, pp. 1631-1642, 2013. URL https:://ccaattee
D13-11700
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for BERT model compression. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing, EMNLP- IJCNLP 2019, pp. 4322 4331, 2019. URL /tLLtttts///cccattoooo http:/ooooooyooo11111110111
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuging Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: successor to transformer for large language models. arxiv preprint arxiv:2307.08621 2023.
/2tp://iw..rr//abb/3.....00 08621
Raphael Tang, Yao Lu, Linqing Liu, Lili Mou, Olga Vechtomova, and Jimmy Lin. Distilling task- specific knowledge from BERT into simple neural networks. arxiv preprint arxiv: 1rxv::103.121336, 2019. URL httpltttp/aaaaaa
org/ abs 1903.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: survey. ACM Comput. Surv. 1566::199: 2023. URL hww.l:ttss/s/ec.m :/ooo..././oo// org/ 10. 1145/353081]
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R. Bowman. GLUE; A multi-task benchmark and analysis platform for natural language understanding. In 7th International Conference on Learning Representations, ICLR 2019. OpenReview.net, 2019, URL
htttps://openreviev. net
Junxiong Wang. Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The Mamba in the Llama: Distilling and accelerating hybrid models. In Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, 2024.
14

---

Published as a conference paper at ICLR 2025
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arxiv preprint arxiv:2006.07688, 2020. URL htttps://axxiv..ott abs/2005 04
Sarthak Yadav and Zheng-Hua Tan. Audio Mamba: Selective state spaces for self-supervised audio representations. arxiv preprint arxiv:2406.021175 2024. URL https://arxiv... org.iii.rrrggaae /bs/ 2406. 02178
Shu-Wen Yang. Ha-Han Chi, Yung-Sung Chuang, Cheng-I Jeff Lai, Kushal Lakhotia, Yist Y. Lin, Andy T. Liu, Jiatong Shi. Xuankai Chang. Guan-Ting Lin, Tzu-Hsien Huang. Wei-Cheng Tseng, Ko-tik Lee, Da-Rong Liu, Zili Huang, Shuyan Dong, Shang- Wen Li, Shinji Watanabe, Ab- delrahman Mohamed, and Hung-yi Lee. SUPERB: Speech processing universal performance benchmark. In 22nd Anmual Conference of the International Speech Communication Associa- tion, Interspeech 2021, pp. 1194-1198. ISCA, 2021. URL https: http:////i.o.o...././/1.///////.
Interspeech 20tppeeehh.2211111 1775
Songlin Yang, Bailin Wang, Yikang Shen. Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efficient training. In Forty-first International Conference on Machine Learning. JCML 2024. OpenReview.net, 2024. /RLlhttpps//opeereeeee net forum?
id=ia5XvxFUse
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big Bird: Transformers for longer sequences. In Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, 2020.
Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh M. Susskind. An attention free transformer. arXiv preprint arxiv:2105. 2021. URL
https://arxiv. org. wttporgggaba/..oaaa 14103
Michael Zhang, Kush Bhatia, Hermann Kumbong, and Christopher Ré. The hedgehog & the
porcupine: Expressive linear attentions with softmax mimicry. In The Twelfth International Conference on Learning Representationss, ICLR 2024. OpenReview.net, 2024a. URL https:
Xiangyu Zhang, Qiquan Zhang, Hexin Liu, Tianyi Xiao, Xinyuan Qian, Beena Ahmed, Eliathamby Ambikairajah, Haizhou Li, and Julien Epps. Mamba in speech: Towards an alternative to self- attention. arxiv preprini arxiv:2405. /2609, 2024b.
12609
A IMPLEMENTATION DETAILS
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 aKD 1 in ASR experiments. On other experiments, we find that the distillation on output probabilities actually leads to suboptimal results, and thus set aND O. We always use acE = 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 converge, generally around 30% of the total training steps.
For each target task, we train the models based on its common training recipe for standard transform- ers, and further search and determine hyperparameters and the model checkpoint 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 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
15

---

Published as a conference paper at ICLR 2025
Table 4: Hyperparameters for language processing tasks.
QNLI
QQP
SST2
IMDB
Std. RoBERTa
LR
3e-5
3e-5
3e-5
le-5
LR
le-5
le-5
3e-5
le-5
Unguided
Share mode
KV
Layer
KV
KV
LR
le-4
le-4
lc-4
3e-4
Target Guided
aLD
25
20
10
15
Share mode
Layer
Layer
KV
KV
LR
3e-4
lc-4
le44
3e-4
Traj. Guided
OLD
30
25
30
40
Share mode
Layer 5
KV
KV
KV
3
LR
le-4
le-4
lc-4
3e-4
Waypoint Guided
LLD
10
20
25
Share mode
KV
Layer
KV
KV
Language modeling: We train the model using learning rate of 3e-4 and 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 aLD 15 in the experiments. FP16 mixed precision training is adopted.
Speech processing: We train the model using 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 For ASR, IC, and SID, we set OLD 15, 15, and 30, respectively. BF16 mixed precision training is adopted. We find that these hyperparameters work well across different modes of distillation.
B DISTILLATION ALGORITHMS
Here we provide the pseudo-code to describe the training algorithms under different modes of distil- lation to facilitate understanding. Please refer Section for the notation scheme use, including the model output terms and loss functions.
C COMPUTATIONAL COSTS
The proposed method is presented process on the target architecture, as the model pretraining is notoriously expensive and generally infeasible with academic computation. For example, Wav2Vec2-large pretraining took 5.2 days on 128 V100 GPUs (Baevski et al. 2020), while our target guided CALD experiment to convert Wec2Vec2-largg into Mamba2 for ASR takes only 1.6 days on single RTX3090 with just 24GB memory. Without the computation on the teacher model, the unguided training will be roughly 40% faster, but the accuracy degradation is unacceptable. The hybrid approach is faster owing more than half of the training steps being under the unguided mode, while reaching better performance than the target guided one. The waypoint guided model doesn't involve any extra computation compared to the target guided one, and hence takes roughly the same time. In NLP experiments, the trajectory guided models enjoy better performance at the cost more computation. For example, on QNLI with layer-shi layer-shared Linformer projection, the target waypoint guided training takes around 3.4 hours under our training configuration with a single RTX3090, while the trajectory guided train- ing takes 3.5 hours (Ta, 5) to 4.0 hours (Ta, 1). The unguided training takes around 3.0 hours.
efficient alternative re-doing the whole pretraining
16

---

Published as a conference paper at ICLR 2025
Algorithm CALD training under the target guided or hybrid mode.
Require: mode: training mode (hybrid or target guided); Tą: number of guided steps in the hybrid mode; T: total number of training steps; dataloader: loader of training data; OCE, OKD,
ar.p: hyperparameters for loss term scaling.
1: Build and initialize the teacher model Mr and the student model Ms with the teacher target model.
2: Replace the attention layers in Ms.
3: Build the optimizer for Mg.
4: for i = to T do
5: (x. + next (dataloader)
6: (y()),(()) t Ms(x)
7: if mode Hybrid and i >> T, then
8:
L+ aCE (
9:
else
10:
(()) - Mr(x)
+ acE aCE ( ( y) + akd (xD(v (0) + ALD (
11:
end if
13:
[.backward()
Update Ms parameters using the optimizer.
15: end for
Algorithm 2 CALD training under the trajectory guided mode.
Require: T.: interval between teacher model updates; T: total number of training steps; dataloader: loader of training data; acE, AKD. aLD: hyperparameters for loss term scal- ing.
1: Build and initialize the teacher model My and the student model Ms with the teacher source model (the original pretrained model).
2: Replace the attention layers in Ms.
3: Build the optimizer for Ms and Mr.
4: for = T do
5: (x.8) < next (dataloader:
6:
(() + (x
7:
mod T, then
8:
Cx = (c
9:
br.backward
10:
Update Mt parameters using the optimizer.
11:
end if
12:
(y("H(()) + Ms(x)
13:
L+ acE LcE(y!"), y) aKD (*)()("), (()) + oLD (LD(H))),, (1(0)
14:
[.backward
15:
Update Ms parameters using the optimizer.
16:
end for
17

---

Published as a conference paper at ICLR 2025
Algorithm 3 CALD training under the waypoint guided mode.
Require: W: list of waypoint models; Tw: interval between waypoint switching; T: total number of training steps; dataloadez loader of training data; OCE. OKD, appi hyperparameters for loss term scaling.
1: Build and initialize the teacher model Mr with W1, t 1.
2: Build and initialize the student model Ms from the teacher source model (the original pretrained model).
3: Replace the attention layers Ms.
4: Build the optimizer for Ms.
5: for = 1 to T do
6:
(xr +- next (dataloader)
7:
(y("), H())) Ms(x)
8:
- Mr(x)
C+ OcE ( v) OKD LxD(y + oLD LLD(H()), H())
9:
10:
C.backward ()
Update Ms parameters using the optimizer.
12: i mod Te = 0 then
Wi W4 +1
14:
Re-initialize Mr with W
end if
16: end for
1600
Transformer
1400
1200
Time cast (md]
1000
200
Input audio length (6)
Figure 5: Time costs for Mamba2 and transformer based models performing ASR on audio samples with different lengths, averaged by 5 runs on single RTX3090.
To sum up. the CALD approach saves much computation compared re-pretraining. and the extra cost compared to the unguided training is limited.
We also provide brief evaluation the inference speed of transformer-based versus Mamba2- based ASR models we used in Section 4.4 as shown in Figure 5 With asymptotic linear time complexity, the Mamba2 model will be much faster given sufficiently long inputs compared the quadratic time transformers, although the actual speedup depends on the exact implementation; an investigation on this topic falls out of the scope of this paper.
18


Reviews:
[{"review_id": "VSM99PClel", "reviewer": "Reviewer_hriZ", "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. \n\nThe 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.", "questions": "Please see the above \"weakness\" section.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 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.", "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. \n\nThere 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.\" \n\nIt 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. \n\nFor 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1tT4vOaSJZ", "reviewer": "Reviewer_L6QU", "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.\nThe paper highlights that CALD effectively minimizes performance loss compared to directly converted models, demonstrating its efficacy across various language and speech processing tasks.", "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?\n\nQ2 - 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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n\n2.) 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.\n\n3.) 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.\n\n4.) 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.\n\n5.) The paper provides a solid theoretical explanation for why trajectory-guided distillation can retain pre-training knowledge and ensure better downstream task performance.", "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.\n\nComputational 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cwkLzV1N8H", "reviewer": "Reviewer_QBSN", "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.", "questions": "- How do you choose hyper parameters in equ. 7?\n- In Table 1, Src. init results are worse, why?\n- 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 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.\n- Multiple knowledge distillation methods are proposed and examined.\n- 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.", "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? \n- More detailed analysis about different distillation methods would be helpful.\n- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4Fyn58M9X9", "reviewer": "Reviewer_VgsV", "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).\nHowever, 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.\nMoreover, for non-text modalities such as speech, large scale Mamba 2 models are simply not available yet.\nHence authors propose leveraging off-the-shelf pre-trained (transformer) models and converting them into the target (mamba 2) model with linear complexity.\nMore 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.\nMultiple 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.", "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 ?\n\nQ: can we further speed-up inference through quantization (with added benefits)?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "-new approach to address model conversion/distillation from transformer (quadratic) architectures to mamba 2 (linear) architectures\n\n-not only applied to written language but also to speech with convincing results (3 types of tasks overall: NPL, LM and Speech)", "weaknesses": "-the obtained model is no longer a general purpose model (converted for a specific downstream task)\n\n-some inference speed evaluations would have been a plus (we know that obtained models with linear complexity should be faster though)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)