SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Matthew Lyon
matthew.s.lyon@.manchester.ac.uk
Paul Armitage
p.armitage@sheffield.ac.uk
Mauricio A Alvarez
maurizio.alvarezlopez@manchester.ac.uk
Diffusion MRI (dMRI) is a widely used imaging modality, but requires long scanning times to acquire high resolution datasets. By leveraging the unique geometry present within this domain, we present a novel approach to dMRI angular super-resolution that extends upon the parametric continuous convolution (PCCConv) framework. We introduce several additions to the operation including a Fourier feature mapping, global coordinates, and domain specific context. Using this framework, we build a fully parametric continuous convolution network (PCCNN) and compare against existing models. We demonstrate the PCCNN performs competitively while using significantly fewer parameters. Moreover, we show that this formulation generalises well to clinically relevant downstream analyses such as fixel-based analysis, and neurite orientation dispersion and density imaging.
Diffusion magnetic resonance imaging (dMRI) is a clinical imaging modality that is routinely used to diagnose white matter pathology within the human brain. Advancements in downstream analyses using dMRI data continuously expand the diagnostic capabilities of the modality, but often require high resolution data that is clinically infeasible to acquire [28]. However, recent research demonstrates that deep learning can be used to predict high resolution images using low resolution dMRI data [26, 44, 41], a technique known as super-resolution. These advancements therefore pave the way for clinically accessible advanced dMRI analyses.
dMRI data can be represented using a non-Euclidean geometry. Specifically, a dMRI dataset consists of a series of three-dimensional volumes where each volume measures the water diffusivity in a given direction. Here, the direction is quantified by a 3D vector known as the 'b-vector', and is sparsely sampled from a unit sphere. Each three-dimensional volume, however, is regularly sampled in a discrete grid. Therefore, dMRI data has two distinct resolution types: spatial resolution which denotes the regularly sampled grid size, and angular resolution which denotes the number of diffusion directions. Despite this relatively unique geometry, most dMRI deep learning methods use operations grounded in a Euclidean framework, such as convolutional neural networks (CNNs) [43, 44, 41].
This presents an opportunity as typical CNN architectures do not fully utilise the geometric properties present in dMRI data. For example, implicit within the formulation of the CNN is the assumption that data are densely and regularly sampled in a discrete manner. This is only true when considering the three spatial dimensions of dMRI data, whilst the other dimensions would be more suited using approaches like graph convolutional networks (GCNs) [46], spherical CNNs [7], and point cloud CNNs [22]. Examples of approaches that develop geometrically motivated convolutions in the dMRI
domain include those for tissue segmentation [21], lesion segmentation [29], and fibre orientation distribution (FOD) reconstruction [5, 32].
Another promising candidate for a framework capable of utilising the unique dMRI geometry is the parametric continuous convolution (PCConv) introduced by Wang et al. [39], which extended the discrete convolution into the continuous domain. This was achieved via an inner parameterised hypernetwork [16] responsible for generating convolutional weights. Here, the convolutional weights were not learnt directly but were instead sampled from the hypernetwork given kernel coordinates as input. This operation could therefore be used to model arbitrary geometries.
In this work, we build upon the PCConv operation to convolve across both spatial and angular dimensions of dMRI data via a new, flexible, coordinate embedding. We subsequently create a deep parametric continuous CNN (PCCNN) that is well-suited for inference within this domain. Furthermore, we develop additional modifications to the PCCNN by including supplemental prior information in the coordinate embedding. These modifications include the application of Fourier feature mappings, the inclusion of prior domain information, and the incorporation of global coordinate components.
We demonstrate the effectiveness and flexibility of this approach through the task of angular superresolution of dMRI data in various sampling schemes, including single-shell and multi-shell data [18]. To assess the performance of our method, we compare results across clinically relevant downstream analysis tasks such as fixel-based analysis (FBA) [30], and neurite orientation dispersion and density imaging (NODDI) [45]. We show that the PCCNN performs competitively against a similar recurrent CNN (RCNN) network trained on the same data whilst using approximately a tenth of the number of parameters. Additionally, we demonstrate the generalisability of the network by comparing it against models trained specifically for downstream analysis tasks.
Whilst there is a body of literature within spatial super-resolution methods in the natural image domain [40], and more specifically the dMRI domain [1, 36, 6, 25], this work focuses on angular superresolution as a means of increasing image quality. Here, the task of inferring high angular resolution data from low angular resolution dMRI acquisitions can broadly be split into two categories: methods that super-resolve raw dMRI data, and those that super-resolve downstream analyses. The former allows for more flexible use of the data and analysis method, but the unconstrained nature of the task makes inference more challenging. The latter constrains the problem by leveraging assumptions made within the analysis method, however the mapping from raw dMRI data to downstream analysis is in general not invertible, therefore the super-resolved data cannot be used to derive other downstream analysis metrics.
Notable works in angular super-resolution of raw dMRI data include Yin et al. [43], whereby a preliminary study constructed a 1D CNN autoencoder that only convolved across the angular dimension, disregarding possible spatial correlations within the data. Lyon et al. [26] developed a 3D RCNN that benefited from parameter sharing in the spatial dimensions, but required angular dimension shuffling to overcome sequence order bias. Both models concatenated the b-vectors as channels within the input, therefore treating the b-vectors as additional features, rather than coordinates of a manifold for which the data lies on. Finally, Ren et al. [31] proposed a conditional generative adversarial network (GAN) that additionally used structural T1w and T2w data, as well as b-vector information to predict unseen dMRI intensity data.
As there exists a myriad of downstream analyses for dMRI, this work focuses on two commonly used methods: FBA [30] and NODDI [45]. FBA relies on segmenting FODs [18] into ‘fixels’ or fibres within a voxel. Relevant works in FOD super-resolution include Lucena et al. [24] and Zeng et al. [44]. Both of these studies regressed low angular resolution derived FODs with their high resolution counterparts through 3D CNNs. NODDI estimates micro-structural complexity of the underlying brain tissue through modelling axon and dendrite populations within a voxel. Approaches to infer high angular resolution derived NODDI parameters from lower resolution raw dMRI data include Golkov et al. [15], Ye et al. [41], and Ye et al. [42]. These studies use various sub-sampling schemes and network architectures, including a 3D CNN within Ye et al. [41].
This section introduces parametric continuous convolutions, and extends the framework into the dMRI domain. We highlight our key contributions to the PCConv operation, then finally provide implementation details used to perform training and inference within the task of dMRI superresolution.
3.1 Parametric Continuous Convolutions
A discrete convolution within the context of deep learning is defined as follows
$$
h [ n ] = (f * g) [ n ] = \sum_ {m = - M} ^ {M} f [ m ] g [ n - m ],
$$
where the data function  $f: \mathcal{N} \to \mathbb{R}$  and kernel  $g: \mathcal{G} \to \mathbb{R}$  are defined over the support domain for the finite integer set  $\mathcal{N} = \mathbb{Z}^D$  and  $\mathcal{G} = \{-M, -M + 1, \dots, M - 1, M\}^D$  respectively. Conversely, a continuous convolution is defined as
$$
h (\mathbf {y} _ {j}) = (f * g) (\mathbf {y} _ {j}) = \int_ {- \infty} ^ {\infty} f (\mathbf {x}) g (\mathbf {y} _ {j} - \mathbf {x}) d \mathbf {x},
$$
where both  $f$  and  $g$  are continuous functions over the support domain  $\mathcal{N} = \mathcal{G} = \mathbb{R}^D$ . Typically in deep learning, this integration is analytically intractable but can be approximated as in Wang et al. [39] using the following form:
$$
h (\mathbf {y} _ {j}) = \int_ {- \infty} ^ {\infty} f (\mathbf {x}) g (\mathbf {y} _ {j} - \mathbf {x}) d \mathbf {x} \approx \frac {1}{N} \sum_ {i = 1} ^ {N} f (\mathbf {x} _ {i}) g (\mathbf {y} _ {j} - \mathbf {x} _ {i}), \tag {1}
$$
where a finite number of input points  $\mathbf{x}_i$  are sampled from the support domain for each output point  $\mathbf{y}_j$ . In the parametric continuous convolution framework, we choose to use a multi-layer perceptron (MLP) as the kernel function  $g(\mathbf{e};\pmb {\theta}) = \mathrm{MLP}(\mathbf{e};\pmb {\theta})$ , where  $\mathbf{e}$  is a coordinate embedding that depends on  $\mathbf{x}_i$  and  $\mathbf{y}_j$ , which is described in detail in Section 3.2. The kernel function  $g(\mathbf{e};\pmb {\theta}): \mathbb{R}^{D}\to \mathbb{R}$  spans the full continuous domain while being parameterised by a finite number of parameters  $\pmb{\theta}$ . In contrast to standard convolutional operations, the weights of the PCCov are not learnt directly, and are instead sampled from  $g(\mathbf{e};\pmb {\theta})$ . As the kernel function is itself a fully differentiable function, it can be trained through back-propagation.
In typical convolutional layers the set of output points form a subset of the input points set, provided there is no padding. However, within the parametric continuous framework, these two sets can be disjoint. This is key within the task of dMRI angular super-resolution, where the angular component of  $\mathbf{y}_j$ , i.e. the target b-vectors, is not contained within the set of input b-vectors.
Each PCConv layer requires three inputs to produce the output features  $h_{k,j}(\mathbf{y}_j)$ : the input features  $f_{c,i}(\mathbf{x}_i)$ , the set of input points  $\mathbf{x}$ , and the set of output points  $\mathbf{y}$ . Within each layer, following on from the right-hand side of Equation (1), the PCConv operation is defined as
$$
h _ {k, j} \left(\mathbf {y} _ {j}\right) = \sum_ {c = 1} ^ {C} \sum_ {i = 1} ^ {N} f _ {c, i} \left(\mathbf {x} _ {i}\right) g _ {c, k} (\mathbf {e}; \boldsymbol {\theta}). \tag {2}
$$
Here,  $C$  and  $K$  denote the input and output feature dimensions respectively. Within each layer the number of input points  $N$  is fixed, therefore the kernel function  $g$  will encapsulate the  $\frac{1}{N}$  factor in Equation (1) during training.
3.2 Convolutions in Q-Space
We define q-space  $\mathbb{Q}$  as the joint spatial-angular domain within dMRI where  $\mathbb{Q} = \{(u,v,w,\rho ,\vartheta ,\varphi)^{\top}\mid u,v,w,\rho \in \mathbb{R},0\leq \vartheta < 2\pi ,0\leq \varphi \leq \pi \}$  and  $\mathbf{x}_j,\mathbf{y}_i\in \mathbb{Q}$ . Here,  $(u,v,w)^{\top}$  denote the spatial components, and  $(\rho ,\vartheta ,\varphi)^{\top}$  denote the angular components. Specifically,  $\vartheta$  and  $\varphi$  refer to the azimuthal and polar angles, respectively, which form the set of points on the surface of a sphere of radius  $\rho$ . In this context,  $\rho$  is the diffusion weighting or b-value. Figure 1 visualises
Figure 1: Visualisation of dMRI parametric continuous convolution. Points  $\mathbf{x}_i$  and  $\mathbf{y}_j$  are of the form  $(u, v, w, \rho, \vartheta, \varphi)^\top$  with spatial components  $(u, v, w)$  and angular components  $(\rho, \vartheta, \varphi)$ . Left: The raw dMRI data  $f(\mathbf{x})$ . Points  $\mathbf{x}_4$  and  $\mathbf{x}_6$  are located in the same region on the q-space sphere and thus have the same angular components. Similarly,  $\mathbf{x}_6$  and  $\mathbf{x}_{20}$  have the same spatial coordinates but different angular components. Middle: the parametric continuous kernel, whose geometry is defined by the coordinate embedding e. Right: Sampled points from q-space are pointwise multiplied with the kernel and summed over to produce one output feature  $h(\mathbf{y}_j)$ .
convolutions across q-space through Equation (2). To perform convolutions using q-space, we first calculate a coordinate embedding  $\mathbf{e} \in \mathbb{R}^{2LE}$  via
$$
\mathbf {e} = \left[ \boldsymbol {\gamma} \left(p _ {1}\right), \boldsymbol {\gamma} \left(p _ {2}\right), \dots , \boldsymbol {\gamma} \left(p _ {E}\right) \right] ^ {\top}, \tag {3}
$$
$$
\gamma (p _ {m}) = \left[ \sin \left(2 ^ {0} \pi p _ {m}\right), \cos \left(2 ^ {0} \pi p _ {m}\right), \dots , \sin \left(2 ^ {L - 1} \pi p _ {m}\right), \cos \left(2 ^ {L - 1} \pi p _ {m}\right) \right]. \tag {4}
$$
Here  $p_m$  is the  $m$ th component of the coordinate vector  $\mathbf{p} \in \mathbb{R}^E$  given by  $\mathbf{p} = [u_i - u_j, v_i - v_j, w_i - w_j, \rho_i - \rho_j, d_r]$ . We have opted to enforce rotational invariance across the sphere by omitting  $\vartheta$  and  $\varphi$  directly, and instead use  $p_5 = d_r(\vartheta_j, \varphi_j, \vartheta_i, \varphi_i)$  where  $d_r(\vartheta_j, \varphi_j, \vartheta_i, \varphi_i) = \sin(\vartheta_j) \sin(\vartheta_i) + \cos(\vartheta_j) \cos(\vartheta_i) \cos(\varphi_j - \varphi_i)$ . The function  $\gamma: \mathbb{R} \to \mathbb{R}^{2L}$  is a Fourier feature mapping motivated by Tancik et al. [35] and Mildenhall et al. [27] which is used to improve the hypernetwork's ability to learn high-frequency functions. In practice,  $L$  is a hyperparameter of the PCConv layer. Figure 2 provides examples of weights sampled from a trained PCConv layer with varying  $d_r$ . The channels selected within the Figure demonstrate a variety of functions across the sphere, highlighting the network's ability to learn both high and low frequency kernels.
An important aspect of the discrete convolution is its finite support domain  $M$ , thus allowing the network to share parameters across the image domain. Within the PCConv framework this is achieved by limiting the number of points to convolve across, given some constraint. When convolving over q-space, the spatial dimensions follow the same constraint as in discrete convolutions. That is, for each dimension given a kernel size  $n$ , only the  $n$  closest points are selected. As the angular components represent sparsely sampled data across a sphere, a kernel size  $n$  selects the  $n$  closest points, as determined by  $d_r$ . We refer to the angular kernel size as  $k_q$  from this point onward. There is an additional constraint with the angular components that only points within a fixed radius  $d_{\mathrm{max}}$ , such that  $d_r \leq d_{\mathrm{max}}$ , are included. To achieve this, a binary mask is applied after the pointwise multiplication of the kernel weight with the input feature. The PCConv is a flexible operation that enables convolutions over non-Euclidean spaces. Additionally, the framework allows for kernel weights to be dependent upon coordinates that are not defined in terms of the relative distance between an input and output point. In this work, we additionally introduce a variant of the PCCNN that uses one such encoding. Denoted as 'PCCNN-Bv', this model uses  $\mathbf{p} = [u_i - u_j, v_i - v_j, w_i - w_j, \rho_i, \rho_j, d_r]$ .
Figure 2: Visualisation of kernel weights sampled from  $g_{c,k}(\mathbf{e};\pmb{\theta})$  within a PCConv layer, where  $c$  and  $k$  index the input and output feature dimensions respectively. To generate the spheres,  $d_r$  is varied whilst all other components of  $\mathbf{p}$  are kept constant.
The inclusion of this model is motivated by the non-linear relationship between different shells within multi-shell data, which may not be appropriately captured by the relative difference of b-values.
3.3 Global Information
Discrete convolutional layers are computed over a limited kernel size, and therefore have a limited field of view. This affords the convolutional layer computational efficiency, as the kernel size is typically much smaller than the image. However, this comes at the cost of losing global information. Within dMRI, structures within the brain are highly spatially dependent and follow a similar pattern in each individual. Non-local coordinates, and therefore global context, can be incorporated within the PCConv layer through the coordinate embedding. For example, the relative position of the kernel within an image could be incorporated into the coordinate embedding. This idea is inspired by how Vision transformers (ViTs) [12] actively extract global information through correlations between image patches. The availability of global context allows ViTs to show competitive performance in image recognition tasks, where accurate synthesis of semantic meaning more heavily relies on non-local information [19].
In clinical practice, neuroimaging data are routinely aligned to a reference space, such as Talairach space [3]. As such, coordinates within reference spaces are correlated across different subjects, and therefore can be used to provide context for the location of a voxel within the brain. In this work, after normalising the coordinates with respect to the subject brain mask size, we use the centroid of the training patches, as discussed in Section 3.6, as additional coordinates within  $\mathbf{p}$ . Denoted as an '-Sp' suffix, both 'PCCNN-Sp' and 'PCCNN-Bv-Sp' models within Section 4 have this modification.
3.4 Factorised Convolutions
dMRI is routinely stored in a dense four-dimensional format, with three spatial dimensions and one angular dimension whose coordinates are defined by  $\mathbf{b}$ . Performing a convolution in this high-dimensional space can be computationally expensive. We mitigate this cost by using factorised convolutions [34] within the PCConv framework. For example, in the two-dimensional case, an  $(n\times m)$  kernel would be factorised into two sequential convolutions of size  $(n\times 1)$  and  $(1\times m)$ . As in the non-factorised case, the PCConv layer's weights are sampled from one hypernetwork  $g(\mathbf{e};\pmb {\theta})$ . Following the approach used in Wang et al. [39], we further factorise the kernel weight tensor from  $W\in \mathbb{R}^{N\times O\times C\times K}$  into two sequential operations: an input convolution  $W_{\mathrm{in}} = \mathbb{R}^{N\times O\times C}$  followed by a forward projection matrix  $W_{\mathrm{out}} = \mathbb{R}^{C\times K}$ , where  $O$  is the number of output points.
3.5 Implementation
Since PCConv layers are capable of convolving over both spatial and angular dimensions, the PCCNN models are constructed exclusively of PCConv layers and non-linear activation functions. In practice, this involves a sequence of spatially pointwise  $(1\times 1\times 1\times k_q)$  PCConv layers, followed by blocks of parallel convolutions consisting of spatially pointwise and  $(3\times 3\times 3\times k_q)$  factorised convolutions. The hyperparameters for each PCCNN, including the number of PCConv layers, can be found in Figure 5 in the Appendix. Within this work, each PCConv layer has its own independent hypernetwork. This configuration allows for significant weight sharing capabilities, as the number of parameters for each hypernetwork does not depend on the kernel size or dimensionality. Each
PCConv layer or residual PCConv block is followed by a rectified linear unit (ReLU), excluding the final layer. Each hypernetwork is composed of two dense layers, each followed by a leaky ReLU with a negative slope of 0.1, and a final dense layer with output size 1 and no subsequent activation. The code for this work is available at github.com/m-lyon/dmri-pcconv.
3.6 Data and Training
Given the high dimensionality of dMRI data, using complete acquisitions as individual training examples is prohibitively expensive in terms of memory. To mitigate this issue, training sets were derived from input image patches  $\mathbf{X}_{\mathrm{in}} \in \mathbb{R}^{10 \times 10 \times 10 \times q_{\mathrm{in}}}$ , and target image patches  $\mathbf{X}_{\mathrm{out}} \in \mathbb{R}^{10 \times 10 \times 10 \times q_{\mathrm{out}}}$  where  $q_{\mathrm{in}}$  and  $q_{\mathrm{out}}$  denote the number of input and target b-vectors respectively. During training, to sample the angular dimension of the input sets, an initial b-vector  $\mathbf{q}_0$  was randomly selected. The next  $n$ , up to  $q_{\mathrm{in}}$ , b-vectors  $\mathbf{q}$  were chosen using  $\arg \max_j (\min_i (d_r(\mathbf{q}_i, \mathbf{q}_j))), \forall i, j \in \{0, 1, \dots, n-1\}$ . Similarly, this process was repeated to sample the output set up to size  $q_{\mathrm{out}}$ .
To reduce inter-subject variability, each dMRI dataset was normalised such that  $99\%$  of the data lay between  $[-1, 1]$ . Initially, a normalisation range of  $[0, 1]$  was used; however, this resulted in training instabilities within the hypernetwork. This was likely due to a greater mismatch between the data distribution and the weight initialisation scheme within the hypernetwork, which followed  $w \sim \mathcal{U}(-\sqrt{\frac{1}{k_{\mathrm{in}}}}, \sqrt{\frac{1}{k_{\mathrm{in}}}})$ , where  $k_{\mathrm{in}}$  was the number of input features in a given layer. dMRI data from the WU-Minn Human Connectome Project (HCP) [38] were used for training, validation and testing. This data were initially processed with the standard HCP processing pipeline [14], as well as denoised using the 'patch2self' algorithm [13]. Models were trained on twenty-seven subjects from the HCP dataset, while three subjects were used for validation during development. Hyperparameters for the PCCNN were selected through a random grid search with RayTune [20].
Each acquisition within the HCP dataset had shells of  $b_{\mathrm{full}} = \{1000, 2000, 3000\} \, \mathrm{s/mm^2}$ , each with 90 directions. Training examples contained randomly selected shells  $b_{\mathrm{in}}, b_{\mathrm{out}} \sim U(b_{\mathrm{full}})$ . This sampling scheme trained each network on single-shell and multi-shell examples simultaneously. The input angular dimension size was determined via  $q_{\mathrm{in}} \sim U(q_{\mathrm{sample}})$ ,  $q_{\mathrm{sample}} = \{6, 7, ..., 19, 20\}$ . This approach allowed each network to learn the data distribution for a range of angular input sizes. To implement this, training examples had a fixed  $q_{\mathrm{in}} = 20$ , with zero-filled entries for input size less than 20. Models were trained using 4 NVIDIA A100's with a batch size of 16, and an  $\ell_1$  loss function, for 200,000 iterations using AdamW [23]. Each PCCNN model took approximately a week to train. Details of pre-processing and training for other models can be found in the Appendix. The model weights used for test inference in Section 4 were the best performing in the validation dataset.
We evaluated the performance in dMRI angular super-resolution using eight, previously unseen subjects within the HCP dataset. Our results included four variants of the PCCNN, as listed in Table 1. As dMRI data is most commonly used for downstream analyses, we compared the results in commonly used techniques for both single-shell and multi-shell experiments. Our results included other methods where appropriate, such as spherical harmonic (SH) interpolation for single-shell experiments, FOD-Net [44] for multi-shell FBA, and SR-q-DL [41] for multi-shell NODDI. FOD-Net is a 3D CNN that maps low resolution FOD data to high resolution FOD data. SR-q-DL is a 3D CNN that maps low resolution dMRI data to high resolution NODDI data. RCNN is a 3D recurrent CNN that maps low resolution dMRI data to high resolution dMRI data. When applicable, we also