# Predicting Poverty Level from Satellite Imagery using Deep Neural Networks

Varun Chitturi
Mission San Jose High School
<EMAIL_ADDRESS>
Zaid Nabulsi
Stanford Univeristy
<EMAIL_ADDRESS>

###### Abstract

Determining the poverty levels of various regions throughout the world is
crucial in identifying interventions for poverty reduction initiatives and
directing resources fairly. However, reliable data on global economic
livelihoods is hard to come by, especially for areas in the developing world,
hampering efforts to both deploy services and monitor/evaluate progress. This
is largely due to the fact that this data is obtained from traditional door-
to-door surveys, which are time consuming and expensive. Overhead satellite
imagery contain characteristics that make it possible to estimate the region’s
poverty level. In this work, I develop deep learning computer vision methods
that can predict a region’s poverty level from an overhead satellite image. I
experiment with both daytime and nighttime imagery. Furthermore, because data
limitations are often the barrier to entry in poverty prediction from
satellite imagery, I explore the impact that data quantity and data
augmentation have on the representational power and overall accuracy of the
networks. Lastly, to evaluate the robustness of the networks, I evaluate them
on data from continents that were absent in the development set.

##  1 Introduction

The Sustainable Development Goals lists its first of 17 goals as ending all
forms of poverty globally [9]. However, despite this being a foremost goal of
the United Nations, there are still extensive regions across the world where
extreme levels of poverty are endemic. Additionally, advancement towards that
goal is impeded by a consistent lack of data regarding key economic
indicators, especially in the developing world. Having accurate measures of
poverty helps for two main reasons:

  1. 1.

Accurate measurements of poverty level allows philanthropic agencies and
governments to identify where resources and interventions are needed, and
helps guide the direction of financial aid.

  2. 2.

Frequent and reliable data on poverty levels and distribution allows agencies
to better track progress on the Sustainable Development Goals.

Currently, poverty is formally calculated by numerous philanthropic agencies
including the World Bank. One of the reasons why data on poverty is sparse in
the developing world is because it is infrequently collected due to the high
cost associated with on-the-ground surveys. For instance, there is an average
of more than four years between national wealth surveys in a majority of
African countries [29]. Furthermore, at current rates of these surveys, a
household in the African continent will appear less than once every 1000 years
on average, which is roughly 100 times scarcer than a household in the United
States [35]. In order to increase the reliability of wealth data in regions
like countries in Africa through these surveys, an infeasible amount of
capital will need to be spent. Thus, accurately and efficiently evaluating the
level of need of the developing regions of the word and in turn directing
resources accordingly necessitates a method that can assess critical lifestyle
metrics without these costly door-to-door surveys.

Recent advancements in deep learning present an exciting opportunity for
application to poverty prediction. More specifically, both daytime and
nighttime satellite imagery of regions can be used to estimate poverty in
certain regions. [35] Deep learning has been a main factor behind recent
breakthroughs in numerous computer vision tasks such as image classification,
segmentation, and object detection [15, 10, 24, 36]. However, while the
emergence of deep learning has been a boon for countless areas, training an
accurate neural network requires vast amounts of labeled data [12, 27], a
luxury that the intersection between satellite imagery and poverty data does
not have. While there are no comprehensive datasets with labels that currently
exist for predicting poverty from satellite imagery, there does exist a
plethora of satellite imagery spanning the entire world, which can be paired
with the reliable poverty data that does exist.

In this paper, I test the hypothesis that deep learning can leverage satellite
imagery to reliably predict the poverty level of a region. I assemble a
dataset of 88,386 images from 44,193 cities spanning Africa, South America,
Asia, Europe, and the Caribbean. For each city, I obtain a daytime satellite
image, a nighttime satellite image, and the city’s wealth index. I then train
deep neural networks (DNNs) to predict a city’s wealth index, given a
satellite image. I leverage techniques such as data augmentation, pretraining
and transfer learning, regularization, and cross validation in developing and
evaluating the networks. I evaluate the networks using cross-country and
within-country wealth data, and also explore the impact of data quantity on
the representational power of the networks.

##  2 Related Work

Deep learning has been applied to satellite imagery in the context of
structure classification [33], map segmentation [31] [6], and pattern
identification [5]. Using satellite imagery to predict poverty levels of
cities is not a new task. Previous work has explored predicting humanitarian
indicators using satellite imagery. [13] uses satellite imagery to predict
$NO_{x}$ emissions, and [32] [11] [19] utilize Nightlight and RGB satellite
imagery to predict poverty outcomes with promising results.

Furthermore, deep learning has been used to infer both spatial and temporal
differences in local-level economic well-being using multiple sources of
satellite imagery [35]. Jean etl. Al. use CNNs to predict poverty from high
resolution satellite imagery of terrain10 and Perez et. al. and Xie et. al.
utilize transfer learning and multi-modal networks to predict poverty from low
resolution satellite imagery [20, 34]. More recently, Yeh et. al. presents the
utility of satellite imagery for research and policy, demonstrating their
potential to create a wealth map for the entire world [35]. Interestingly,
Uzkent et. al. presented an approach pairing satellite imagery with geolocated
Wikipedia articles to accurately predict the wealth index of a region [30].
Lastly, generative models have been used to discern the distribution of
satellite images and identify forgery [16].

To the best of my knowledge, there have been no prior work to this date that
explores the impact of data quantity on deep neural networks, or evaluates
such networks on data from different continents than they were trained on.
Furthermore, to the best of my knowledge, there have been no prior works that
explore both nighttime and daytime imagery.

##  3 Methods

I developed DNN to predict the wealth index of a region given an overhead
satellite imagery. I explore DNNs that take in as input both nighttime images
or daytime images. Formally, the task of poverty prediction can be defined as
predicting $Y$, which is the wealth index, given a daytime satellite image,
$X_{Day\\_Image}$, or nighttime satellite image, $X_{Nigttime\\_Image}$. The
DNN will output $\hat{Y}$, the predicted wealth index of the region. Because
the wealth index, $Y$, is a continuous number in contrast to a categorized ,
this task is framed as a regression. All in all, the goal is to predict some
$Y$ by training a DNN to predict $Y|X_{Nigttime\\_Image}$ or
$Y|X_{Day\\_Image}$.

###  3.1 Dataset

I obtain my development and evaluation satellite images from the Google Earth
Engine API [4]. Daytime satellite imagery is obtained from the LANDSAT
satellite [1] and nighttime satellite imagery is obtained from the VIIRS
satellite [2]. To ingest the images, I use the Google Earth Engine API to
obtain a daytime and a nighttime satellite image for 44,193 points that I had
poverty levels for (see Ground Truth). For each of these points, I obtain the
closest satellite image to the geographic coordinate of the wealth index. All
in all, I collected 88,386 satellite images. Examples of both daytime and
nighttime images are presented in Figure 1.

All satellite images had initial spatial dimensions of 890x890 pixels but were
downsampled to 256x256 pixel images. The nighttime images were single-channel
(intensity) grayscale images, while the daytime images were 3-channel RGB
images.

I divided the initial images into training, tuning, and testing sets with the
80:10:10 split. The splits were executed on the country level to ensure that
data from the same country is not in more than one set. Furthermore, the
splits were consistent across daytime and nighttime images meaning that
daytime and nighttime images of the same region were in the same split,
ensuring consistency in reporting metrics across daytime and nighttime
experiments.

All evaluation is done on held-out test regions that were not used in model
development. This mimics the real-world usage of such a network, where the
network is expected to make predictions on regions unseen during training.

Figure 1: Sample satellite images, obtained randomly from the development
dataset. On top, we see the daytime satellite images, with the corresponding
nighttime images of the same regions shown directly below.

###  3.2 Poverty Ground Truth Labels

The ground truth data is the poverty level of a region, given as a wealth
index. I obtain this wealth index data from the UN World Bank datasets [3] for
a total of 44,193 coordinates. I then normalize the wealth data, ending up
with numbers in the interval [-2, 2] ,where higher (less negative or more
positive) numbers indicate greater wealth. Zero indicates the median wealth
index in the world, negative numbers indicate below average wealth, and
positive numbers indicate above average wealth. Because the wealth index is a
continuous number, this formulation is framed as a regression task. The
distribution of the wealth indices across the splits of my dataset is shown in
Figure 2. Table 1 shows some key statistics of the wealth indices on the
country level.

Note that in obtaining the wealth indices’ ground truth data, in order to
protect the anonymity of each interview subject, a 5km jitter
(error/uncertainty) was added to every coordinate in a rural area, and 2km on
coordinates in urban areas.

Figure 2: Distribution of wealth measurements in the entire dataset (across all splits), and for each split. The values are normalized and range from -2 to 2. Note that the data is bi-modal, with a median close to 0. Furthermore, note that the distributions are consistent across data splits.  country | average | median | variance
---|---|---|---
Albania | 1.23173 | 1.26679 | 0.0524784
Armenia | 1.22638 | 1.27885 | 0.0379816
Bangladesh | -0.310995 | -0.409245 | 0.395392
Bolivia | 0.509215 | 0.763213 | 0.653936
Burkina Faso | -0.480204 | -0.666471 | 0.345656
Burundi | -0.87925 | -1.08369 | 0.320024
Cambodia | -0.162819 | -0.399031 | 0.48493
Cameroon | -0.415385 | -0.486957 | 0.493745
Central African Republic | -1.17838 | -1.16831 | 0.0453396
Comoros | 0.308073 | 0.350365 | 0.249185
Cote d’Ivoire | 0.00953811 | 0.0686464 | 0.393958
Democratic Republic of the Congo | -0.918128 | -1.20888 | 0.425719
Dominican Republic | 0.773491 | 0.878884 | 0.299624
Egypt | 1.10081 | 1.20375 | 0.130868
Ethiopia | -0.925519 | -1.23838 | 0.427344
Gabon | 0.289272 | 0.505296 | 0.509286
Ghana | -0.00517709 | 0.0100654 | 0.523037
Guinea | -0.458707 | -0.706828 | 0.525901
Guyana | 0.417429 | 0.650645 | 0.513695
Haiti | -0.277625 | -0.428153 | 0.483068
Kenya | -0.426766 | -0.577282 | 0.471848
Kyrgyz Republic | 1.21326 | 1.25745 | 0.0469557
Lesotho | -0.4049 | -0.536814 | 0.385506
Liberia | -0.811382 | -0.878786 | 0.231922
Madagascar | -0.735995 | -1.00382 | 0.425281
Malawi | -0.775441 | -0.953094 | 0.229014
Mali | -0.290661 | -0.500998 | 0.425105
Moldova | 0.996866 | 1.03015 | 0.146392
Mozambique | -0.560905 | -0.888259 | 0.616949
Myanmar | 0.147347 | 0.111866 | 0.292677
Namibia | 0.0822124 | -0.127053 | 0.9082
Nepal | -0.185337 | -0.185715 | 0.416537
Niger | -1.06198 | -1.34356 | 0.259844
Nigeria | -0.0903742 | -0.0900875 | 0.537997
Pakistan | 0.142732 | 0.146731 | 0.516279
Peru | 0.191104 | 0.310998 | 0.741152
Philippines | 0.694071 | 0.820905 | 0.342544
Rwanda | -0.812009 | -0.992644 | 0.248862
Senegal | -0.174405 | -0.318213 | 0.534392
Sierra Leone | -0.622954 | -0.862833 | 0.37829
Swaziland | 0.179458 | 0.0660039 | 0.433119
Tajikistan | 0.959175 | 1.04378 | 0.173927
Tanzania | -0.57875 | -0.789753 | 0.37443
Timor | -0.685584 | -0.904893 | 0.515151
Togo | -0.133553 | -0.308391 | 0.433274
Uganda | -0.847543 | -1.01354 | 0.241416
Zambia | -0.479619 | -0.764645 | 0.619943
Zimbabwe | -0.127609 | -0.444624 | 0.722577
Table 1: This table summarizes the wealth statistics of the countries used
during the experiments. Here, the average, median, and variance of wealth
index is provided for each country.

###  3.3 Neural Network Development

I trained multiple DNNs to predict the wealth index of a region (representing
the poverty level).The DNNs had a single output each. Firstly, I flattened the
input image, and then passed it through two multi-layer perceptrons, each with
512 nodes before finally outputting a wealth index score. The same
architecture is used for daytime and nighttime images, as shown in Figure 3.

An SGD optimizer with momentum [22] is used with a learning rate of $10^{-7}$
for all experiments on nighttime imagery and a learning rate of $10^{-9}$ for
all experiments on daytime imagery. A momentum value of $0.9$ was used for all
experiments. I trained the DNNs using mean squared error (MSE) loss, as shown
in Equation 1.

To ensure fair comparison across trials and experiments, I hold all
hyperparameters and training procedures constant, with the exception of
learning rate which differed for nighttime and daytime imagery as detailed
above. However, within all networks trained on daytime images and all networks
trained on nighttime images, the learning rate was held constant. Furthermore,
all networks were trained for the 10 epochs. Because I empirically determined
run-to-run variance for the same experiment to be low, I only ran each
experiment once.

I developed the networks using Pytorch version 1.8 [18] using a single NVIDIA
Tesla V100 graphics processing unit for training and evaluation.

Figure 3: Architecture of the neural network used in experiments

$L=\frac{1}{n}\sum_{i=1}^{n}(Y_{i}-\hat{Y}_{i})^{2}$ (1)

###  3.4 Checkpoint Selection

All networks were trained on images only from the train split (See _Dataset_
section). After training, a checkpoint was selected based on what had the
smallest RMSE value on the tune set, which consisted of images only from the
tune split. Finally, using that checkpoint, inference was run on the test set,
which consisted of images only from the test split. Only test set performance
is reported.

###  3.5 Evaluation and Statistical Analysis

My main evaluation metric used in the analysis is root mean squared error
(RMSE) which is one of the most commonly used metrics to assess performance on
regression tasks [7]. The goal is to minimize RMSE values.

Confidence intervals (CI) for all evaluation metrics are calculated using the
non-parametric bootstrap method with n=1000 permutations at the example level.
95% confidence intervals are reported for all root mean squared errors (RMSE).

##  4 Experiments

To test my hypothesis that deep NNs can be trained to leverage satellite
imagery to accurately identify the poverty level (wealth index) of the region
along with several sub-hypotheses, I designed and executed numerous
experiments. The experiments are aimed at evaluating the following:

  1. (a)

How well can a deep neural network predict the poverty level of a region given
an overhead nighttime satellite image with a limited dataset?

  2. (b)

How well can a deep neural network predict the poverty level of a region given
an overhead daytime satellite image with a limited dataset?

  3. (c)

What is the impact of data quantity (i.e. more or less data) on the
representational power and overall accuracy of a deep neural network trained
to predict the poverty level of a region given an overhead nighttime satellite
image?

  4. (d)

What is the impact of data quantity (i.e. more or less data) on the
representational power and overall accuracy of a deep neural network trained
to predict the poverty level of a region given an overhead daytime satellite
image?

  5. (e)

What is the impact of data augmentation on the representational power and
overall accuracy of a deep neural network trained to predict the poverty level
of a region given an overhead nighttime satellite image? How does data
augmentation compare with a larger dataset?

  6. (f)

What is the impact of data augmentation on the representational power and
overall accuracy of a deep neural network trained to predict the poverty level
of a region given an overhead nighttime daytime image?

  7. (g)

How well can a deep neural network trained on a single continent’s data
generalize to data from other continents?

See Figure 4 for a schematic of all experiments run.

(a) Experiment setup for baseline experiments.

(b) Experiment setup for data quantity experiments.

(c) Experiment setup for data augmentation experiments.

(d) Experiment setup for generalization to unseen continents experiments.

Figure 4: The four experiment setups presented in this work.

###  4.1 Baseline: Using a Limited Dataset of Daytime and Nighttime Satellite
Images

To evaluate my top level hypothesis and establish a baseline for the remaining
experiments, I ran experiments evaluating how well a deep neural network can
predict the poverty level of a region based on a single daytime or nighttime
satellite image, with a limited dataset. I restricted my training dataset to
800 images, and evaluated my network on 100 images, each randomly chosen from
their respective splits (see _Dataset_ under _Methods_). See Figure 4(a) for
an illustration of the experiment. I ran one experiment using daytime images,
and one experiment using nighttime images and results are shown in Table 2.

###  4.2 Impact of Data Quantity

Secondly, I wanted to better understand the importance of data quantity on the
performance of the network for both daytime and nighttime image regression. It
is a widely known and accepted fact that deep neural networks require a lot of
data, and performance is often strongly correlated with data quantity [25].
However, different tasks often require different amounts of data depending on
the characteristics and difficulty of the tasks and the data available [28].
To evaluate the impact of data quantity on the representational power and
overall accuracy of the DNN in the task of poverty prediction, I trained deep
DNNs with the following amount of data:

  1. 1.

800 images

  2. 2.

1600 images

  3. 3.

2400 images

I ran the previous experiments once using daytime images and once using
nighttime images for a total of 6 experiments. See Figure 4(b) for an
illustration of the experiment. Results are presented in Table 3.

###  4.3 Impact of Data Augmentation

For deep learning tasks in the low data regime where obtaining labeled data is
expensive or infeasible, a technique called data augmentation is often used to
enlarge the dataset [26]. Existing works have explored data augmentation for
satellite imagery [8] with respectable results. To evaluate the impact of data
augmentation on the representational power and overall accuracy of the DNNs
for this specific task, I trained DNNs with various amounts of data
augmentation, for both daytime images and nighttime images.

The specific augmentations I used were random flipping (180 degree rotation
followed by a mirror across the vertical axis through the center of the
image), random rotations (90 degrees counter-clockwise), and random Gaussian
noise with a mean of zero and standard deviation of 2.5 [23]. See Figure 4(c)
for an illustration of the experiment. Results for experiments with data
augmentation for both daytime and nighttime satellite imagery are presented in
Table 4.

###  4.4 Generalization to Unseen Continents

While DNNs have been shown to match and even surpass human performance in
various tasks ranging to medical image diagnosis [21] to emotion recognition
[14], a major issue these networks face is inability to generalize to data not
seen during training. Thus, to robustly evaluate a performance of a DNN, it is
important to assess its performance on unseen, "new" data to see how well it
generalizes. For example, a neural network trained to distinguish between
normal and abnormal chest X-rays, needs to be evaluated on its performance in
the presence of new patient populations and new diseases [17].

In the specific task of poverty prediction from satellite imagery, in order to
robustly evaluate DNNs, it is similarly important to assess performance on
"new" data. I explore this generalization by defining "new" data as data from
a separate continent. To evaluate how well a DNN trained on satellite images
from a single continent to predict poverty can generalize to another
continent, I ran three experiments. For each experiment, I picked a continent
at random to train on and a separate continent to evaluate on. Specifically,
the three experiments are:

  1. 1.

Train on data from Africa, evaluate on data from Asia

  2. 2.

Train on data from Europe, evaluate on data from Africa

  3. 3.

Train on data from Asia, evaluate on data from Europe.

Note that these experiments are only conducted on nighttime images because of
the better performance. See Figure 4(d) for an illustration of the experiment.
Results of the above experiments can be found in Table 5.

##  5 Results

Results for the experiments detailed above are presented in Tables 2-5, and
detailed in the respective sections below.

###  5.1 Baselines

For the task of predicting poverty, a DNN trained on nighttime images achieves
a RMSE of 1.222 (95% CI: 1.033-1.423) and a DNN trained on daytime images
achieves a RMSE of 1.611 (95% CI: 1.399-1.832). In both daytime and nighttime
experiments, each DNN was trained on 800 images and evaluated on 100 images,
each chosen from their respective data splits. Results are presented in Table
2.

Experiment Name | RMSE (95% CI)
---|---
Baseline Night | 1.222 (1.033-1.423)
Baseline Day | 1.611 (1.399-1.832)
Table 2: This table summarizes the results of the model when simple day and
night images were used to train and test the model. To evaluate the
performance of the model 800 train images and 100 test images were used.

###  5.2 Data Quantity and Data augmentation

Results for experiments on data quantity and data augmentation are presented
in Table 3 and Table 4, respectively. A DNN trained on nighttime images
achieves RMSE’s of 1.222 (95% CI: 1.033-1.423), 1.132 (95% CI: 0.963-1.354),
and 1.106 (95% CI: 0.956-1.283) when trained on 800, 1600, and 2400 images
respectively.

A DNN trained on daytime images achieves RMSE’s of 1.611 (95% CI:
1.399-1.832), 1.459 (95% CI: 1.275-1.633), and 1.477 (95% CI: 1.203-1.894)
when trained on 800, 1600, and 2400 images respectively.

The same evaluation dataset of 100 images was used across all experiments for
each of daytime and nighttime experiments for consistent comparisons.

Experiment Name | RMSE (95% CI)
---|---
Night, 800 images | 1.222 (1.033-1.423)
Night, 1600 images | 1.132 (0.963-1.354)
Night, 2400 images | 1.106 (0.956-1.283)
Day, 800 images | 1.611 (1.399–1.832)
Day, 1600 images | 1.459 (1.275-1.633)
Day, 2400 images | 1.477 (1.203-1.894)
Table 3: This table summarizes the results of the model when the amount of data used to train the model was varied. Experiment Name | RMSE (95% CI)
---|---
Baseline Night | 1.222 (1.033-1.423)
Night with Augmentation | 0.956 (0.872-1.040)
Baseline Day | 1.611 (1.399-1.832)
Day with Augmentation | 1.463 (1.221-1.705)
Table 4: This table summarizes the results of the model when the data used to
train the model was partially augmented.

###  5.3 Performance in the Setting of Data from Unseen Continents

Table 5 shows results for my generalization experiments. A DNN trained on data
from Africa and evaluated on data from Asia achieves RMSE of 0.926 (95% CI
0.814-1.031). Additionally, a DNN trained on data from Europe and evaluated on
data from Africa achieves RMSE of 3.182 (95% CI 1.981-4.321). Lastly, a DNN
trained on data from Asia and evaluated on data from Europe achieves RMSE of
1.156 (95% CI 1.014-1.285).

All three experiments here only used nighttime images. For all three
experiments, the same evaluation set of 100 images was used.

Experiment Name | Train Continent | Test Continent | RMSE (95% CI)
---|---|---|---
Experiment 1 | Africa | Asia | 0.926 (0.814-1.031)
Experiment 2 | Europe* | Africa | 3.182 (1.981-4.321)
Experiment 3 | Asia | Europe | 1.156 (1.014-1.285)

*The train set for Europe was slightly smaller due to insufficient data with 672 images compared to 800 images for the rest of the experiments.

Table 5: This table summarizes the result of the model when separate
continents were used to train and test the model. For this experiment, only
night time images were used.

###  5.4 Distributional Shifts between Countries/Continents

As expected, there was a fairly significant data distribution shift across
countries and continents. Table 1 presents key statistics regarding wealth
indices on the country level, where the distribution shift is dramatic.

Median wealth indices for each country range from -1.238 in Ethiopia to 1.279
in Armenia. Variance of the wealth index also differed dramatically across
countries, ranging from 0.908 in Namibia to 0.0380 in Armenia. Variance in the
wealth index is a proxy for wealth inequality within a country where higher
variance indicates greater inequality.

##  6 Discussion

In this paper, I developed a multitude of deep neural networks under various
circumstances with the goal of predicting a region’s poverty level from
overhead satellite images. Both daytime and nighttime satellite images were
used as input. Furthermore, different deep learning methods were utilized in
order to better understand the task. I designed and ran experiments to
determine the impact of data quantity and data augmentation on the accuracy of
the DNNs. Lastly, I designed and ran experiments to evaluate how well DNNs
generalize to data from continents that were absent in the training set.

For the task of poverty prediction, the results show that a DNN can be trained
to estimate a region’s poverty from a satellite image fairly accurately, with
RMSE values of only slightly above 1. Furthermore, from Table 2, I see that a
DNN can better estimate a region’s povery from the nighttime satellite image
rather than a daytime satellite image. This is expected, given that a
nighttime satellite image shows the amount of light in a region, which is
often correlated with an active economy (see Figure 1 for sample daytime and
nighttime satellite images). The better performance on nighttime imagery also
shows that a nighttime satellite image holds more information regarding a
region’s poverty level than a daytime satellite image. However, that is not to
say that a daytime satellite image contains no clues regarding the wealth of a
given region as the performance of a DNN trained on daytime images has RMSE as
low as 1.477, which is surprisingly accurate. The main takeaway here is that
both daytime and nighttime satellite imagery can be used to effectively
estimate the poverty level of a region.

Experiments on data quantity and data augmentation showed that, as expected,
DNN performance improved with more data for both daytime and nighttime
satellite imagery (either through data augmentation methods or through
obtaining more organic data). Significant increases in performance are seen
when the data is augmented from 800 to 1600, decreasing by 0.09 RMSE (7.36%)
and 0.152 RMSE (9.44%) for nighttime and daytime networks respectively.
Performance started to stagnate past a training dataset size of 1600. A
decrease of 0.026 RMSE (2.30%) was observed for the nighttime network and an
increase of 0.018 (1.23%) was observed for the daytime network, indicating
that more data is not necessarily useful past 1600 images. While seeing the
trend of stagnating performance with continuous training data increases is
common across deep learning, the number of images this happens is normally
much higher than 1600 images. This indicates that training DNNs for poverty
prediction from satellite imagery may not need as much data. Furthermore, this
also indicates that the at 1600 images, the constraint is no longer amount of
data, but rather the representational power of the DNN. For more improvements,
a more complicated architecture (i.e. more layers) is likely needed.

Similar to training data increases, train-time data augmentation provided a
comparable a benefit. For the nighttime network, RMSE decreased by 0.266
(21.8%). For the daytime network, RMSE decreased by 0.148 (9.19%). These
results demonstrate that in addition to obtaining more labeled data, using
data augmentation methods are also beneficial. Similarly to the results
obtained by increasing training data size, data augmentation seems to provide
more of a boost to nighttime networks. Not only does data augmentation
synthetically increase the dataset size, but it also serves as an important
regularizer.

Interestingly, data augmentation seemed to provide a larger benefit to the
nighttime network than increasing the development dataset size (achieving a
decrease of 0.266 RMSE compared to 0.116 RMSE). For the daytime networks, the
performance change is comparable (roughly 0.15 decrease in RMSE). The fact
that performance increases are more dramatic in the nighttime images than the
daytime images indicates that the DNN’s representational power can be
increased by having a more complex architecture, specifically for daytime
images. For nighttime images, performance may not have fully saturated with
the current network.

The results of the generalization experiments are also insightful. Two of the
three experiments show that the DNN can generalize well to data from unseen
continents. A DNN trained on data from only Africa achieves an RMSE of 0.926
on data from Asia, better than what the general-purpose model performs on a
general-purpose evaluation set. This demonstrates great generalization power
of the DNN. A similar trend is seen in the DNN trained on data from Asia and
evaluated on data from Europe, achieving an RMSE 1.156. However, a completely
different result is observed on the network trained on Europe and evaluated on
Africa, where an RMSE of 3.182 is noted. While the train set in this
experiment is slightly smaller (672 images compared to 800 for the prior two),
this smaller dataset likely is not sufficient to explaining the huge
performance drop, indicating that data in Europe is very different than data
in Africa.

As alluded to earlier, in this work, a simple neural network architecture was
used in this paper, consisting of only two layers of 512 nodes each (see
Figure 3). This was done by design, as if a more complex architecture was
used, it would be harder (or even infeasible) to saturate the network’s
performance based on its representational power with the the available data.
Furthermore, the goal of this work was not to develop the most accurate
network to estimate poverty level, but rather evaluate how different deep
learning techniques impact performance. Because the DNN architecture and
hyperparameters were held constant across experiments, the impact of each
technique could be better assessed.

##  7 Limitations and Future Work

While this work presents an important advancement in the intersection of
computer vision and poverty prediction from satellite imagery, there are a few
limitations worth discussing. First of all, the data is imperfect by nature.
The way I paired wealth indices with satellite imagery is not exact, as the
wealth index is given at a specific coordinate, but satellite images are
somewhat sparse. The closest daytime and nighttime satellite image to the
coordinates of the wealth index was chosen, but sometimes this point was not
as close as ideal. However, since neighboring cities tend to have similar
wealth indices, this may not be as big of an issue in practice. Furthermore,
only a subset of the available data was used. while this was necessary to
better analyze the differences and impact of each deep learning technique,
using more data would likely result in more result networks.

On the modeling front, only a single architecture was used for all
experiments, and no extensive hyperparameter search was conducted due to
computational constraints. For example, a convolutional neural network (CNN)
would be expected to perform significantly better. In practice, this likely
would have resulted in improved performances, but this was not the main goal
of this work. An additional limitation is that while both increasing data
quantity and data augmentation were both assessed, the conjunction was never
tried, and could lead to insightful results.

While the results presented in this work are exciting, there is still plenty
of work to be done in order to realize the impact on a global scale. Two main
ideas immediately stem from this work. The first is that a network can be
trained to take in as input both a daytime and a nighttime satellite image to
reap the benefits of both. Furthermore, since oftentimes multiple satellite
images of a single city are often available, these can all be provided to a
network as input (see Figure 5.1

Figure 5: Architecture of a multi-image model CNN.

##  8 Conclusion

The goal of this research was to examine a novel approach to poverty
prediction using both daytime and nighttime satellite images. While there is
more to do, I think my work can provide a basis for poverty prediction as well
as show its potential. Identifying regions of poverty is the first step on its
reduction and I believe my work has made a contribution.

##  9 Acknowledgements

I would like to acknowledge the contributions of Zaid Nabulsi for his guidance
and mentorship in tackling this problem, as well as his help in exploring
previous research efforts. I’d also like to acknowledge the contributions of
Professor Stefano Ermon from Stanford and his Sustainability and Artificial
Intelligence Lab for inspiring this research and early efforts.

## References

  * [1] https://landsat.gsfc.nasa.gov/. Accessed: 2021-09-12.
  * [2] https://earthdata.nasa.gov/earth-observation-data/near-real-time/download-nrt-data/viirs-nrt. Accessed: 2021-09-12.
  * [3] https://data.worldbank.org/. Accessed: 2021-09-12.
  * [4] Google earth engine.  https://developers.google.com/earth-engine.  Accessed: 2021-09-12.
  * [5] A. Albert, J. Kaur, and M. C. Gonzalez.  Using convolutional networks and satellite imagery to identify patterns in urban environments at a large scale.  In Proceedings of the 23rd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pages 1357–1366. ACM, 2017.
  * [6] N. Audebert, B. Le Saux, and S. Lefèvre.  Joint learning from earth observation and openstreetmap data to get faster better semantic maps.  In EARTHVISION 2017 IEEE/ISPRS CVPR Workshop. Large Scale Computer Vision for Remote Sensing Imagery, 2017.
  * [7] T. Chai and R. R. Draxler.  Root mean square error (rmse) or mean absolute error (mae)? – arguments against avoiding rmse in the literature.  Geoscientific Model Development, 7(3):1247–1250, 2014.
  * [8] P. R. Dave and H. A. Pandya.  Satellite image classification with data augmentation and convolutional neural network.  In Lecture Notes in Electrical Engineering, Lecture notes in electrical engineering, pages 83–92. Springer Singapore, Singapore, 2020.
  * [9] dpicampaigns and fangweizhao.  Goal 1: End poverty in all its forms everywhere.  https://www.un.org/sustainabledevelopment/poverty/. Accessed: 2021-09-27.
  * [10] J. Han, D. Zhang, G. Cheng, N. Liu, and D. Xu.  Advanced deep-learning techniques for salient and category-specific object detection: A survey.  IEEE Signal Processing Magazine, 35(1):84–100, 2018.
  * [11] N. Jean, M. Burke, M. Xie, W. M. Davis, D. B. Lobell, and S. Ermon.  Combining satellite imagery and machine learning to predict poverty.  Science, 353(6301):790–794, 2016.
  * [12] M. I. Jordan and T. M. Mitchell.  Machine learning: Trends, perspectives, and prospects.  Science (New York, 349(6245):255, 2015.
  * [13] S.-W. Kim, A. Heckel, S. McKeen, G. Frost, E.-Y. Hsie, M. Trainer, A. Richter, J. Burrows, S. Peckham, and G. Grell.  Satellite-observed us power plant nox emission reductions and their impact on air quality.  Geophysical Research Letters, 33(22), 2006.
  * [14] S. Li and W. Deng.  Deep facial expression recognition: A survey.  IEEE Transactions on Affective Computing, page 1–1, 2020.
  * [15] T.-Y. Lin, M. Maire, S. Belongie, L. Bourdev, R. Girshick, J. Hays, P. Perona, D. Ramanan, C. L. Zitnick, and P. Dollár.  Microsoft coco: Common objects in context, 2015.
  * [16] M. Mirza and S. Osindero.  Conditional generative adversarial nets.  Nov. 2014.
  * [17] Z. Nabulsi, A. Sellergren, S. Jamshy, C. Lau, E. Santos, A. P. Kiraly, W. Ye, J. Yang, R. Pilgrim, S. Kazemzadeh, J. Yu, S. R. Kalidindi, M. Etemadi, F. Garcia-Vicente, D. Melnick, G. S. Corrado, L. Peng, K. Eswaran, D. Tse, N. Beladia, Y. Liu, P.-H. C. Chen, and S. Shetty.  Deep learning for distinguishing normal versus abnormal chest radiographs and generalization to two unseen diseases tuberculosis and COVID-19.  Sci. Rep., 11(1):15523, Sept. 2021.
  * [18] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Köpf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala.  PyTorch: An imperative style, high-performance deep learning library.  Dec. 2019.
  * [19] A. Perez, C. Yeh, G. Azzari, M. Burke, D. Lobell, and S. Ermon.  Poverty prediction with public landsat 7 satellite imagery and machine learning.  11 2017.
  * [20] A. Perez, C. Yeh, G. Azzari, M. Burke, D. Lobell, and S. Ermon.  Poverty prediction with public landsat 7 satellite imagery and machine learning.  Nov. 2017.
  * [21] P. Rajpurkar, J. Irvin, K. Zhu, B. Yang, H. Mehta, T. Duan, D. Ding, A. Bagul, C. Langlotz, K. Shpanskaya, M. P. Lungren, and A. Y. Ng.  Chexnet: Radiologist-level pneumonia detection on chest x-rays with deep learning, 2017.
  * [22] S. Ruder.  An overview of gradient descent optimization algorithms.  arXiv preprint arXiv:1609.04747, 2016.
  * [23] E. Rusak, L. Schott, R. S. Zimmermann, J. Bitterwolf, O. Bringmann, M. Bethge, and W. Brendel.  A simple way to make neural networks robust against diverse image corruptions, 2020.
  * [24] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein, A. C. Berg, and L. Fei-Fei.  Imagenet large scale visual recognition challenge, 2015.
  * [25] M.-A. Schulz, B. T. T. Yeo, J. T. Vogelstein, J. Mourao-Miranada, J. N. Kather, K. Kording, B. Richards, and D. Bzdok.  Different scaling of linear models and deep learning in UKBiobank brain images versus machine-learning datasets.  Nat. Commun., 11(1):4238, Aug. 2020.
  * [26] C. Shorten and T. M. Khoshgoftaar.  A survey on image data augmentation for deep learning.  J. Big Data, 6(1), Dec. 2019.
  * [27] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov.  Dropout: A simple way to prevent neural networks from overfitting.  Journal of Machine Learning Research, 15(56):1929–1958, 2014.
  * [28] V. Sze, Y.-H. Chen, T.-J. Yang, and J. Emer.  Efficient processing of deep neural networks: A tutorial and survey, 2017\.
  * [29] L. S. Tusting, D. Bisanzio, G. Alabaster, E. Cameron, R. Cibulskis, M. Davies, S. Flaxman, H. S. Gibson, J. Knudsen, C. Mbogo, F. O. Okumu, L. von Seidlein, D. J. Weiss, S. W. Lindsay, P. W. Gething, and S. Bhatt.  Mapping changes in housing in sub-saharan africa from 2000 to 2015.  Nature, 568(7752):391–394, Apr. 2019.
  * [30] B. Uzkent, E. Sheehan, C. Meng, Z. Tang, M. Burke, D. Lobell, and S. Ermon.  Learning to interpret satellite images in global scale using wikipedia.  May 2019.
  * [31] J. Wang, Q. Qin, Z. Li, X. Ye, J. Wang, X. Yang, and X. Qin.  Deep hierarchical representation and segmentation of high resolution remote sensing images.  In Geoscience and Remote Sensing Symposium (IGARSS), 2015 IEEE International, pages 4320–4323. IEEE, 2015.
  * [32] W. Wang, H. Cheng, and L. Zhang.  Poverty assessment using dmsp/ols night-time light satellite imagery at a provincial scale in china.  Advances in Space Research, 49(8):1253–1264, 2012.
  * [33] G.-S. Xia, X. Bai, J. Ding, Z. Zhu, S. Belongie, J. Luo, M. Datcu, M. Pelillo, and L. Zhang.  Dota: A large-scale dataset for object detection in aerial images.  In Proc. CVPR, 2018.
  * [34] M. Xie, N. Jean, M. Burke, D. Lobell, and S. Ermon.  Transfer learning from deep features for remote sensing and poverty mapping.  Sept. 2015.
  * [35] C. Yeh, A. Perez, A. Driscoll, G. Azzari, Z. Tang, D. Lobell, S. Ermon, and M. Burke.  Using publicly available satellite imagery and deep learning to understand economic well-being in africa.  Nat. Commun., 11(1):2583, May 2020.
  * [36] S. S. A. Zaidi, M. S. Ansari, A. Aslam, N. Kanwal, M. Asghar, and B. Lee.  A survey of modern deep learning based object detection models, 2021.


11institutetext: 1Department of Electrical Engineering, Yale University
11email<EMAIL_ADDRESS>
2Department of Radiology and Biomedical Imaging, Yale University
3Department of Biomedical Engineering, Yale University
4Department of Statistics and Data Science, Yale University

# Implicit Anatomical Rendering for Medical Image Segmentation with Stochastic
Experts

Chenyu You 1(✉) Weicheng Dai 22 Yifei Min 44 Lawrence Staib 112233
James S. Duncan 11223344

###### Abstract

Integrating high-level semantically correlated contents and low-level
anatomical features is of central importance in medical image segmentation.
Towards this end, recent deep learning-based medical segmentation methods have
shown great promise in better modeling such information. However, convolution
operators for medical segmentation typically operate on regular grids, which
inherently blur the high-frequency regions, i.e., boundary regions. In this
work, we propose MORSE, a generic implicit neural rendering framework designed
at an anatomical level to assist learning in medical image segmentation. Our
method is motivated by the fact that implicit neural representation has been
shown to be more effective in fitting complex signals and solving computer
graphics problems than discrete grid-based representation. The core of our
approach is to formulate medical image segmentation as a rendering problem in
an end-to-end manner. Specifically, we continuously align the coarse
segmentation prediction with the ambiguous coordinate-based point
representations and aggregate these features to adaptively refine the boundary
region. To parallelly optimize multi-scale pixel-level features, we leverage
the idea from Mixture-of-Expert (MoE) to design and train our MORSE with a
stochastic gating mechanism. Our experiments demonstrate that MORSE can work
well with different medical segmentation backbones, consistently achieving
competitive performance improvements in both 2D and 3D supervised medical
segmentation methods. We also theoretically analyze the superiority of MORSE.

###### Keywords:

Medical Image Segmentation Implicit Neural Representation Stochastic Mixture-
of-Experts.

##  1 Introduction

Medical image segmentation is one of the most fundamental and challenging
tasks in medical image analysis. It aims at classifying each pixel in the
image into an anatomical category. With the success of deep neural networks
(DNNs), medical image segmentation has achieved great progress in assisting
radiologists in contributing to a better disease diagnosis.

Until recently, the field of medical image segmentation has mainly been
dominated by an encoder-decoder architecture, and the existing state-of-the-
art (SOTA) medical segmentation models are roughly categorized into two
groups: (1) convolutional neural networks (CNNs) [25, 4, 19, 1, 29, 38, 11,
12, 6, 34, 33], and (2) Transformers[2, 5, 35]. However, despite their recent
success, several challenges persist to build a robust medical segmentation
model: ❶ Classical deep learning methods require precise pixel/voxel-level
labels to tackle this problem [37, 36, 30, 31, 32]. Acquiring a large-scale
medical dataset with exact pixel- and voxel-level annotations is usually
expensive and time-consuming as it requires extensive clinical expertise [16,
14, 13, 20, 10, 34]. Prior works [15, 7] have used point-level supervision on
medical image segmentation to refine the boundary prediction, where such
supervision requires well-trained model weights and can only capture discrete
representations on the pixel-level grids. ❷ Empirically, it has been observed
that CNNs inherently store the discrete signal values in a grid of pixels or
voxels, which naturally blur the high-frequency anatomical regions, i.e.,
boundary regions. In contrast, implicit neural representations (INRs), also
known as coordinate-based neural representations, are capable of representing
discrete data as instances of a continuous manifold, and have shown remarkable
promise in computer vision and graphics [22, 27, 28]. Several questions then
arise: how many pixel- or voxel-level labels are needed to achieve good
performance? how should those coordinate locations be selected? and how can
the selected coordinates and signal values be leveraged efficiently?

Orthogonally to the popular belief that the model architecture matters the
most in medical segmentation (i.e., complex architectures generally perform
better), this paper focuses on an under-explored and alternative direction:
towards improving segmentation quality via rectifying uncertain coarse
predictions. To this end, we propose a new INR-based framework, MORSE
(iMplicit anatOmical Rendering with Stochastic Experts). The core of our
approach is to formulate medical image segmentation as a rendering problem in
an end-to-end manner. We think of building a generic implicit neural rendering
framework to have fine-grained control of segmentation quality, i.e., to
adaptively compose coordinate-wise point features and rectify uncertain
anatomical regions. Specifically, we encode the sampled coordinate-wise point
features into a continuous space, and then align position and features with
respect to the continuous coordinate.

We further hinge on the idea of mixture-of-experts (MoE) to improve
segmentation quality. Considering our goal is to rectify uncertain coarse
predictions, we regard multi-scale representations from the decoder as
experts. During training, experts are randomly activated for features from
multiple blocks of the decoder, and correspondingly the INRs of multi-scale
representations are separately parameterized by a group of MLPs that compose a
spanning set of the target function class. In this way, the INRs are acquired
across the multi-block structure while the stochastic experts are specified by
the anatomical features at each block.

In summary, our main contributions are as follows: (1) We propose a new
implicit neural rendering framework that has fine-grained control of
segmentation quality by adaptively composing INRs (i.e., coordinate-wise point
features) and rectifying uncertain anatomical regions; (2) We illustrate the
advantage of adopting mixture-of-experts that endows the model with better
specialization of features maps for improving the performance; (3) Extensive
experiments show that our method consistently improves performance compared to
2D and 3D SOTA CNN- and Transformer-based approaches; and (4) Theoretical
analysis verifies the expressiveness of our INR-based model. Code is released
at here.

##  2 Method

Figure 1: Illustration of the MORSE pipeline.

Let us assume a supervised medical segmentation dataset
$\mathcal{D}=\\{(x,y)\\}$, where each input $x=x_{1},x_{2},...,x_{T}$ is a
collection of $T$ 2D/3D scans, and $y$ refers to the ground-truth labels.
Given an input scan ${x}\in\mathbb{R}^{H\times W\times d}$, the goal of
medical segmentation is to predict a segmentation map ${\hat{y}}$. Fig. 1
illustrates the overview of our MORSE. In the following, we first describe our
baseline model $f$ for standard supervised learning, and subsequently present
our MORSE. A baseline segmentation model consists of two main components: (1)
encoder module, which generates the multi-scale feature maps such that the
model is capable of modeling multi-scale local contexts, and (2) decoder
module that makes a prediction $\hat{y}$ using the generated multi-block
features of different resolution. The entire model $M$ is trained end-to-end
using the supervised segmentation loss $\mathcal{L}_{\text{sup}}$ [35] (i.e.,
equal combination of cross-entropy loss and dice loss).

###  2.1 Stochastic Mixture-of-Experts (SMoE) Module

Motivation We want a module that encourages inter- and intra-associations
across multi-block features. Intuitively, multi-block features should be
specified by anatomical features across each block. We posit that due to the
specialization-favored nature of MoE, the model will benefit from explicit use
of its own anatomical features at each block by learning multi-scale
anatomical contexts with adaptively selected experts. In implementation, our
SMoE module follows an MoE design [21], where it treats features from multiple
blocks of the decoder as experts. To mitigate potential overfitting and enable
parameter-efficient property, we further randomly activate experts for each
input during training. Our approach makes three major departures compared to
[21] (i.e., SOTA segmention model): (1) implicitly optimized during training
since it greatly trims down the training cost and the model scale; (2) using
features from the decoder instead of the encoder tailored for our refinement
goal; and (3) empirically showing that “self-slimmable” attribute delivers
sufficiently exploited expressiveness of the model.

Modulization We first use multiple small MLPs with the same size to process
different block features and then up-sample the features to the size of the
input scans, i.e., $H\times W\times d$. With $N$ as the total number of layers
(experts) in the decoder, we treat these upsampled features
$[\bm{F}_{1},\bm{F}_{2},...,\bm{F}_{N}]$ as expert features. We then train a
gating network $\mathcal{G}$ to re-weight the features from activated experts
with the trainable weight matrices $[\bm{W}_{1},\bm{W}_{2},...,\bm{W}_{N}]$,
where $\bm{W}\in\mathbb{R}^{H\times W\times d}$. Specifically, the gating
network or router $\mathcal{G}$ outputs these weight matrices satisfying
$\sum_{i}\bm{W}_{i}=\mathbf{1}^{H\times W\times d}$ using a structure depicted
as follows:

$\bm{W}_{i}=[\texttt{Softmax}(\texttt{Conv}([\bm{F}_{1},\bm{F}_{2},...,\bm{F}_{N}]))]_{i},\quad\textnormal{for}\
\ i\in[N].$ (1)

The gating network first concatenates all the expert features along channels
and uses several convolutional layers to get
$\texttt{Conv}([\bm{F}_{1},\bm{F}_{2},...,\bm{F}_{N}])\in\mathbb{R}^{C\times
H\times W\times d\times N}$, where $C$ is the channel dimension. A softmax
layer is applied over the last dimension (i.e., $N$-expert) to output the
final weight maps. After that, we feed the resultant output $x_{\text{out}}$
to another MLP to fuse multi-block expert features. Finally, the resultant
output $x_{\text{out}}$ (i.e. the coarse feature) is given as follows:

$x_{\text{out}}=\texttt{MLP}(\,\sum_{i=1}^{N}\bm{W}_{i}\cdot\bm{F}_{i}\,),$
(2)

where $\cdot$ denotes the pixel-wise multiplication, and
$x_{\text{out}}\in\mathbb{R}^{C\times H\times W\times d}$.

Stochastic Routing The prior MoE-based model [21] are densely activated. That
is, a model needs to access all its parameters to process all inputs. One
drawback of such design often comes at the prohibitive training cost.
Moreover, the large model size suffers from the representation collapse issue
[26], further limiting the model’s performance. Our proposed SMoE considers
randomly activated expert sub-networks to address the issues. In
implementation, we simply apply standard dropout to multiple experts with a
dropping probability $\alpha$. For each training iteration, there are dropout
masks placed on experts with the probability $\alpha$. That is, the omission
of experts follows a $\texttt{Bernoulli}(\alpha)$ distribution. As for
inference, there is no dropout mask and all experts are activated.

###  2.2 Implicit Anatomical Rendering (IAR)

The existing methods generally assume that the semantically correlated
information and fine anatomical details have been captured and can be used to
obtain high-quality segmentation quality. However, CNNs inherently operate the
discrete signals in a grid of pixels or voxels, which naturally blur the high-
frequency anatomical regions, i.e., boundary regions. To address such issues,
INRs in computer graphics are often used to replace standard discrete
representations with continuous functions parameterized by MLPs [28, 27]. Our
key motivation is that the task of medical segmentation is often framed as a
rendering problem that applies implicit neural functions to continuous
shape/object/scene representations [22, 27]. Inspired by this, we propose an
implicit neural rendering framework to further improve segmentation quality,
i.e., to adaptively compose coordinate-wise point features and rectify
uncertain anatomical regions.

Point Selection Given a coarse segmentation map, the rendering head aims at
rectifying the uncertain boundary regions. A point selection mechanism is thus
required to filter out those pixels where the rendering can achieve maximum
segmentation quality improvement. Besides, point selection can significantly
reduce computational cost compared to blindly rendering all boundary pixels.
Therefore, our MORSE selects $N_{p}$ points for refinement given the coarse
segmentation map using an uncertainty-based criterion. Specifically, MORSE
first uniformly randomly samples $k_{p}N_{p}$ candidates from all pixels where
the hyper-parameter $k_{p}\geq 1$, following [9]. Then, based on the coarse
segmentation map, MORSE chooses $\rho N_{p}$ pixels with the highest
uncertainty from these candidates, where $0.5<\rho<1$. The uncertainty for a
pixel is defined as
$\texttt{SecondLargest}(\mathbf{v})-\texttt{max}(\mathbf{v})$, where
$\mathbf{v}$ is the logit vector of that pixel such that the coarse
segmentation is given as Softmax$(\mathbf{v})$. The rest $(1-\rho)N_{p}$
pixels are sampled uniformly from all the remaining pixels. This mechanism
ensures the selected points contain a large portion of points with uncertain
segmentation which require refinement.

Positional Encoding It is well-known that neural networks can be cast as
universal function approximators, but they are inferior to high-frequency
signals due to their limited learning power [23, 18]. Unlike [9], we explore
using the encoded positional information to capture high-frequency signals,
which echoes our theoretical findings in Appendix 0.A. Specifically, for a
coordinate-based point $(x,y)\in[H]\times[W]$, the positional encoding
function is given as:

$\displaystyle\psi(x,y)$
$\displaystyle=[\sin(2\pi(w_{1}\tilde{x}+v_{1}\tilde{y})),\cdots,\sin(2\pi(w_{L}\tilde{x}+v_{L}\tilde{y})),$
$\displaystyle\qquad\cos(2\pi(w_{1}\tilde{x}+v_{1}\tilde{y})),\cdots,\cos(2\pi(w_{L}\tilde{+}v_{L}\tilde{y}))],$
(3)

where $\tilde{x}=2x/H-1$ and $\tilde{y}=2y/W-1$ are the standardized
coordinates with values in between $[-1,1]$. The frequency
$\\{w_{i},v_{i}\\}_{i=1}^{L}$ are trainable parameters with Gaussian random
initialization, where we set $L=128$ [3]. For each selected point, its
position encoding will then be concatenated with the coarse features of that
point (i.e., $x_{\textnormal{out}}$ defined in Sec. 2.1), to output the fine-
grained features.

Rendering Head The fine-grained features are then fed to the rendering head
whose goal is to rectify the uncertain predictions with respect to these
selected points. Inspired by [9], the rendering head adopts 3-layer MLPs
design. Since the rendering head is designed to rectify the class label of the
selected points, it is trained using the standard cross-entropy loss
$\mathcal{L}_{\text{rend}}$.

Adaptive Weight Adjustment Instead of directly leveraging pre-trained weights,
it is more desirable to train the model from scratch in an end-to-end way. For
instance, we empirically observe that directly using coarse masks by pre-
trained weights to modify unclear anatomical regions might lead to suboptimal
results (See Sec. 3.1). Thus, we propose to modify the importance of
$\mathcal{L}_{\text{rend}}$ as:

$\lambda_{t}=\lambda_{\textnormal{rend}}\cdot\left[\mathbbm{1}\\{t>T/2\\}\cdot\left(\frac{t-T/2}{T}\right)\right],$
(4)

where $t$ is the index of the iteration, $T$ denotes the total number of
iterations, and $\mathbbm{1}\\{\cdot\\}$ denotes the indicator function.

Training Objective As such, the model is trained in an end-to-end manner using
total loss
$\mathcal{L}_{\text{total}}=\mathcal{L}_{\text{sup}}+\lambda_{t}\times\mathcal{L}_{\text{rend}}$.

##  3 Experiments

Table 1: Quantitative comparisons for multi-organ segmentation on the Synapse
multi-organ CT dataset. The best results are indicated in bold.

Method Average Aorta Gallbladder Kidney (L) Kidney (R) Liver Pancreas Spleen
Stomach DSC $\uparrow$ Jaccard $\uparrow$ 95HD $\downarrow$ ASD $\downarrow$
UNet (Baseline) [25] 70.11 59.39 44.69 14.41 84.00 56.70 72.41 62.64 86.98
48.73 81.48 67.96 \+ PointRend [9] 71.52 61.34 43.19 13.70 85.74 57.14 75.42
63.27 87.32 50.16 81.82 71.29 \+ Implicit PointRend [3] 67.33 59.73 52.44
22.15 76.32 51.99 70.28 70.36 81.69 43.77 77.18 67.05 \+ Ours (MoE) 72.83
62.64 40.44 13.15 86.11 59.51 75.81 67.10 87.82 52.11 83.48 70.86 \+ Ours
(SMoE) 74.86 64.94 37.69 12.66 86.39 63.99 77.96 68.93 88.88 53.62 86.12 72.98
\+ Ours (IAR) 73.11 62.98 34.01 12.67 86.28 60.25 76.58 65.34 88.32 52.12
83.47 72.51 \+ Ours (IAR+MoE) 75.37 65.65 33.34 11.43 87.00 64.45 78.14 70.13
89.32 52.33 85.20 76.40 \+ Ours (MORSE) 76.59 66.97 32.00 10.67 87.28 64.73
80.58 71.87 90.04 54.60 86.67 76.93 TransUnet (Baseline) [2] 77.49 64.78 31.69
8.46 87.23 63.13 81.87 77.02 94.08 55.86 85.08 75.62 \+ PointRend [9] 78.30
65.88 34.17 8.62 87.93 63.96 83.47 77.23 94.86 56.45 85.76 76.75 \+ Implicit
PointRend [3] 71.92 60.62 41.42 18.55 78.39 61.64 79.59 73.20 89.61 50.01
80.17 62.75 \+ Ours (MoE) 77.85 65.30 32.75 7.90 87.40 63.46 82.34 77.88 94.14
56.12 85.24 76.25 \+ Ours (SMoE) 78.68 65.98 31.86 7.00 87.60 66.21 82.62
78.12 94.88 57.59 85.97 76.48 \+ Ours (IAR) 79.37 66.50 30.13 7.25 88.63 66.76
83.70 79.50 95.26 57.10 86.90 77.10 \+ Ours (IAR+MoE) 79.60 66.99 27.59 6.54
88.73 66.83 83.85 80.19 95.98 57.12 86.92 77.21 \+ Ours (MORSE) 80.85 68.53
26.61 6.46 88.92 67.53 84.83 81.68 96.83 59.70 87.73 79.58

Dataset We evaluate the models on two important medical segmentation tasks.
(1) Synapse multi-organ
segmentation111https://www.synapse.org/#!Synapse:syn3193805/wiki/217789:
Synapse multi-organ segmentation dataset contains 30 abdominal CT scans with
3779 axial contrast-enhanced abdominal clinical CT images in total. Each
volume scan has variable volume sizes $512\times 512\times 85$ $\sim$
$512\times 512\times 198$ with a voxel spatial resolution of
$([0.54\\!\sim\\!0.54]\times[0.98\\!\sim\\!0.98]\times[2.5\\!\sim\\!5.0])$mm3.
For a fair comparison, the data
split222https://github.com/Beckschen/TransUNet/tree/main/lists/lists_Synapse
is fixed with 18 (2211 axial slices) and 12 patients’ scans for training and
testing, respectively. The entire dataset has a high diversity of aorta,
gallbladder, spleen, left kidney, right kidney, liver, pancreas, spleen, and
stomach.
(2) Liver segmentation: Multi-phasic MRI (MP-MRI) dataset is an in-house
dataset including 20 patients, each including T1 weighted DCE-MRI images at
three-time phases (i.e., pre-contrast, arterial, and venous). Here, our
evaluation is conducted via 5-fold cross-validation on the 60 scans. For each
fold, the training and testing data includes 48 and 12 cases, respectively.

Implementation Details We use AdamW optimizer [17] with an initial learning
rate $5e^{-4}$, and adopt a polynomial-decay learning rate schedule for both
datasets. We train each model for 30K iterations. For Synapse, we adopt the
input resolution as 256$\times$256 and the batch size is 4. For MP-MRI, we
randomly crop 96$\times$96$\times$96 patches and the batch size is 2. For
SMoE, following [21], all the MLPs have hidden dimensions $[256,256]$ with
ReLU activations, the dimension of expert features
$[\bm{F}_{1},\bm{F}_{2},...,\bm{F}_{N}]$ are 256. We empirically set $\alpha$
as 0.7. Following [9], $N_{p}$ is set as 2048, and 8192 for training and
testing, respectively, and $k_{p}$, $\rho$ are 3, 0.75. We follow the same
gating network design [21], which includes four $3\times 3$ convolutional
layers with channels $[256,256,256,N]$ and ReLU activations.
$\lambda_{\text{rend}}$ are set to 0.1. We adopt four representative models,
including UNet [25], TransUnet [2], 3D-UNet [4], UNETR [5]. Specifically, we
set $N$ for UNet [25], TransUnet [2], 3D-UNet [4], UNETR [5] with 5, 3, 3, 3,
respectively. We also use Dice coefficient (DSC), Jaccard, 95% Hausdorff
Distance (95HD), and Average Surface Distance (ASD) to evaluate 3D results. We
conduct all experiments in the same environments with fixed random seeds
(Hardware: Single NVIDIA RTX A6000 GPU; Software: PyTorch 1.12.1+cu116, and
Python 3.9.7).

Table 2: Quantitative comparisons for liver segmentation on the Multi-phasic
MRI dataset. The best results are indicated in bold.

Method Average Method Average DSC $\uparrow$ Jaccard $\uparrow$ 95HD
$\downarrow$ ASD $\downarrow$ DSC $\uparrow$ Jaccard $\uparrow$ 95HD
$\downarrow$ ASD $\downarrow$ 3D-UNet (Baseline) [4] 89.19 81.21 34.97 10.63
UNETR (Baseline) [5] 89.95 82.17 24.64 6.04 \+ PointRend [9] 89.55 81.80 30.88
10.12 \+ PointRend [9] 90.49 82.36 21.06 5.59 \+ Implicit PointRend [3] 88.01
79.83 37.55 12.86 \+ Implicit PointRend [3] 88.72 80.18 26.63 10.58 \+ Ours
(MoE) 89.81 82.06 29.96 10.15 \+ Ours (MoE) 90.70 82.80 15.31 5.93 \+ Ours
(SMoE) 90.16 82.28 28.36 9.79 \+ Ours (SMoE) 91.02 83.29 15.12 5.64 \+ Ours
(IAR) 91.22 83.30 27.84 8.89 \+ Ours (IAR) 91.63 83.83 14.25 4.99 \+ Ours
(IAR+MoE) 92.77 83.94 26.57 7.51 \+ Ours (IAR+MoE) 93.01 84.70 13.29 4.84 \+
Ours (MORSE) 93.59 84.62 19.61 6.57 \+ Ours (MORSE) 93.85 85.53 12.33 4.38

###  3.1 Comparison with State-of-the-Art Methods

We adopt classical CNN- and transformer-based models, i.e., 2D-based {UNet
[25], TransUnet [2]} and 3D-based {3D-UNet [4], UNETR [5]}, and train them on
{2D Synapse, 3D MP-MRI} in an end-to-end manner 333All comparison experiments
are using their released code..

Main Results The results for 2D synapse multi-organ segmentation and 3D liver
segmentation are shown in Tables 1 and 2, respectively. The following
observations can be drawn: (1) Our MORSE demonstrates superior performance
compared to all other training algorithms. Specifically, Compared to UNet,
TransUnet, 3D-UNet, and UNETR baselines, our MORSE with all experts selected
obtains 3.36%$\sim$6.48% improvements in Dice across two segmentation tasks.
It validates the superiority of our proposed MORSE. (2) The stochastic routing
policy shows consistent performance benefits across all four network backbones
on 2D and 3D settings. Specifically, we can observe that our SMoE framework
improves all the baselines, which is within expectation since our model is
implicitly “optimized” given evolved features. (3) As is shown, we can observe
that IAR consistently outperforms PointRend across all the baselines (i.e.,
UNet, TransUnet, 3D-UNet, and UNETR) and obtain {1.59%, 1.07%, 2.03%, 1.14%}
performance boosts on two segmentation tasks, highlighting the effectiveness
of our proposal in INRs. (4) With Implicit PointRend [3] equipped, all the
models’ performances drop. We find: adding Implicit PointRend leads to
significant performance drops of -2.78%, -5.57%, -1.18%, and -1.23%
improvements, compared with the SOTA baselines (i.e., UNet, TransUnet,
3D-UNet, and UNETR) on two segmentation tasks, respectively. Importantly, we
find that: [3] utilizes INRs for producing different parameters of the point
head for each object with point-level supervision. As this implicit function
does not directly optimize the anatomical regions, we attribute this drop to
the introduction of additional noise during training, which leads to the
representation collapse. This further verifies the effectiveness of our
proposed IAR. In Appendix Figs. 2 and 3, we provide visual comparisons from
various models. We can observe that MORSE yields sharper and more accurate
boundary predictions compared to all the other training algorithms.

Visualization of IAR Modules To better understand the IAR module, we visualize
the point features on the coarse prediction and refined prediction after the
IAR module in Appendix Fig. 4. As is shown, we can see that IAR help rectify
the uncertain anatomical regions for improving segmentation quality.

Table 3: Effect of stochastic rate $\alpha$ and expert number $N$.

$\alpha$ |  | DSC[%]$\uparrow$ |  | ASD[voxel]$\downarrow$ |  | $N$ |  | DSC[%]$\uparrow$ |  | ASD[voxel]$\downarrow$
---|---|---|---|---|---|---|---|---|---|---
0.1 |  | 75.41 |  | 11.96 |  | 1 (No MoE) |  | 75.11 |  | 11.67
0.2 |  | 75.68 |  | 11.99 |  | 2 |  | 75.63 |  | 11.49
0.5 |  | 76.06 |  | 10.43 |  | 3 |  | 75.82 |  | 11.34
0.7 |  | 76.59 |  | 10.67 |  | 4 |  | 76.16 |  | 11.06
0.9 |  | 74.16 |  | 11.32 |  | 5 |  | 76.59 |  | 10.67

Table 4: Ablation studies of the Adaptive Weight Adjustment (AWA).

Method |  | DSC[%]$\uparrow$ |  | ASD[voxel]$\downarrow$
---|---|---|---|---
w/o AWA & train w/ $\mathcal{L}_{\text{rend}}$ from scratch |  | 70.56 |  | 14.89
w/o AWA & train w/ $\mathcal{L}_{\text{rend}}$ in $\frac{T}{2}$ |  | 75.42 |  | 12.00
w/ AWA |  | 76.59 |  | 10.67

###  3.2 Ablation Study

We first investigate our MORSE equipped with UNet by varying $\alpha$ (i.e.,
stochastic rate) and $N$ (i.e., experts) on Synapse. The comparison results of
$\alpha$ and $N$ are reported in Table 4. We find that using $\alpha=0.7$
performs the best when the expert capacity is $N=5$. Similarly, when reducing
the expert number, the performance also drops considerably. This shows our
hyperparameter settings are optimal.

Moreover, we conduct experiments to study the importance of Adaptive Weight
Adjustment (AWA). We see that: (1) Disabling AWA and training
$\mathcal{L}_{\text{rend}}$ from scratch causes unsatisfied performance, as
echoed in [9]. (2) Introducing AWA shows a consistent advantage compared to
the other. This demonstrates the importance of the Adaptive Weight Adjustment.

##  4 Conclusion

In this paper, we proposed MORSE, a new implicit neural rendering framework
that has fine-grained control of segmentation quality by adaptively composing
coordinate-wise point features and rectifying uncertain anatomical regions. We
also demonstrate the advantage of leveraging mixture-of-experts that enables
the model with better specialization of features maps for improving the
performance. Extensive empirical studies across various network backbones and
datasets, consistently show the effectiveness of the proposed MORSE.
Theoretical analysis further uncovers the expressiveness of our INR-based
model.

## References

  * [1] Chen, C., Qin, C., Qiu, H., Ouyang, C., Wang, S., Chen, L., Tarroni, G., Bai, W., Rueckert, D.: Realistic adversarial data augmentation for mr image segmentation. In: MICCAI (2020)
  * [2] Chen, J., Lu, Y., Yu, Q., Luo, X., Adeli, E., Wang, Y., Lu, L., Yuille, A.L., Zhou, Y.: Transunet: Transformers make strong encoders for medical image segmentation. arXiv preprint arXiv:2102.04306 (2021)
  * [3] Cheng, B., Parkhi, O., Kirillov, A.: Pointly-supervised instance segmentation. In: CVPR (2022)
  * [4] Çiçek, Ö., Abdulkadir, A., Lienkamp, S.S., Brox, T., Ronneberger, O.: 3d u-net: learning dense volumetric segmentation from sparse annotation. In: MICCAI (2016)
  * [5] Hatamizadeh, A., Tang, Y., Nath, V., Yang, D., Myronenko, A., Landman, B., Roth, H., Xu, D.: Unetr: Transformers for 3d medical image segmentation. In: WACV (2022)
  * [6] He, Y., Lin, F., Tzeng, N.F., et al.: Interpretable minority synthesis for imbalanced classification. In: IJCAI (2021)
  * [7] Huang, R., Lin, M., Dou, H., Lin, Z., Ying, Q., Jia, X., Xu, W., Mei, Z., Yang, X., Dong, Y., et al.: Boundary-rendering network for breast lesion segmentation in ultrasound images. Medical Image Analysis (2022)
  * [8] Jacot, A., Gabriel, F., Hongler, C.: Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems 31 (2018)
  * [9] Kirillov, A., Wu, Y., He, K., Girshick, R.: Pointrend: Image segmentation as rendering. In: CVPR (2020)
  * [10] Lai, Z., Oliveira, L.C., Guo, R., Xu, W., Hu, Z., Mifflin, K., Decarli, C., Cheung, S.C., Chuah, C.N., Dugger, B.N.: Brainsec: Automated brain tissue segmentation pipeline for scalable neuropathological analysis. IEEE Access (2022)
  * [11] Lai, Z., Wang, C., Cheung, S.c., Chuah, C.N.: Sar: Self-adaptive refinement on pseudo labels for multiclass-imbalanced semi-supervised learning. In: CVPR. pp. 4091–4100 (2022)
  * [12] Lai, Z., Wang, C., Gunawan, H., Cheung, S.C.S., Chuah, C.N.: Smoothed adaptive weighting for imbalanced semi-supervised learning: Improve reliability against unknown distribution data. In: ICML. pp. 11828–11843 (2022)
  * [13] Lai, Z., Wang, C., Hu, Z., Dugger, B.N., Cheung, S.C., Chuah, C.N.: A semi-supervised learning for segmentation of gigapixel histopathology images from brain tissues. In: EMBC. IEEE (2021)
  * [14] Lai, Z., Wang, C., Oliveira, L.C., Dugger, B.N., Cheung, S.C., Chuah, C.N.: Joint semi-supervised and active learning for segmentation of gigapixel pathology images with cost-effective labeling. In: ICCV (2021)
  * [15] Li, H., Yang, X., Liang, J., Shi, W., Chen, C., Dou, H., Li, R., Gao, R., Zhou, G., Fang, J., et al.: Contrastive rendering for ultrasound image segmentation. In: MICCAI. Springer (2020)
  * [16] Lin, F., Yuan, X., Peng, L., Tzeng, N.F.: Cascade variational auto-encoder for hierarchical disentanglement. In: ACM CIKM (2022)
  * [17] Loshchilov, I., Hutter, F.: Decoupled weight decay regularization. In: ICLR (2019)
  * [18] Mildenhall, B., Srinivasan, P.P., Tancik, M., Barron, J.T., Ramamoorthi, R., Ng, R.: Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM (2021)
  * [19] Oktay, O., Schlemper, J., Folgoc, L.L., Lee, M., Heinrich, M., Misawa, K., Mori, K., McDonagh, S., Hammerla, N.Y., Kainz, B., et al.: Attention u-net: Learning where to look for the pancreas. arXiv preprint arXiv:1804.03999 (2018)
  * [20] Oliveira, L.C., Lai, Z., Siefkes, H.M., Chuah, C.N.: Generalizable semi-supervised learning strategies for multiple learning tasks using 1-d biomedical signals. In: NeurIPS 2022 Workshop on Learning from Time Series for Health (2022)
  * [21] Ou, Y., Yuan, Y., Huang, X., Wong, S.T., Volpi, J., Wang, J.Z., Wong, K.: Patcher: Patch transformers with mixture of experts for precise medical image segmentation. In: MICCAI. Springer (2022)
  * [22] Park, J.J., Florence, P., Straub, J., Newcombe, R., Lovegrove, S.: Deepsdf: Learning continuous signed distance functions for shape representation. In: CVPR (2019)
  * [23] Rahaman, N., Baratin, A., Arpit, D., Draxler, F., Lin, M., Hamprecht, F., Bengio, Y., Courville, A.: On the spectral bias of neural networks. In: ICML. PMLR (2019)
  * [24] Rahimi, A., Recht, B.: Random features for large-scale kernel machines. Advances in neural information processing systems 20 (2007)
  * [25] Ronneberger, O., Fischer, P., Brox, T.: U-net: Convolutional networks for biomedical image segmentation. In: MICCAI (2015)
  * [26] Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., Dean, J.: Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 (2017)
  * [27] Sitzmann, V., Martel, J., Bergman, A., Lindell, D., Wetzstein, G.: Implicit neural representations with periodic activation functions. In: NeurIPS (2020)
  * [28] Tancik, M., Srinivasan, P., Mildenhall, B., Fridovich-Keil, S., Raghavan, N., Singhal, U., Ramamoorthi, R., Barron, J., Ng, R.: Fourier features let networks learn high frequency functions in low dimensional domains. In: NeurIPS (2020)
  * [29] Xue, Y., Xu, T., Zhang, H., Long, L.R., Huang, X.: Segan: Adversarial network with multi-scale l 1 loss for medical image segmentation. Neuroinformatics (2018)
  * [30] You, C., Dai, W., Liu, F., Su, H., Zhang, X., Staib, L., Duncan, J.S.: Mine your own anatomy: Revisiting medical image segmentation with extremely limited labels. arXiv preprint arXiv:2209.13476 (2022)
  * [31] You, C., Dai, W., Min, Y., Liu, F., Zhang, X., Clifton, D.A., Zhou, S.K., Staib, L.H., Duncan, J.S.: Rethinking semi-supervised medical image segmentation: A variance-reduction perspective. arXiv preprint arXiv:2302.01735 (2023)
  * [32] You, C., Dai, W., Min, Y., Staib, L., Sekhon, J., Duncan, J.S.: Action++: Improving semi-supervised medical image segmentation with adaptive anatomical contrast. arXiv preprint arXiv:2304.02689 (2023)
  * [33] You, C., Dai, W., Staib, L., Duncan, J.S.: Bootstrapping semi-supervised medical image segmentation with anatomical-aware contrastive distillation. In: IPMI (2023)
  * [34] You, C., Yang, J., Chapiro, J., Duncan, J.S.: Unsupervised wasserstein distance guided domain adaptation for 3d multi-domain liver segmentation. In: Interpretable and Annotation-Efficient Learning for Medical Image Computing (2020)
  * [35] You, C., Zhao, R., Liu, F., Chinchali, S., Topcu, U., Staib, L., Duncan, J.S.: Class-aware generative adversarial transformers for medical image segmentation. In: NeurIPS (2022)
  * [36] You, C., Zhao, R., Staib, L.H., Duncan, J.S.: Momentum contrastive voxel-wise representation learning for semi-supervised volumetric medical image segmentation. In: MICCAI (2022)
  * [37] You, C., Zhou, Y., Zhao, R., Staib, L., Duncan, J.S.: Simcvd: Simple contrastive voxel-wise representation distillation for semi-supervised medical image segmentation. IEEE Transactions on Medical Imaging (2022)
  * [38] Zhou, Z., Rahman Siddiquee, M.M., Tajbakhsh, N., Liang, J.: Unet++: A nested u-net architecture for medical image segmentation. In: Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support (2018)

## Appendix

Figure 2: Visualization results on Synpase. MORSE yields more accurate
predictions, especially for small regions. Figure 3: Visualization results on
MP-MRI. MORSE outputs more accurate segmentation results, especially for
boundaries. Figure 4: Visualization of sampled point location. As is shown,
IAR significantly improves the segmentation quality.

##  Appendix 0.A Theoretical Analysis

In this section, we theoretically analyze the expressiveness of MORSE, and
demonstrate the approximation power of INR features combined with MLPs.

We first introduce the kernel method which is commonly used in analyzing
neural networks. We define the kernel function $\kappa$:
$\mathbb{R}^{d}\times\mathbb{R}^{d}\to\mathbb{R}$, which is a symmetric
function measuring the similarity between two vectors in $\mathbb{R}^{d}$.
Kernel functions are used to approximate unknown functions from data.
Specifically, given a training dataset
$\\{\mathbf{x}_{i},\mathbf{y}_{i}\\}_{i=1}^{n}$ where
$\mathbf{y}_{i}=\mu(\mathbf{x}_{i})$ for some unknown function $\mu(\cdot)$,
an estimate $\hat{\mu}$ can be constructed using kernel function as
$\hat{\mu}(\mathbf{x})=\sum_{i=1}^{n}(\mathbf{K}^{-1}\mathbf{Y})_{i}\kappa(\mathbf{x}_{i},\mathbf{x})$
, where $\mathbf{K}$ is the $n\times n$ matrix with
$\mathbf{K}_{i,j}=\kappa(\mathbf{x}_{i},\mathbf{x}_{j})$, and
$\mathbf{Y}=[\mathbf{y}_{1},\cdots,\mathbf{y}_{n}]^{\top}$.

The kernel method is related to MLPs through a kernel function called Neural
Tangent Kernel [8]. For an MLP $\Pi_{\theta}$ with trainable parameters
$\theta$, it has been shown that, under certain conditions, the model
$\Pi_{\theta}$ trained with stochastic gradient descent will converge to the
estimate generated by the kernel method through Neural Tangent Kernel, which
is defined as
$\kappa_{\text{NTK}}(\mathbf{x}_{i},\mathbf{x}_{j})=\mathbb{E}_{\theta\sim\mathcal{N}}\left\langle\frac{\partial\Pi_{\theta}(\mathbf{x}_{i})}{\partial\theta},\frac{\partial\Pi_{\theta}(\mathbf{x}_{j})}{\partial\theta}\right\rangle$.
Importantly, for $\mathbf{x}_{i},\ \mathbf{x}_{j}$ on the unit sphere, NTK is
an inner product kernel, i.e.,

$\displaystyle\kappa_{\text{NTK}}(\mathbf{x}_{i},\mathbf{x}_{j})=k_{\text{NTK}}(\mathbf{x}_{i}\cdot\mathbf{x}_{j}),$
(5)

for some function $k_{\text{NTK}}$.

With the kernel function explained, we now show how positional encoding helps
with expressiveness. By construction, the positional encoding maps pixel
coordinates into sinusoidal vectors. This is closely related to random Fourier
features which are provably able to approximate the family of shift-invariant
kernel functions [24]. Specifically, a kernel function $\kappa$ is shift-
invariant if for any $\mathbf{x}_{1},\ \mathbf{x}_{2}\in\mathbb{R}^{d}$, it
holds that
$\kappa(\mathbf{x}_{1},\mathbf{x}_{2})=\kappa(\mathbf{x}_{1}-\mathbf{x}_{2})$
with a slight abuse of notation. In other words, the value of
$\kappa(\mathbf{x}_{1},\mathbf{x}_{2})$ only depends on the difference
$\mathbf{x}_{1}-\mathbf{x}_{2}$. It is clear that shift-invariance is an ideal
property in imaging tasks.

To understand how the Fourier features generated by positional encoding can
approximate shift-invariant kernel, we rewrite positional encoding (Eqn. 2.2)
as:

$\displaystyle\psi(\mathbf{x})=[\sin(2\pi\mathbf{w}_{1}\cdot\mathbf{x}),\cdots,\
\sin(2\pi\mathbf{w}_{L}\cdot\mathbf{x}),\cos(2\pi\mathbf{w}_{1}\cdot\mathbf{x}),\cdots,\
\cos(2\pi\mathbf{w}_{L}\cdot\mathbf{x})],$

where $\mathbf{x}=(\tilde{x},\tilde{y})$ , and $\mathbf{w}_{i}=[w_{i},v_{i}]$
for $i=1,\cdots,L$. Given any shift-invariant kernel
$\kappa(\mathbf{x}_{1},\mathbf{x}_{2})=\kappa(\mathbf{x}_{1}-\mathbf{x}_{2})$,
we define a distribution $\mathcal{P}$ over $\mathbf{w}$ as
$\mathcal{P}(\mathbf{w})=\frac{1}{2\pi}\int e^{-2\pi
i\mathbf{w}^{\top}\mathbf{x}}\kappa(\mathbf{x})d\mathbf{x}$, which is the
Fourier transform of the kernel $\kappa$. Suppose that
$\mathbf{w}_{1},\cdots,\mathbf{w}_{L}$ are $i.i.d.$ samples from
$\mathcal{P}$. Then it holds that [24]:

$\displaystyle\textnormal{Pr}\left[\sup_{\mathbf{x}_{1},\mathbf{x}_{2}}\left|\frac{1}{L}\psi(\mathbf{x}_{1})^{\top}\psi(\mathbf{x}_{2})-\kappa(\mathbf{x}_{1},\mathbf{x}_{2})\right|\geq\epsilon\right]\leq
2^{8}\frac{\sigma_{\mathcal{P}}^{2}}{\epsilon^{2}}\exp\left(-\frac{L\epsilon^{2}}{16}\right),$

where
$\sigma_{\mathcal{P}}^{2}=\mathbb{E}_{\mathcal{P}}(\mathbf{w}^{\top}\mathbf{w})$.
The above result indicates that with high probability, any shift-invariant
kernel function can be approximated with Fourier features. Therefore, this
demonstrates the expressive power of INR features.

Finally, we show that $\psi(\cdot)$ combined with the MLP $\Pi_{\theta}$ forms
a shift-invariant kernel. We define the positional encoding kernel
$\kappa_{\text{pe}}$ as
$\kappa_{\text{pe}}(\mathbf{x}_{1},\mathbf{x}_{2})=\psi(\mathbf{x}_{1})^{\top}\psi(\mathbf{x}_{2})$.
It can be shown that
$\kappa_{\text{pe}}(\mathbf{x}_{1},\mathbf{x}_{2})=\sum_{j=1}^{L}\cos(2\pi\mathbf{w}_{j}\cdot(\mathbf{x}_{1}-\mathbf{x}_{2}))$
[28], which is shift-invariant. Combining with Eqn. 5, our positional encoding
combined followed by MLP approximately equals to
$k_{\text{NTK}}(\kappa_{\text{pe}}(\mathbf{x}_{1}-\mathbf{x}_{2}))$, which is
a shift-invariant kernel. This demonstrates the expressiveness of MORSE.