Structured Masked Attention (SMA)

28 July 2025 · Download PDF · Download code

Masked Attention (MA) is given by the relation [1]

where $𝐋$ is a mask applied to the matrix $𝐁$ using the elementwise Hadamard product (denoted by $\odot$). When the mask is structured one can most often apply multiplication with $𝐌$ efficiently, and we refer to the masked attention as Structured Masked Attention (SMA). In the simplest case of $𝐋$ being a lower triangular matrix filled with ones the SMA reduces to

which can be viewed as a weighted cumulative sum. While the above is a structured computation, it is not efficient unless $𝐁$ itself has some structure. An example of when $𝐁$ is structured is when it is of low-rank (i.e., $𝐁=𝐔{𝐕}^{\top }$). The resulting SMA is a semiseparable matrix for which multiplication can be applied in linear time [2].

In general if $𝐋$ is a low-rank matrix (of rank $𝑝$) then multiplication with $𝐌$ scales as $𝑝$-times the scaling of multiplication with $𝐁$ as

for which we see that we need to perform $𝑝$ multiplications with $𝐁$ as well as $2𝑝$ diagonal multiplications.

1
using LinearAlgebra
2
using Random
3
using Test
4
Random.seed!(42)
5
n, p = 10, 2
6
U, V = randn(n,p), randn(n,p)
7
B = rand(n,n)
8
M = B .* (U*V')
9
x = randn(n)
10
@test M*x ≈ sum(i -> Diagonal(U[:,i])*(B*(Diagonal(V[:,i])) * x), 1:p)

1. State-space models as structured matrices

First we recap that a state-space model is given by the equations

An alternative way to view the state-space model is through the lens of structured matrices. We take ${𝐡}_0={𝐁}_0{𝐱}_0$, corresponding to a zero state before the first input. In particular a state-space model up until time step $𝑇$ can also be written using block matrices as well as block bidiagonal matrices as

We know that by simply iterating forwards in time we can compute the states in linear time. For the following generator representation, assume that each transition matrix ${𝐀}_{𝑖}$, $𝑖=1,\dots ,𝑇-1$, is invertible. A structured representation of the inverse of the bidiagonal matrix $𝐀$ can be constructed and applied in linear time for fixed state dimension, using

This type of matrix structure is called semiseparable [2]. Using this structured representation of ${𝐀}^{-1}$ we can compute the hidden states efficiently as

Similarly, the output $𝐲$ can be computed by applying the output matrices ${𝐂}_{𝑖}$ in a blocked fashion to the hidden states, i.e.

1
using SymSemiseparableMatrices
2
using BlockBandedMatrices
3
using LinearAlgebra
4
using Random
5
using Test
6
Random.seed!(42)
7
8
T = 10          # Sequence length
9
n = 6           # State size
10
input_dim = 1   # Dimension of forcing term
11
12
A_blks = [rand()*I(n) for _ in 1:T-1]       # Diagonal dynamics in Mamba 2
13
D_blks = [I(n)/1.0 for i in 1:T]            # Diagonal blocks are identity
14
C_blks = [rand(n,1) for _ in 1:T]           # Measurements are scalars
15
B_blks = [rand(n,input_dim) for _ in 1:T]   # Inputs are scalars
16
# Defining zeros blocks for the matrices
17
A_zero_blks = [zeros(n,n) for b in A_blks]
18
C_zero_blks = [zeros(n,1) for _ in 1:T-1]
19
B_zero_blks = [zeros(n,input_dim) for _ in 1:T-1]   
20
# Defining the block matrices
21
A = BlockTridiagonal(-A_blks,D_blks,A_zero_blks)
22
C = BlockTridiagonal(C_zero_blks,C_blks,C_zero_blks)
23
B = BlockTridiagonal(B_zero_blks,B_blks,B_zero_blks)
24
25
# Computing states by iterating forward
26
x_blks = [randn(input_dim) for _ in 0:T-1]  # input data
27
h_blks = [B_blks[1]*x_blks[1]]              # initial hidden state
28
for i in 2:T
29
    push!(h_blks, A_blks[i-1]*h_blks[i-1] + B_blks[i]*x_blks[i])
30
end
31
y_blks = [C'*h for (C,h) in zip(C_blks,h_blks)] 
32
33
# Computing states using semiseparable matrices
34
Ai_blks = [prod(reverse(A_blks[1:i-1]),init=1.0*I(n)) for i in 1:T]
35
U = vcat(Ai_blks...)
36
V = vcat(inv.(transpose.(Ai_blks))...)
37
# "SymSemiseparableCholesky" represents the matrix A^{-1} = tril(UV') 
38
Ai = SymSemiseparableCholesky(U', V')
39
x = vcat(x_blks...) # Collecting input data
40
41
@test Ai*(B*x) ≈ vcat(h_blks...)        # Checking hidden states
42
@test C'*(Ai*(B*x)) ≈ vcat(y_blks...)   # Checking measurement

2. State-space models as SMAs

In this section we aim to derive that the simplified SSM in the Mamba-2 paper is a special case of Structured Masked Attention (SMA) [1], i.e., that the multiplication

can be written differently as

The SSM in the Mamba-2 paper restricts the dynamics of $𝐀$ to be scalar-times-identity dynamics in order for the mask to be structured [1]. In short this restriction means that the dynamics for all hidden states are independent but equal.

From a practical point-of-view this means that we can collect each index of the dynamics and treat them separately. The resulting $𝐀$-matrix can be described by a Kronecker product (depending on how we organize the states it’s either $𝐀=𝐚\otimes {𝐈}_{𝑛}$ or $𝐀={𝐈}_{𝑛}\otimes 𝐚$). In the following we choose to separate the states, resulting in $𝐀$ having the form

For the following scalar generator representation, assume that ${𝑎}_{𝑖}\ne 0$ for $𝑖=1,\dots ,𝑇-1$. Using that the inverse of a Kronecker product is the Kronecker product of the inverses it follows that

where we further used that the inverse of a bidiagonal matrix is semiseparable. Furthermore, we have to re-arrange $𝐁$ and $𝐂$ which results in

where ${𝑏}_{𝑡}^{\left(𝑘\right)}$ and ${𝑐}_{𝑡}^{\left(𝑘\right)}$ denote the $𝑘$th entries of ${𝐛}_{𝑡}$ and ${𝐜}_{𝑡}$. The final multiplication will therefore look as

Finally, using the properties of the Hadamard product we get to the Structured Masked Attention form that we were looking for

This means that the SSM dynamics can be interpreted as a structured masked attention mechanism. Note that in the case of the dynamics being independent but different (i.e., we would have ${𝐚}^{\left(𝑘\right)}$ rather than just $𝐚$) the SSM dynamics would result in a sum of masked attentions, i.e.,

1
using SymSemiseparableMatrices
2
using LinearAlgebra
3
using Random
4
using Test
5
Random.seed!(42)
6
7
T = 10          # Sequence length
8
n = 6           # State size
9
10
# Here we treat the matrices in terms of their states and not sequence lengths
11
a_blks = rand(T-1)
12
a = Bidiagonal(ones(T), -a_blks, :L)
13
A = kron(I(n),a)        # Kronecker product with identity
14
15
# The blocks are now size equal to the sequence and a block for each state!
16
B_blks = [rand(T) for _ in 1:n]
17
C_blks = [rand(T) for _ in 1:n]  
18
# Collecting the blocks into the B and C
19
B = vcat(Diagonal.(B_blks)...)
20
C = vcat(Diagonal.(C_blks)...) 
21
22
# We want to see if the full matrix M = C'*(A\B) can be written as
23
# structured masked attention, i.e.,
24
# M = (CB')\circ a^{-1}. For this we start by computing a^{-1}
25
ai = inv(a) # We ignore here that inv(a) is semiseparable
26
@test C'*(A\B) ≈ sum(i-> Diagonal(C_blks[i])*ai*Diagonal(B_blks[i]), 1:n)
27
@test C'*(A\B) ≈ sum(i->(C_blks[i]*B_blks[i]') .* ai, 1:n)
28
29
# We can collect the terms and write it as Structured Masked Attention!
30
Cn = hcat(C_blks...)
31
Bn = hcat(B_blks...)
32
@test C'*(A\B) ≈ (Cn*Bn') .* ai
33
34
# We can apply the semiseparable structure of the inverse when multiplying!
35
ai_blks = [prod(a_blks[1:i-1],init=1.0) for i in 1:T]
36
u = vcat(ai_blks...)
37
v = vcat(inv.(ai_blks)...)
38
# "SymSemiseparableCholesky" represents the matrix A^{-1} = tril(UV') 
39
ais = SymSemiseparableCholesky(u',v')
40
# Efficient products using the structure of "ais" and diagonal B and C
41
x = randn(T)
42
@test C'*(A\(B*x)) ≈ sum(i-> C_blks[i].*(ais*(B_blks[i].*x)), 1:n)

@misc{paltorp2025structuredmaskedattention,
  author = {Paltorp, Mikkel},
  title = {{Structured Masked Attention (SMA)}},
  year = {2025},
  month = jul,
  url = {https://mipals.github.io/pdf/2025/07-28-structured-masked-attention/sma.pdf},
  note = {Blog post}
}