State-Space Large Audio Language Models

Structured state space sequence models (S4) [16] are inspired by classical state-space models such as Kalman filters and Hidden Markov Models. The state-space models map a 1-D sequence $x(t)\in\mathbb{R}\rightarrow y(t)\in\mathbb{R}$ through a hidden state $h(t)\in\mathbb{R}^{N}$ via linear ordinary differential equations as follows:

where $\mathbf{A}\in\mathbb{R}^{N\times N}$, $(\mathbf{B}\in\mathbb{R}^{N\times 1},\mathbf{C}\in\mathbb{R}^{1\times N})$ are called the evolution and projection parameters respectively.

A discretization step transforms the continuous parameters, $\mathbf{A},\mathbf{B}$ to discrete parameters $\bar{\mathbf{A}},\bar{\mathbf{B}}$. A commonly used method for discretization, zero-order hold uses a timescale parameter $\Delta$ to discretize as follows:

After the discretization step, the state-space equations can written as:

The current view of state-space models is analogous to RNNs where the output of the current time step only depends on the previous hidden state and current input. Since the state-space parameters are not time or input-dependent, SSM can also be viewed as a convolution $y_{t}=(x_{0},x_{1},...,x_{t})*(\mathbf{C}\overline{\mathbf{B}},\mathbf{C}\overline{\mathbf{A}\mathbf{B}},...,\mathbf{C}\overline{\mathbf{A}}^{M-1}\overline{\mathbf{B}})=\mathbf{x}*\overline{\mathbf{K}}$ where $\overline{\mathbf{K}}$ is the called the global convolutional kernel.

SSMs can either be viewed as CNNs or RNNs depending on the task. During training, the convolutional view is used to enable parallel training. During inference, the recurrent view allows faster inference and unbounded context. However, the linear time-invariant nature of these models limits their performance on content-based reasoning tasks.

Gu et al. [17] proposed a parametrization method to make the state-space parameters input-dependent. However, this breaks the convolutional view and poses a challenge for efficient computation. To address this, a hardware-aware parallel algorithm is used to efficiently compute the output.

DASS [19] was one of the first attempts to use a pure state-space-based model to classify audio signals. DASS uses AST [31], a transformer-based teacher model, to guide and train a state-space audio classifier. DASS combines the best of both worlds: it outperforms the transformer-based models and retains the computational advantages of state-space models.

The DASS model can be divided into four groups and each group consists of a state-space block and the first three groups also contain a downsampling layer. Each group progressively reduces the sequence lengths and increases the feature size. A pooling method generates the final embedding that summarizes the input spectrogram. DASS shows remarkable duration scalability: even a model trained with ten-second utterances can infer information from hour-long audio.

In this work, we use DASS pretrained on AudioSet-2M [35] dataset with the classification layer removed as the audio features extractor for the input audio. It takes a spectrogram of size 1024*128 as input and generates a feature map of size 32*4*768 as output. To further reduce the spatial dimension, we use a two-dimensional convolution with a kernel size of 3 and stride 2 and then use a linear layer to map the features from 768 dimensions to the input size of the language model i.e. 4096 for the LLaMa and 2560 for the state-space based LLM.

We use the following LLMs in this work:

LLaMA: We follow LTU and use Vicuna instruction tuned LLaMA-7B LLM [36]. LLaMa is pretrained on large amounts of natural language and code corpora in a self-supervised manner. Vicuna [37] is trained on instruction-following language prompts generated by GPT models which improves the models’ performance on reasoning and generation tasks.

State-space LLM: State-space LLM-2.8B [17] trained on the Pile dataset. The state-space LLM outperforms similar-size transformer-based models such as GPT-Neo 2.7B.

Low-rank Adapters: Instead of finetuning all the weights of LLMs on our task, we use Low-rank (LoRA) [38] adapters to finetune the LLMs. LoRA adds a small set of learnable weights on top of the pre-trained weights from the LLM. The learnable weights can be decomposed into a product of two low-rank matrices. This allows us to modify the large parameter matrices of the LLMs without adding a lot of learnable parameters. The final LLM parameters are the addition of frozen parameters and the low-rank learnable matrices.

For LLaMa, we add LoRA adapters (rank=8 and $\alpha$=16) to the key and query projection layers in all the self-attention layers of LLaMa models. This step introduces 4.2M learnable parameters. For state-space LLM, we add LoRA adapters (rank=8 and $\alpha$=16) to the input projection layers of the state-space block. This step adds 6.5M learnable parameters.

We train our models using the next token prediction task conditioned on the input audio and past tokens. We maximize the following probability, $P(x_{t}|x_{1:t-1},A)$, by using cross-entropy loss for all the text tokens $1<t\leq T$ in the input text tokens and reference audio. For generation, we use the following settings: Temperature=0.1, Top K=500, and Top P=0.95 with a repetition penalty of 1.1.

We train our models on OpenAQA dataset [9]. This dataset contains tuples of Audio, question, answer where the models take Audio and question as input and generate answer as the output.

We follow the same training pipeline from LTU for training the models proposed in this paper. We use 4× RTX A6000 GPUs for training the models. The hybrid LALM is trained for about 3 days. For the state-space LALM, we can increase the batch size from 4 to 16. We use gradient accumulation to ensure the effective batch size is the same i.e. 256. This speeds up the training and allows us to train the model in less than two days. However, the reason for the training speed up is unclear: it could be simply due to the state-space LLM being smaller or the state-space model being computationally efficient. In the future, when larger state-space LLMs become publicly available, we plan to increase the model size and compare the computational efficiencies of the state-space and Transformer-based LALMs.

The small and medium hybrid-LALM contains 42M and 61M trainable parameters out of 6.8B parameters respectively. The small and medium ssLALM contains 43M and 62M trainable parameters out of a total of 2.8B parameters respectively.

To reason and understand the audio, the LALM must be able to first recognize the input audio. One important step in bench-marking LALMs is the performance comparison on close-ended tasks where the output labels are predefined. We follow the evaluation pipeline from LTU [9] and compare our models with existing LALMs on 8 audio classification benchmarks and 2 audio captioning benchmarks.

Audio Classification: LALMs do not directly predict the class index but instead output the audio label names or descriptions. To compute the performance of these models, we first encode the LALM output and the evaluation label using a text encoder and then we compute a cosine similarity between the LALM output and the label. For single-label classification tasks, we use the label with the highest similarity score and compute accuracy or F1-score and for multi-label classification tasks, we use the cosine similarity as the prediction score and compute the mAP. We used the prompt “write an audio caption describing the sound” for the classification tasks.

Audio Captioning: For the caption generation tasks, we use the prompt “write an audio caption describing the sound” and take the LALM output as the prediction. We AudioCaps and Clotho datasets and use SPICE as the evaluation metric.

As seen in Table 1, our proposed model outperforms the other existing models such as SALMONN, Pengi, and AudioGPT. For SALMONN [8], Pengi [7], AudioGPT [11], we use the results reported in the GAMA paper. For both the audio classification and audio captioning task our proposed models outperform the existing models in most of the datasets and overall average performance. For the audio captioning task, our models perform similarly to the best-supervised systems.

Our proposed models match the performance of LTU [9] which our approaches build upon. The audio encoder i.e. AST used in LTU is first pretrained on audio-visual data and then finetuned on the AudioSet-2M dataset. In contrast, the DASS model is trained on the audio data from AudioSet-2M. DASS trained on only AudioSet-2M outperforms and shows more robustness and duration-scalability than AST trained on the same dataset [19].

The state-space LLM used in ssLALM is much smaller and is not instruction-trained unlike LLaMA used in hybrid LALMs or transformer LALMs such as LTU. Despite being smaller and having the audio encoder and LLM trained on less data, the ssLALMs perform competitively with the best transformer-based LALMs.

Although all the LALMs, including the ssLALMs, perform poorly on multi-label classification task on AudioSet. We believe it is because LALMs mainly predict the prominent class and underestimate the likelihood of non-prominent sound classes in the input audio which results in low mAP scores. We also show some open-ended question-answering abilities of the models in Table 2. All the models can infer the information from audio and generate reasonable answers.