Discrete Diffusion Language Model for Long Text Summarization

A state-space model represents a system’s dynamics using a set of input, output, and state variables defined through linear differential or difference equations involving system matrices Brogan (1974); Gu et al. (2022); Fu et al. (2023). The model computes the output by applying the state and input variables to the output equation involving the system matrices. Mamba Gu and Dao (2023), which belongs to the family of state-space models, has demonstrated significant capability in handling long sequences across a wide range of application domains. For instance, VisionMamba Zhu et al. (2024) effectively leverages the Mamba kernel to encode images, achieving robust performance in image classification tasks. In the video domain, recent works Chen et al. (2024); Liu et al. (2024) demonstrate Mamba’s proficiency in managing image classification and complex spatiotemporal dynamics, offering both superior performance and favorable efficiency-performance trade-offs. In summarization task, we make the first attempt to integrate Mamba model to solve this complex language understanding task, competing with Transformer-based models.

The application of diffusion modeling to discrete data can be categorized into two main groups. The first group consists of methods that embed discrete structures into a continuous space and then apply Gaussian diffusion Chen et al. (2022); Dieleman et al. (2022); Gulrajani and Hashimoto (2024); Han et al. (2022); Li et al. (2022); Strudel et al. (2022); Lovelace et al. (2024).

Methods that define a diffusion process directly on discrete structures have greater potential for substantial improvements in speed. The D3PM framework Austin et al. (2021) introduces a Markov forward process by the multiplication of transition matrices over discrete time steps. Extending this framework to continuous time, as done in Eq. 1, utilizes continuous time Markov chain (CTMC) theory Campbell et al. (2022). The CTMC framework further generalizes the score-matching perspective on diffusion modeling Song and Ermon (2019) to discrete data Lou et al. (2024); Sun et al. (2022). Notably, SEDD Lou et al. (2024) integrates score-based approaches with ELBO maximization, allowing for effective likelihood-based training of score-based models.

Abstractive summarization involves compressing a longer input text into a shorter output summary that retains the essential information and main ideas using new phrases and sentences rather than simply extracting key phrases or sentences from the original content. Transformer-based models have dominated this field due to the ability to capture long-range dependencies and contextual relationships within the text, thanks to self-attention mechanism Liu and Lapata (2019); Lewis et al. (2019); Zhang et al. (2020). However, these models fail on long abstractive summarization benchmarks due to quadratic complexity of self-attention block, which limits the number of tokens these models can handle Keles et al. (2022). Consequently, recent works have attempted to address this issue by incorporating new attention mechanisms Guo et al. (2022); Zaheer et al. (2021). Our work tackles this problem by leveraging the linear time complexity of the Mamba model while also maintaining comparable performance with Transformer-based models on summarization benchmarks.

The D3PM framework Austin et al. (2021) introduces a Markov forward process $q(z_{t}|z_{t-1})=\text{Cat}(z_{t};Q_{t}z_{t-1})$ which is defined by the multiplication of matrices $Q_{t}$ over $T$ discrete time steps. This process results in the following marginal distributions:

These marginals represent the discrete-state form of equation 1. Specifically, each token in the sequence either remains unchanged or transitions to [MASK] with a certain probability $\beta$. The transition matrix at time step $t$ is defined as:

As the target sequence grows longer, the random noising process makes the conditional probability of generating tokens unpredictable. In DiffusionBERT He et al. (2023), the spindle noise schedule is introduced to estimate the probability that the $i$-th token remains unchanged at step $t$. This probability, denoted as $\overline{\alpha}_{t}^{i}$, is computed using the following equation $\overline{\alpha}_{t}^{i}=1-\frac{t}{T}-S(t)\cdot\widetilde{H}(x_{o}^{i})$ where $\widetilde{H}$ represents the entropy, which measures the information content of a random variable, $x_{i}$ denotes the $i$-th token in the sequence, and $n$ denotes the length of the sequence. However, this approach requires extracting the frequencies of words in the text corpus and does not have versatility across different tasks.

State Space Models (SSMs) are built on continuous systems that transform a 1D function or sequence, $x(i)\in\mathbb{R}^{L}$ into $y(i)\in\mathbb{R}^{L}$ through an internal state $h(i)\in\mathbb{R}^{N}$. Mathematically, SSMs utilize the following ordinary differential equation (ODE) to represent the input data:

where $A\in\mathbb{R}^{N\times N}$ is the system’s evolution matrix, and $B\in\mathbb{R}^{N\times 1},C\in\mathbb{R}^{N\times 1}$ are the projection matrices. This continuous ODE is typically discretized in modern SSMs. Mamba Gu and Dao (2023) represents a discrete variant of the continuous system, incorporating a timescale parameter $\Delta$ to convert the continuous parameters $A,B$ into their discrete forms $\tilde{A},\tilde{B}$. This conversion is generally done using the zero-order hold (ZOH) method, described by:

Mamba features a Selective Scan Mechanism (S6) as its primary SSM operator. The parameters $B\in\mathbb{R}^{B\times L\times N},C\in\mathbb{R}^{B\times L\times N},\Delta\in\mathbb{R}^{B\times L\times D}$, are directly derived from the input data $x\in\mathbb{R}^{B\times L\times D}$ as:

with $s_{B}(x)=\text{Linear}_{N}(x)$, $s_{C}(x)=\text{Linear}_{N}(x)$, $s_{\Delta}(x)=\text{Broadcast}_{D}(\text{Linear}_{1}(x))$, and $\tau_{\Delta}=\text{softplus}$, where $\text{Linear}_{d}$ is a parameterized projection to dimension $d$. The choice of $s_{\Delta}$ and $\tau_{\Delta}$ is motivated by their connection to RNN gating mechanisms.

CMLM Ghazvininejad et al. (2019) deploy a linear layer as a length predictor to predict the length of the target $L$ to avoid generating [PAD] tokens, and we utilize this predictor to adapt the cross-attention mechanism to create cross-Mamba. In detail, we first use Conv1d layers to compress the encoder’s output according to the ratio of max source length and max target length. Let $N$ be the length of the encoder’s output after compression, if $N<L$, we pad the sequence to the same length $L$; otherwise, we take the last $L$ tokens from the encoder’s output to create the matrices $B_{c}$ and $C_{c}$. The two matrices $B_{c}$ and $C_{c}$ are used to compute the target sequence in equation 6.

We set the number diffusion timestep $T$ in training to $T=50$ and inference for evaluation to $T=10$. We construct the encoder and decoder with 8 layers for each. Our model with the Transformer backbone has about 90M parameters and the Mamba backbone has roughly 85M parameters. We train the model using the AdamW optimizer Loshchilov and Hutter (2017) for 100,000 training steps, with a learning rate of $5\times 10^{-5}$. During the initial 10,000 steps, we employ a linear warmup schedule starting from a learning rate of $5\times 10^{-8}$. All experiments are conducted on 2 NVIDIA RTX 3090 GPUs and we use 1 for sampling.

Our quantitative results are presented in Table 1, showcasing ROUGE-1 (unigram), ROUGE-2 (bigram), and ROUGE-L (longest common subsequence) scores as done in prior text summarization work Lewis et al. (2019). Generally, all previous diffusion discrete models have been unable to conditionally generate sequences for the CNN/DailyMail dataset. In contrast, our proposed methods significantly outperform them, achieving improvements of up to 32 and 30 points in ROUGE-1 and ROUGE-L scores, respectively. Although semantic-aware noising continues to struggle with the Arxiv dataset, our Cross-Mamba method maintains consistent performance on this dataset, attaining respectable scores of 21.4 in ROUGE-1 and 20.4 in ROUGE-L.

In this section, we study how the semantic-aware noising process influences both the decoding stage and the training stage.

Figure 3 demonstrates that with the implementation of semantic-aware noising, the training process converges significantly faster on the QQP dataset compared to D3PM using random absorbing. At 20,000 training steps, the semantic-aware noising scheduler demonstrates performance comparable to that of random noising scheduler trained for 40,000 steps. Furthermore, at 40,000 training steps, it surpasses the random noising scheduler trained on 60,000 steps by a large margin in terms of BLEU score on QQP dataset. This finding suggests that discrete diffusion models can achieve enhanced performance through the development of appropriate generation strategies.

To understand more about the design of CrossMamba, we compared it with other prominent techniques that utilize RNN-based models, including seq2seq and Diffuseq. We chose the QQP dataset for this experiment because the paraphrasing task it presents is simpler compared to tasks like summarization. Table 4 demonstrates that our method excels at connecting the source and target sequences, and almost matches the attention mechanism whereas seq2seq suffers from an information bottleneck problem, and Diffuseq requires the model to reconstruct the input.

Intuitively, the attention mechanism computes a categorical distribution from $K,Q,V$ across the sequence, whereas Mamba’s $B$ and $C$ matrices are derived from the corresponding input tokens and encapsulate the sequence information into hidden states. Therefore, we hypothesize that Mamba’s kernels are more independent than the attention kernel, enabling it to perform better during random noise processing.

Without Similarity Loss: Without the similarity loss, there is no guarantee that the attention scores are consistent with the semantic meaning of the target and the noising process remains random, failing to dismantle the sequence in a structured manner. As shown in 6, removing similarity loss causes R-1 score drops by 6.6 points, R-2 score drops by 3.8 points, and R-L score drops by 5.8 points