Tensor networks and operations on them are described using an intuitive graphical notation, introduced in Penrose (1971). Figure 2 (left) shows the commonly used notations for a scalar $S$, vector $V^{i}$, matrix $M^{ij}$ and a general order-$3$ tensor $T^{ijk}$. The number of dimensions of a tensor is captured by the number of edges emanating from the nodes denoted by the edge indices. For instance, the vector $V^{i}$ is an order-$1$ tensor indicated by the edge with index $i$ and an order-$3$ tensor has three indices $(i,j,k)$ depicted by the three edges, and so on.

Tensor notations can also be used to capture operations on higher-order tensors succinctly as shown in Figure 2 (center) where matrix product is depicted, which is also known as tensor contraction. The edge between the tensor nodes $X^{i}_{j}$ and $Y^{k}_{j}$ is the dimension subsumed in matrix multiplication resulting in the matrix $Z^{ik}$. More thorough introduction to tensor notations can be found in Bridgeman and Chubb (2017).

A linear model in a sufficiently high dimensional space can be very powerful (Novikov et al., 2016). In SVMs, this is accomplished by the implicit mapping of the input data into an infinite dimensional space using radial basis function kernels (Hofmann et al., 2008). In this section, we describe the procedure followed in recent tensor network based works, including in ours, to map the input data into an exponentially high dimensional space.

Consider an input vector ${\mathbf{x}}\in[0,1]^{N}$, which can be obtained by flattening a 2D or 3D image with $N$ pixels with intensity values that are normalized in the interval $[0,1]$. A commonly used feature map for tensor networks is obtained from the tensor product of pixel-wise feature maps (Stoudenmire and Schwab, 2016):

where the local feature map acting on a pixel $x_{j}$ is indicated by $\phi^{i_{j}}(\cdot)$. The indices $i_{j}$ run from $1$ to $d$ where $d$ is the dimension of the local feature map. The feature maps are usually simple non-linear functions restricted to have unit norm such that the joint feature map in Eq. (1) also has unit norm. A widely used local feature map with $d=2$ inspired from quantum wave function analysis is (Stoudenmire and Schwab, 2016) is shown in (2), and a simpler local feature map from Efthymiou et al. (2019) in Eq. (3) with similar properties (but non-unit norm):

The joint feature map $\Phi({\mathbf{x}})$ in Eq. (1) is an order-$N$ tensor due to tensor product of the $N$ order-$1$ local feature maps of dimension $d$ in Eq. (2). The joint feature map $\Phi({\mathbf{x}})$ can be seen as mapping each image to a vector in the $d^{N}$ dimensional feature space (Stoudenmire and Schwab, 2016). For multi-channel inputs, such as RGB images or other imaging modalities, with $C$ input channels the local feature map can be applied to each channel separately such that the resulting space is of dimension $(d\cdot C)^{N}$ (Efthymiou et al., 2019).

Given this high dimensional joint feature map $\Phi({\mathbf{x}})$ of Eq. (1) for the input data ${\mathbf{x}}$, a linear decision rule for a multi-class classification task can be formulated as:

where $m=[0,1,\dots M-1]$ are the $M$ classes, and the weight tensor $W^{m}$ is of order-($N+1$) with the output tensor index $m$.

The tensor notation representation of the linear model of Eq. (5) is shown in Figure 3 (Step 1) where the first column of gray nodes are the individual pixel feature maps of order-$1$ with feature dimension $d$. The local feature maps are connected to the order-($N+1$) weight tensor $W^{m}$ along $N$ edges with one edge for output of dimension $M$ marked with index $m$.

The order-($N$+1) weight tensor $W^{m}$ results in a total of $M\cdot d^{N}$ weight elements. Even for a relatively small gray scale image, say of size $100\times 100$, the total number of weight elements in $W^{m}$ can be massive: $2\cdot 2^{10000}\approx 10^{3010}$ which is exponentially more than the number of atoms in the universe⁵⁵5https://en.wikipedia.org/wiki/Observable_universe! The $d^{N}$ lift in quantum physics is sometimes seen as a convenient illusion (Poulin et al., 2011; Orús, 2014) as the most interesting behavior of systems can be captured with fewer degrees of freedom in this high dimensional space. This is analogous to using fewer than input feature dimensions after dimensionality reduction operations. In the next section we will see how tensor networks can access useful sub-spaces represented by such high dimensional tensors, with parameters that grow linearly with $N$ instead of growing exponentially with $N$.

Consider two vectors, $T^{i}$ and $U^{j}$ with indices $i$ and $j$ respectively. The tensor product ⁶⁶6Note that the tensor product of two order-1 tensors is the same as the vector outer product. of these two vectors yields an order-$2$ tensor or a matrix $X^{ij}$. The matrix product state (MPS) (Perez-Garcia et al., 2006) is a type of tensor network that expands on this notion of tensor products allowing the factorization of an order-$N$ tensor (with $N$ edges) into a chain of lower-order tensors. Specifically, the MPS tensor network can factorise an order-$N$ tensor with a chain of order-$3$ (with three edges) except on the borders where they are of order-$2$, as shown in Figure 3 (Step 2). Consider a tensor of order-$N$ with indices $i_{i},i_{2},\dots i_{N}$, using MPS it can be approximated using lower-order tensors as

where $A^{i_{j}}$ are the lower-order tensors. The subscript indices $\alpha_{j}$ are virtual indices that are contracted and are of dimension $\beta$ referred as the bond dimension. The components of these intermediate lower-order tensors $A^{i_{j}}$ form the tunable parameters of the MPS approximation. Note that any $N$ dimensional tensor can be represented exactly using an MPS tensor network if $\beta=d^{N/2}$. In most applications, however, $\beta$ is fixed to a small value or allowed to adapt dynamically when the MPS tensor network is used to approximate higher-order tensors (Perez-Garcia et al., 2006; Miller, 2019).

One of the drawbacks of using MPS tensor networks is that they operate along one dimension (as a chain). This is the primary reason that two dimensional image data has to be first flattened to a vector when working with tensor networks such as the MPS. Tensor networks that can work on arbitrary graphs, which might be more suitable for image data like the projected entangled pair states (PEPS) (Verstraete and Cirac, 2004), are not as well understood and do not yet have efficient algorithms like the MPS.
Dot product with MPS: The decision function in Eq. (5) comprising the dot product of the order-($N$+1) weight tensor $W^{m}$ and the joint feature map $\Phi({\mathbf{x}})$ in Eq. (1) can be efficiently computed using the MPS approximation in Eq. (6), depicted in Figure 3 (Step 2). The order in which tensor contractions are performed can yield a computationally efficient algorithm. The original MPS algorithm (Perez-Garcia et al., 2006) starts from one of the ends, contracts a pair of tensors to obtain a new tensor which is then contracted with the next tensor and this process is repeated until the output tensor is reached. The cost of this algorithm scales with {$N\cdot\beta^{3}\cdot d$} when compared to the cost without the MPS approximation which scales with $d^{N}$. In this work, we use the MPS implementation in Miller (2019) which performs parallel contraction of the horizontal edges and then proceeds to contract them vertically as depicted in Figure 3 (Step 3). The MPS approximation also reduces the number of tunable parameters by an exponential factor, from {$M\cdot d^{N}$} to {$M\cdot d\cdot N\cdot\beta^{2}$}.

The proposed LoTeNet model operates on small image regions and aggregates them with MPS operations at multiple levels. These small image patches at each level are created using the squeeze operation in two steps, as illustrated in Figure 5 for an order-$2$ tensor i.e an image with two spatial dimensions.

The input images are converted into vectors with inflated feature dimensions by applying kernels of stride $k$ along each of the spatial dimensions. Consider an input image $X^{hij}$ with $S=3$ spatial dimensions (height:H, width:V, depth:D), $N$ pixels and $C=d$ channels as feature dimension: $X^{hij}\in{\mathbb{R}}^{H\times V\times D\times C}$. In the first step, input image is reshaped into smaller patches controlled by the stride $k$: $X^{hij}\in{\mathbb{R}}^{(H/k)\times(V/k)\times(D/k)\times d}$ such that the feature dimension increases to $d=C\cdot k^{S}$. The stride of the kernel $k$ decides the extent of reduction in spatial dimensions and the corresponding increase in feature dimension of the squeezed image.

In the second step, the reshaped image patches with inflated feature dimensions are flattened from order-$S$ tensors to order-$1$ tensors, resulting in $X^{h}\equiv{\mathbf{x}}\in{\mathbb{R}}^{(N/k^{S})}$ with $d=C\cdot k^{S}$. This flattening operation provides spatial information in the feature dimension to LoTeNet, thus retaining additional image level structure when compared to flattening the entire image into a single vector ${\mathbf{x}}\in{\mathbb{R}}^{N}$ with $d=C$ channels (Efthymiou et al., 2019). Further, the increase in the feature dimension $d=C\cdot k^{S}$ makes the tensor network more expressive as it increases the dimensionality of the feature space (Stoudenmire and Schwab, 2016).

The transformations in dimensions due to the squeeze operation parameterised by the stride $k$, $\psi(\cdot;k)$, are summarised below:

The squeeze operation can also be treated as a local feature map due to the increase in feature dimensions, similar to the more formal pixel level feature maps mentioned in Eq. (2) and Eq. (3), operating on image patches instead of individual pixels. Finally note that the squeeze kernel stride can be different at each level $l$ of the model denoted $k_{l}$.

In the current formulation, the images are assumed to be rectangular in each plane. As the main purpose of the squeeze operation is to increase the feature dimensions by flattening small local neighbourhoods, this can be achieved by reshaping any small, non-square region of the image into a vector of fixed feature dimension. The simplest strategy to work with non-square (circular images for instance) would be to pad the regions around the borders to make them rectangular. Another approach could be to partition non-square images into cells using Voronoi tesselations (Du et al., 1999) and squeezing these cells into vectors.

An overview of the proposed LoTeNet model is shown in Figure 6 for an image with two spatial dimensions i.e with $S=2$. The number of squeezed vectors and the corresponding MPS blocks at each layer of the proposed model are also indicated in Figure 6.

The LoTeNet model comprises of layers of MPS blocks interleaved with squeeze operations without any non-linear components. The input to the first layer in Figure 6 is the full resolution input image with $N$ pixels, shown with grids marking the $k\times k$ patches which are then squeezed to obtain $N_{1}=N/k^{S}$ vectors with $d=C\cdot k^{S}$. Each of these squeezed patches are input into MPS blocks which contract the $d=C\cdot k^{S}$ vectors to output a vector with dimension $d=\nu$. MPS blocks operating on squeezed image patches can be interpreted as summarising them with a vector of size $\nu$ using a linear model in a higher dimensional feature space. In LoTeNet, we set $\nu$ to be the same as the MPS bond dimension $\beta$, so that there is only a single hyperparameter to be tuned.

The output vectors from all $N_{l}$ MPS blocks at a given layer $l$ are reshaped back into the $S$ dimensional image space. However, due to the MPS contractions, these intermediate image space representations will be of lower resolution as indicated by the smaller image with fewer patches in Figure 6. This is analogous to obtaining an average pooled version of the intermediate feature maps in traditional CNNs. The intermediate images of lower resolution formed at layer $l$ are further squeezed and contracted in the subsequent layers of the model. This process is continued for $L-1$ layers and the final MPS block acting as layer $L$ contracts the output vector from layer $L-1$ to predict the decisions in Eq. (5). The weights of all the MPS blocks in each of the $L$ layers are the model parameters which are tuned in a supervised setting as described next.

We view the sequence of MPS contractions in successive layers of our model as forward propagation and rely on automatic differentiation to compute the backward computation graph (Efthymiou et al., 2019). Torch MPS package (Miller, 2019) is used to implement MPS blocks and trained in PyTorch (Paszke et al., 2019) to learn the model parameters from training data in an end-to-end fashion. These parameters are analogous to the weights of neural network layers and can be updated in a similar iterative manner by backpropagating a relevant error signal computed between the model predictions and the training labels.

We minimize the cross-entropy loss between the true label $y_{i}\in[0,\dots,M-1]$ for each image ${\mathbf{x}}_{i}\in\mathcal{D}$ and the predicted label $f^{(y_{i})}({\mathbf{x}}_{i})$ in the training set $\mathcal{D}$:

where $\sigma(\cdot)$ is the softmax operation used to obtain normalized scores that can be interpreted as the predicted class probabilities. For binary classes, the loss in Eq. (8) reduces to the binary cross entropy and output dimension $M=1$ can be used with the sigmoid non-linearity to obtain probabilistic predictions in $[0,1]$.

We report experiments on three public datasets for the task of binary classification from 2D histopathology slides, 2D thoracic computed tomography (CT) and 3D T1-weighted MRI scans.
PCam Dataset: The PatchCamelyon (PCam) dataset is a binary histopathology image classification dataset introduced in Veeling et al. (2018). Image patches of size $96\times 96$ px are extracted from the Camelyon16 Challenge dataset (Bejnordi et al., 2017) with positive label indicating the presence of at least one pixel of tumour tissue in the central $32\times 32$ px region and a negative label indicating absence of tumour, as shown in Figure 7(a). In this work, a modified PCam dataset from the Kaggle challenge⁷⁷7https://www.kaggle.com/c/histopathologic-cancer-detection is used that removes duplicate image patches included in the original dataset. This dataset consists of about $220,000$ patches for training and an independent test set of about $57,500$ patches is provided for evaluating the models. Data augmentation is performed during training with $0.5$ probability comprising of rotation, horizontal and vertical flips.
LIDC Dataset: The LIDC-IDRI datatset comprises of 1018 thoracic CT images with lesions annotated by four radiologists (Armato III et al., 2004). We use the $128\times 128$ px 2D slices from Knegt (2018) (shown in Figure 7 (b)) yielding a total of $15,096$ patches. The segmentation masks from the four raters are transformed to binary labels indicating the presence (if two or more radiologists marked a tumour) or absence of tumours (if less than two raters marked tumours). These binary labels capture a majority voting among the radiologists. The dataset is split into $60:20:20$ splits for training, validation and a hold-out test set.
OASIS Dataset: The OASIS dataset consists of T1-weighted MRI scans of 416 subjects aged between 18 to 96 years (Marcus et al., 2007). The scans are graded with clinical dementia rating (CDR) into four classes: non-demented=0, very mild dementia=0.5, mild dementia=1.0, moderate dementia=2.0. We create an Alzheimer’s disease classification dataset according to Wen et al. (2020) with scans of all subjects above 60 years resulting in a dataset with 155 subjects. Binary labels are extracted from the CDR scores: cognitively normal (CN) if CDR=0 and Alzheimer’s disease (AD) if CDR$>$0, yielding 82:73 split between CN and AD cases. The MRI scans are preprocessed using the extensive pipeline described in Wen et al. (2020) comprising bias field correction, non-linear registration and skull stripping using the ClincaDL package⁸⁸8https://clinicadl.readthedocs.io/en/latest/. After preprocessing we obtain volumes of size 128x128x128 voxels (Figure 7-c).

The proposed LoTeNet model is related to the tensor network models used for supervised learning tasks, such as the ones in Stoudenmire and Schwab (2016); Efthymiou et al. (2019). In Stoudenmire and Schwab (2016), the parameters of the tensor network are learned from data using density matrix renormalization group (DMRG) algorithm (McCulloch, 2007), whereas in LoTeNet we use automatic differentiation based on the implementation in Miller (2019) similar to Efthymiou et al. (2019) to optimise the parameters of the model in Eq. (6).

The proposed LoTeNet model also improves upon the computation cost of approximating linear decisions compared to other tensor network models such as in Efthymiou et al. (2019). Consider the computation cost of approximating the linear decision in Eq. (5) for an input image with $N$ pixels of $d$ features with a bond dimension $\beta$. For the tensor network in Efthymiou et al. (2019) it is $\mathcal{O}\left(N\cdot d\cdot\beta^{2}\right)$. For LoTeNet, the computation cost reduces exponentially with the number of layers used: $\mathcal{O}\left(\frac{N}{k^{2\cdot L}}\cdot L\cdot k^{2}\cdot d\cdot\beta^{2}\right)$. The spatial resolution of the input image is reduced with successive layers in LoTeNet (Figure 6) captured as the $k^{2\cdot L}$ term in the denominator. The cost of performing MPS operations on patches of size $k\times k$ in $L$ layers is the additional $L\cdot k^{2}$ term in the numerator. However, for $L\geq 2$ and $k>1$, the exponential reduction in computation cost with $L$ can translate into considerable computational advantage when processing high resolution medical images. Further, the computation cost of LoTeNet can be reduced without noticeable degradation in performance by sharing MPS in each layer across all patches (Selvan et al., 2020).

The performance of LoTeNet across the three datasets is on par with, or superior, to the other methods as reported in Tables 1 and 2. This robustness across tasks is remarkable because the LoTeNet architecture is fixed across the tasks. This behaviour could be attributed to the single model hyperparameter – bond dimension, $\beta$ – which is fixed ($\beta=5$) across all experiments. Recollect that the bond dimension controls the quality of MPS approximations of the corresponding operations in high dimensional spaces as described in Section 2.3. Consistent with the reporting in earlier works (Efthymiou et al., 2019), we also find that the value of $\beta$ after a specific value does not affect the model’s performance. This is illustrated in Figure 8, where we show the sensitivity of the validation set performance on LIDC dataset for different bond dimensions. This is a desirable feature as the models can be expected to yield comparable performance when applied to different datasets without entailing dataset specific cost of elaborate architecture search. Additionally, this robustness could also be attributed to more efficient computation of massive number of parameters in LoTeNet resulting in generalised decision rules, as reported in Table 2 where we notice that the number of parameters in LoTeNet are higher than other methods (except MLP baseline).

In Tables 1 and 2, we report the GPU memory requirement for each of the models. LoTeNet requires only a fraction of the GPU memory utilised by the corresponding baseline methods, including the Densenet (Huang et al., 2017) or Rotation Eq-CNN models (Veeling et al., 2018). One reason for this drastic reduction in GPU memory utilisation is because tensor networks do not maintain massive intermediate feature maps. This behaviour is similar to feedforward neural networks and unlike CNNs which use a large chunk of GPU memory mainly to store intermediate feature maps (Rhu et al., 2016). Further, as LoTeNet is based on contracting input data into smaller tensors in a hierarchical manner the memory consumption with successive contracted layers does not increase. This can be an important feature in medical imaging applications as larger images and larger batch sizes can be processed without compromising the quality of the learned decision rules.

The local feature maps in Equations (2) and (3) are simple transformations which have been used in the tensor network literature (Stoudenmire and Schwab, 2016; Efthymiou et al., 2019). More recently, other intensity based local feature maps such as wavelet transforms have been attempted for 1D signal classification in Reyes and Stoudenmire (2020). While we use the local map in Equation (2) for 2D data and squeeze operation based local map for 3D data, experiments on other possible local feature maps were also performed. For instance, we experimented with jet based feature maps (Larsen et al., 2012) that compute spatial gradients at multiple scales, and also an MLP based local feature map acting on individual pixels. Both these local feature maps did not yield any considerable performance improvement and to adhere to existing tensor network literature we used the simple sinusoidal feature maps which also do not have additional hyperparameters.

Feed-forward neural networks: The class of tensor network models adapted for machine learning tasks (Stoudenmire and Schwab, 2016; Efthymiou et al., 2019) are conceptually most similar to feed-forward neural networks such as the MLP, as both classes of models operate on vector inputs (see Figure 9). The difference, however, arises in the type of decision boundaries optimised by each class of model as illustrated in Figure 1. MLP based feed-forward networks obtain non-linear decision boundaries using several layers of neurons and non-linear activation functions (such as sigmoid, ReLU etc.). Tensor network models, including the proposed model, optimise linear decision boundaries in exponentially high dimensional spaces without using any non-linear activation functions.
Convolutional neural networks: LoTeNet uses MPS blocks on image patches and aggregates these tensor representations in a hierarchical manner to learn the decision rule. While this could appear similar to the use of strided convolution kernels in CNNs, it is indeed closer to feed-forward neural networks. The use of an MPS block per patch perhaps can be interpreted as a form of strided convolution but only in how the local operation is performed on each patch. The primary difference is in the weight sharing; the weights of MPS blocks are not shared across the image. Each $k\times k$ image region has an MPS block acting on it. This is indicated in Figure 6, where we report the number of MPS blocks used per layer.

The computation time reported in Table 2 shows it to be higher for LoTeNet ($45.3$s) compared to the CNN ($12.8$s) or MLP ($4.1$s) baselines as efficient implementations of tensor contractions are not natively supported in frameworks such as PyTorch. There are ongoing efforts to further improve efficient implementations which can accelerate tensor network operations (Fishman et al., 2020; Novikov et al., 2020). Considering that LoTeNet usually converges between 10-20 epochs, the increase in computation time with the current implementation is not substantial as the model can be trained even on 3D data easily under an hour.

One possible drawback of having a single hyperparameter $\beta$ controlling the MPS approximations is the lack of granular control of the model capacity. In the 2D OASIS experiments (Table 2), we further investigated the lower average balanced accuracy score by trying out different $\beta$ values. This experiment revealed that LoTeNet was either under-fitting ($\beta<4$) or over-fitting ($4\leq\beta\leq 8$) to the training set. This could be attributed to the fact that the number of parameters grow quadratically with $\beta$ and these discrete jumps make it harder to obtain models of optimal capacity for some tasks.

Another concern with tensor network based models is their tendency to be over parameterised as the number of parameters scale as {$d\cdot N\cdot\beta^{2}$}. The linear dependency on the number of pixels, $N$, for volumetric data results in a model with massive number of parameters at the outset which is further aggravated by the quadratic dependency on $\beta$. While in our experiments, this has not had an adverse effect on the performance, we have observed the model is capable of over-fitting on small training data quite easily. This can be a concern when dealing with datasets with few samples.

One way of handling these problems would be to introduce weight-sharing between MPS blocks at each layer so that the number of parameters can be reduced and controlled better. A possible future work is to investigate the effect of sharing MPS blocks per layer and study if equivariance to translation can be induced in tensor network based models (Cohen and Welling, 2016).

Finally, the issue of extending LoTeNet type models beyond classification or regression tasks to more diverse medical imaging tasks such as segmentation and registration are open research questions waiting to be explored further. Early work on using tensor networks for 2D segmentation tasks was recently reported in Selvan et al. (2021).