Fast Diagonal Convolutions (in JAX)

Why Convolutions?

In this article, we examine convolutions as discrete matrix operations commonly employed in array manipulation programs. These operations are fundamental to many modern machine learning techniques. The convolution operation, denoted by the symbol 🞶, operates between an input matrix A and a convolutional kernel k.

Detailed discussions of matrix convolutions can be found in various sources, including this article.

Our focus is specifically on 2D diagonal convolutions, which involve only the nonzero elements along the diagonal of the convolutional kernel. These are effectively 1D convolutions applied along the diagonals of A and are particularly useful in time-series models. Due to their 1D nature, they can be computed significantly faster than conventional 2D convolutions — a point we briefly explore.

2D Convolution Computations

There are two equivalent methods to compute the convolution. The first method involves applying the convolutional kernel — using element-wise multiplication — to sub-matrices of the input matrix that are of equal size, and then summing the resulting products.

In the figure above we explicitly compute a 6×6 matrix, however its 3×3 blocks are summed over yielding the 2×2 matrix to the right.

This can be computed simply using a JAX-adapted scipy’s convolve2d.

import jax.numpy as jnp
from jax.scipy.signal import convolve2d

A = jnp.arange(8*9).reshape(8,9) ## example matrix
convolve2d(A, jnp.eye(4), mode='valid')

Alternatively, we can employ Flax’s native CNN layer — built on JAX — to perform this operation, thereby extending convolve2d to support both striding and padding.

import jax
import jax.numpy as jnp
import flax.linen as nn
import optax
from flax.training import train_state
from jax.scipy.signal import convolve2d

def create_train_state(input_shape, kernel:jax.Array, stride:int=1, padding="valid"):
    conv_layer = nn.Conv(features=1, kernel_size=kernel.shape, padding=padding, strides=stride, use_bias=False)
    params     = conv_layer.init(jax.random.PRNGKey(0), jnp.ones(input_shape))['params']
    params['kernel'] = jnp.flip(kernel)[:, :, None, None]
    return train_state.TrainState.create(apply_fn=conv_layer.apply, params=params, tx=optax.adam(0.001))

def nn_convolve2d(IN, kernel:jax.Array, stride:int=1, padding="valid"):
    input_array = IN[None, :, :, None]
    state = create_train_state(input_array.shape, kernel, stride=stride, padding=padding)
    return jnp.squeeze(state.apply_fn({'params': state.params}, input_array)).astype(IN.dtype)

### Example Matrix and convolution size
c = 4
A = jnp.arange(8*9).reshape(8,9)

### comparison: flax & jax.scipy
flax  = nn_convolve2d(A, jnp.eye(c), stride=1, padding=0)
scipy = convolve2d(A, jnp.eye(c), mode='valid').astype(int) 
print( jnp.allclose( flax , scipy ) )

As demonstrated, both of these programs consider the case of dense convolutional-kernel matrices, however in our case, diagonal matrices, e.g. jnp.eye(4) above, we have to sum over many zeros.

Alternatively, one can extract sub-matrices from the input matrix — each sized to match the output matrix — and perform element-wise multiplication with the corresponding elements of the convolutional kernel. In our case, the convolutional kernel is diagonal with three elements; thus, we sum three matrices to obtain the final result, as shown in the figure below.

Therefore, we implement the convolution using only the diagonal sub-matrices, reducing the algorithmic time complexity from O(c²N²) to O(cN²). An example JAX implementation for a general matrix A is shown below.

import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.jit, static_argnums=(2,))
def diag_convolve2d(A:jax.Array, k:jax.Array, stride:int=1):
    """
    Diagonal-kernel convolution with diagonal components given by k
    only sum over the diagonal-intersection, for a given stride.

    Arguments:
      A:jax.Array , 2d-initial Array
      k:jax.Array , 1d-diagonal-kernel (e.g. from a purely diagonal 2d-Array)
      stride:int  , stride
    Returns:
      B:jax.Array , 2d-diagonally convolved Array
    """
    N, M    = A.shape
    s       = stride
    c       = k.size

    x       = int((N - c) // s + 1) ## result size x (hori.)
    y       = int((M - c) // s + 1) ## result size y (vert.)
    end_x   = N - ((N-c) % s)
    end_y   = M - ((M-c) % s)
    start_x = end_x - (1 + (x-1)*s)
    start_y = end_y - (1 + (y-1)*s)

    B       = k[0] * A[ start_x:end_x:s, start_y:end_y:s ]
    for i in range(1, c):
        start_x-=1
        start_y-=1
        end_x  -=1
        end_y  -=1
        B  += k[i] * A[ start_x:end_x:s, start_y:end_y:s ]
    return B

### Example: matrix and convolution size:
c = 4
A = jnp.arange(8*9).reshape(8,9)

### comparison: flax & diag
diag = diag_convolve2d(A, jnp.ones(c), stride=2)
flax = nn_convolve2d(  A, jnp.eye(c) , stride=2)
print(jnp.allclose(diag, flax))

Additionally, if the input matrix can be factored (decomposed) as a product of two matrices, i.e., A = Q @ Q† (where @ denotes the matrix product and † represents the matrix-transpose), then the convolution can be computed without explicitly forming A. This is particularly relevant for low-rank and large matrices and is especially useful in LinearOperator or matrix-free methods. In this formalism, the computational complexity is reduced to linear, O(cNr), where r is the rank of both Q and A.

import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.jit, static_argnums=(3,)) ### k defaults to 1s?
def diag_convolve2d_LO(Q:jax.Array, R:jax.Array, k:jax.Array, stride:int=1):
    """ Diagonal-kernel convolution with diagonal components given by k
    only sum over the diagonal-intersection, for a given stride.

    Arguments:
      A:jax.Array -- 2d-initial Array
      k:jax.Array -- 1d-diagonal-kernel (e.g. from a purely diagonal 2d-Array)
      stride:int  -- stride
    Returns:
      B:jax.Array -- 2d-diagonally convolved Array
    """
    N       = Q.shape[0]
    M       = R.shape[0]
    s       = stride
    c       = k.size

    x       = int((N - c) // s + 1) ## result size x (hori.)
    y       = int((M - c) // s + 1) ## result size y (vert.)
    end_x   = N - ((N-c) % s)
    end_y   = M - ((M-c) % s)
    start_x = end_x - (1 + (x-1)*s)
    start_y = end_y - (1 + (y-1)*s)

    B       = k[0] * Q[ start_x:end_x:s, :] @ ( R.T[:, start_y:end_y:s ] )
    for i in range(1, c):
        start_x-=1
        start_y-=1
        end_x  -=1
        end_y  -=1
        B  += k[i] * Q[ start_x:end_x:s, :] @ ( R.T[:, start_y:end_y:s ] )
    return B

### Example: matrix-factor and convolution size:
N, M = 80, 20
c=4

key = jax.random.key(1701)
Q   = jax.random.normal(key, shape=(N,M))
k   = jax.random.normal(key, shape=(c,))

### comparison: flax & diag
LO_diag = diag_convolve2d_LO(Q, Q, k, stride=3)
flax    = nn_convolve2d(Q @ Q.T , jnp.diag(k), stride=3)
print(jnp.allclose(LO_diag, flax, atol=1e-5))

Kernel-factorization

Matrix factorization of the input matrix is useful, but what about factorizing the convolutional kernel? It is well known that for dense and low-rank convolutional kernels, the 2D convolution operation reduces to sequential 1D convolutions along each direction. In the case of a rank-1 tensor-product decomposition, this reduction follows the equation below.

However, for us diagonal convolutional-kernels are of naïvely of full-rank, and therefore such an approach would seem hopeless. However, there is a niffy-trick! For simplicity, let’s first consider the case of a identity-matrix convolutional-kernel. That is we consider the convolutional-kernel as being the reshaped tensor-product of a ones vector 1, of length of some chosen base b. The shape of k is thus an bⁿ × bⁿ.

from functools import partial
import jax
import jax.numpy as jnp

@partial(jax.jit, static_argnums=(1, 2, 3))
def matrix_sum(Z:jax.Array, a:tuple, A:int, B:int):
    C = Z[a[0] : a[0] + A, a[0] : a[0] + B]
    for i in range(1, len(a)):
        C += Z[a[i] : a[i] + A, a[i] : a[i] + B]
    return C

def i_convolve2d(Z, n, base:int=2):
    """ convolve: Z by jnp.eye( base **n )
    Arguments:
        Z : jax.Array , Input Matrix
        n : jax.Array , power (number of generations)
        base : int , number of convolutions before exponentiation
    Returns:
        B:jax.Array , 2d-diagonally convolved Array
    """

    for generation in range(1, n+1):
        A = int( Z.shape[0] - (((base)**(generation-1))*(base-1)+1) + 1 ) 
        B = int( Z.shape[1] - (((base)**(generation-1))*(base-1)+1) + 1 )
        a = ((base)**(generation-1))*jnp.arange(base) ### upper-left most part of submatrix

        Z = matrix_sum(Z, tuple(a.tolist()), A, B) ### this shrinks Z
    return Z

### Example Matrix and convolution size:
Z    = jax.random.randint(key, (25,25), 0, 10)
n    = 2
base = 3

### comparison: flax & diag
flax = nn_convolve2d(Z , jnp.eye(base**n), stride=1)
icon = i_convolve2d( Z, n, base=base)
print( jnp.allclose(icon, flax) )

Incredible, we are exponentially faster than before, using O(log(c)N²) operations instead of O(cN²)! However, we are restricted to convolutions (convolutional-kernel matrix) of the identity matrix in certain sizes. This is due by our choice of kernel factorization. Let’s now instead consider the case for a collection of 1D factors, that may be outer-producted-and-reshaped into the diagonal of our convolutional-kernel matrix.

To concretely implement this, we begin by defining a function to compute the n-fold outer product of a list of arrays. In our case, these arrays correspond to the outer-product factors of the convolutional kernels. The function is provided below.

from typing import Union, List
import jax.numpy as jnp

@jax.jit
def outers(k:list[jax.Array]):
    out = jnp.ones(1)
    for i in range(len(k)):
        out = jnp.outer(out, k[i]).T.reshape(-1)
    return out

### Example:
k = [jnp.array([1, 2, 3]), jnp.array([7,5]), jnp.array([6,2,1])]

print( outers(k) )

Now we must modify our identity-matrix function, i_convolve2d, to instead include jnp.cumprod as replacement for the cumulative base factor count, and also to include jnp.flip (likely due to the way our custom function outers reshapes the tensor-product) to flip each convolutional-kernel factor. In addition, to other minor modifications we obtain the following JAX code.

from functools import partial
import jax
import jax.numpy as jnp

@jax.jit
def outers(k:list[jax.Array]):
    out = jnp.ones(1)
    for i in range(len(k)):
        out = jnp.outer(out, k[i]).T.reshape(-1)
    return out

@partial(jax.jit, static_argnums=(1, 2, 3, 4))
def matrix_sum(Z:jax.Array, a:tuple, A:int, B:int, cc:jax.Array):
    C = cc[0] * Z[a[0] : a[0] + A, a[0] : a[0] + B]
    for i in range(1, len(a)):
        C += cc[i] * Z[a[i] : a[i] + A, a[i] : a[i] + B]
    return C

def d_convolve2d(Z, k):
    """ convolve: Z by jnp.diag( ⊗k )
    Arguments:
        Z : jax.Array , Input Matrix
        k : List[jax.Array] , List of 1d-diagonal-kernel factors
    Returns:
        B:jax.Array , 2d-diagonally convolved Array
    """

    bs = jnp.cumprod(jnp.array([[1] + [ len(k[i]) for i in range(len(k))]]))

    for generation in range(len(k)):
        c = jnp.flip( k[generation] )
        b = len(c) ## base

        A = int( Z.shape[0] - ((bs[generation])*(b-1)+1) + 1 )
        B = int( Z.shape[1] - ((bs[generation])*(b-1)+1) + 1 )
        a = (bs[generation])*jnp.arange(b) ### upper-left of submatrix

        Z = matrix_sum(Z, tuple(a.tolist()), A, B, tuple(c.tolist())) ### iteratively shrinks Z
    return Z

### Example matrix & convolution-kernel-factors:
Z = jax.random.randint(key, (150,160), 0, 10)
k = [jnp.array([5, 2, 33]), jnp.array([3, 1, 17]), jnp.array([1, 2]), jnp.array([1, 2, 3, 8])]

### comparison: flax & d_convolve2d
flax = nn_convolve2d(Z, jnp.diag(outers(k)), stride=1)
dcon = d_convolve2d( Z, k)
print( jnp.allclose(dcon, flax) )

And there we have it. Now we can compute arbitrary diagonal-convolutions in O(N²log c) time as long as they do not have prime dimensions!

A. Honorable Mention: FFT-based Convolutions

In addition to using clever summations for computing diagonal convolutions, an alternative strategy involves element-wise multiplication of the discrete Fourier transforms (FFT) of the input matrix and the convolutional kernel. The time complexity of this approach scales as O(N²log⁡N) for matrices of size N × N. An example implementation in JAX is provided below, demonstrating that for c = N, this method becomes equivalent in time complexity to our previously discussed algorithm.

import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.jit, static_argnums=(2,3))
def fft_convolve2d(a, k, stride:int=1, padding:int=0):
    """ 2D FFT-based convolution with striding and padding.
    Parameters:
      a: jax.Array, The input array.
      k: jax.Array, The convolution kernel.
      stride: int, default=1, The stride (downsampling factor) for the output.
      padding: int, default=0, The number of zeros to pad on all sides of the input.

    Returns:
      A jax.Array containing the convolution result.
    """

    k1, k2   = k.shape ## kernel dimensions

    ## pad the input array by `padding` on each side
    a_padded = jnp.pad(a, ((padding, padding), (padding, padding)))
    H, W     = a_padded.shape

    ## Compute the output size for the full convolution
    output_shape = (H + k1 - 1, W + k2 - 1)

    ## Compute FFTs for: padded-input & convolution-kernel
    A = jnp.fft.fft2(a_padded, s=output_shape)
    K = jnp.fft.fft2(k, s=output_shape)

    ## Compute the full convolution via iFFT
    conv = jnp.real(jnp.fft.ifft2(A * K))[k1 - 1: H, k2 - 1: W]

    return conv[::stride, ::stride] ## apply striding

### Example Matrix and convolution size:
c = 4
A = jnp.arange(8*9).reshape(8,9)

### comparison: flax & fft
flax = nn_convolve2d( A, jnp.eye(c), stride=2, padding=2)
fft  = fft_convolve2d(A, jnp.eye(c), stride=2, padding=2)
print( jnp.allclose( flax , fft , atol=1e-2) )