# Amazon Last-Mile Delivery Trajectory Prediction Using Hierarchical TSP with
Customized Cost Matrix

Xiaotong Guo, Baichuan Mo, Qingyi Wang 111In alphabetical order by the last
name Department of Civil and Environmental Engineering, Massachusetts
Institute of Technology  77 Massachusetts Ave, Cambridge, MA, USA

###### Abstract

In response to the Amazon Last-Mile Routing Challenge, Team Permission Denied
proposes a hierarchical Travelling Salesman Problem (TSP) optimization with a
customized cost matrix. The higher level TSP solves for the zone sequence
while the lower level TSP solves the intra-zonal stop sequence. The cost
matrix is modified to account for routing patterns beyond the shortest travel
time. Lastly, some post-processing is done to edit the sequence to match
commonly observed routing patterns, such as when travel times are similar,
drivers usually start with stops with more packages than those with fewer
packages. The model is tested on 1223 routes that are randomly selected out of
the training set and the score is $0.0381$. On the 13 routes in the given
model apply set, the score was $0.0375$.

##  1 Introduction

This report presents the thought processes, selected methodology, and expected
results of the Amazon Last-Mile Routing Research Challenge by Team Permission
Denied. In summary, the team went through four phases before arriving at the
final submission.

Descriptive Analysis: Upon receiving the challenge, a thorough descriptive
analysis is done. The first important finding is that, in most circumstances,
the drivers finish all deliveries in one zone before moving on to the stops in
another zone. This rule is only broken when backtracking exists. A further
look at the scores confirms this intuition: assuming the zone sequence and
intra-zonal stop sequence are correct, the loss on the score due to certain
zones being revisited is only 0.009. If the zone sequence is correct and the
stops in each zone are shuffled, the average score is around 0.02. Therefore,
getting the zone sequence correct is the most important, and the team decides
to adopt a hierarchical approach: solving for the zone sequence, and then the
intra-zonal stop sequence. This greatly reduces the scale of the problem since
the majority of the routes have around 150 stops (up to 250), but the number
of zones is between 6 and 47. Second, the zonal transitional probabilities are
investigated. As most of the zones only appear in the training set once, an
attempt at a frequency tabulation is not successful. On the other hand, 74% of
the zonal transitions select the zone that is closest by travel time, making
the step-by-step prediction algorithm potentially successful. Next, the
correlation between package dimensions, package counts, delivery time windows,
and sequence order is investigated but no apparent relationship is found.

Benchmarking: A benchmark model is created to establish an idea of the
solution quality and expected performance. Since most drivers follow the
navigation given by Amazon, a shortest-distance tour becomes a natural
benchmark. The team solves a tour-based (where the start and end stations are
both INIT) to generate zone sequences and a path-based (where the distance
from the last zone to INIT is not counted) Travelling Salesman Problem (TSP)
to generate intra-zonal stop sequences as benchmarks. Inside each zone, a
path-based TSP is generated from the stop closest to the last zone to the stop
closest to the next zone.

Model Attempts: Both naive TSP solutions achieve scores reasonable scores
(around 0.06). To improve the performance, machine learning models are
attempted. First, it is noticed that correctly predicting the first zone would
significantly improve the TSP performance, therefore a neural network is
constructed to predict the first zone based on the travel time, distance,
package count and size, etc. Second, pure machine learning models to generate
sequences are investigated, including myopic approaches that predict the next
element based on previously predicted stops, as well as sequence-to-sequence
(seq2seq) approaches that encode and decode the entire sequence. Third,
different training methods are considered, including the traditional cross-
entropy loss, customized weighted loss, as well as reinforcement learning
using policy gradients. Lastly, some improvements are made to the benchmark
TSP models by adding penalty costs to non-consecutive zone-ids. Due to the
small sample size (6k), machine learning techniques cannot outperform the
benchmark models. After experimenting with various modeling techniques, the
team decides to use the TSP solution as the final submission.

Hyperparameter Searching and Post-Processing: The customized cost matrix
involves hyperparameters that the team searched for over the given training
set. Lastly, some post-processing patterns are identified to further improve
the quality of our solution.

The highlights of the final submitted model are:

  * 1.

Hierarchical modeling - To reduce the size of each optimization problem, the
problem is broken down into zone routing and intra-zonal stop routing.

  * 2.

Customized TSP cost matrix - To account for considerations in addition to
shortest distance, the cost matrix is modified and the TSP performance
improved by almost 0.01.

  * 3.

Post-processing to match behavioral patterns - Some TSP sequences are reversed
to accommodate delivery patterns such as stops with more packages are visited
first instead of last, all else being equal.

  * 4.

Stable hyperparameters - The cost hyperparameters have good generalizability
and do not require re-training.

The rest of the technical report reviews the relevant literature and its
compatibility with the research question; describes the selected model in
detail, and discusses the expected results.

##  2 Literature Review

This problem is essentially a vehicle routing problem, except that the
traditional setup for vehicle routing problems aims for the shortest distance
traveled, but the problem of interest looks for the most similarity with the
observed sequence. Two research communities have extensively studied the
vehicle routing problem: machine learning and operations research. Literature
in both communities is reviewed, with the pros and cons of the algorithms
discussed for the problem of interest.

###  2.1 Operations Research

Given a set of locations one would like to visit, a Traveling Salesman Problem
(TSP) can be solved to find the route with the minimum cost or distance. The
overview and history of the TSP can be found in Applegate et al. [2011].
Although TSP is a well-known NP-hard problem in combinatorial optimization,
off-the-shelf integer optimization solvers (e.g., Gurobi and GLPK) are able to
solve it efficiently for real-world instances. One key approach we utilized
when solving the TSP is the cutting-plane method [Marchand et al., 2002],
which is initially applied to TSP by Dantzig et al. [1954].

###  2.2 Machine Learning

Two types of architectures can be used to re-order the input sequence: step-
by-step or sequence-to-sequence (seq2seq). Step-by-step prediction involves
predicting the stops one by one, given the information from previous stops, as
well as candidate stops. Since the information from candidate stops are
crucial, feed-forward neural networks are not a good candidate since it does
not attribute features to candidates. Instead, a feed-forward neural network
with alternative-specific utility is adopted [Wang et al., 2020]. This
architecture draws the connection between discrete choice models with neural
networks and uses neural networks to generate the utility for each candidate,
and the candidate with the highest ’utility’ is chosen. A sequence is then
formed by repeatedly feeding the selected stop into the algorithm to get the
next stop until the end of the sequence is reached. The advantage of this
algorithm is that it is at the stop level instead of the sequence level.
Therefore, the sample size, which is critical for the success of machine
learning algorithms, is significantly larger than the seq2seq models. The
disadvantage of this algorithm is that it is myopic and only sees the next
step candidates while making a selection.

In recent years, a lot of seq2seq prediction algorithms have been developed,
mainly for natural language processing (NLP) tasks. Compared to step-by-step
prediction, seq2seq models comprise an encoder and a decoder. All elements in
the sequence are encoded before decoding starts, therefore a global view is
attained. The architecture of encoder and decoder often involves variants of
the recurrent neural networks (ex. long-short term memory networks) [Sutskever
et al., 2014], or attention [Vaswani et al., 2017]. Most seq2seq problems are
considered with mapping one sequence to another, whereas the problem of
interest is concerned with re-ordering the input sequence. Pointer network is
proposed to solve this type of problem, where the decoder uses self-attention
to point to one of the input elements [Vinyals et al., 2015]. The authors used
a pointer network to solve TSP and achieved similar performance to TSP
solvers. One drawback of the original pointer network is that it is sensitive
to the order of inputs. The authors, therefore, added another encoding module
to eliminate this influence [Vinyals et al., 2016]. However, in our
experiments, this dependency can be leveraged by arranging the input set in a
meaningful sequence to improve performance. For example, ordering the input
stops according to the TSP sequence would accelerate model convergence and
improve the score. However, in the papers presented above, 1M training samples
were fed into the network. Given that the training set only contains 6000
routes, score improvements on TSP solutions are unsuccessful.

The original pointer network uses cross-entropy loss (supervised learning). In
this problem, the cross-entropy loss is very inefficient due to the way the
score is calculated, since the loss only considers the probability of the
correct position, and the loss for predicting all other positions is the same.
But the scoring function considers similarity in addition to correctness. The
scoring function is not differentiable and cannot be directly used as the loss
function and use gradient descent. An alternative training method is
reinforcement learning based on policy gradients [Ma et al., 2019, Bello et
al., 2019]. Using the well-known REINFORCE algorithm, we can directly optimize
the non-differentiable score function. Researchers have found that this method
has the same sample efficiency and better generalizability for TSP problems
compared to supervised learning [Joshi et al., 2019]. However, training with
reinforcement learning in this particular problem with the sample size and
given information also does not outperform TSP solutions.

###  2.3 Proposed Method

Our proposed method is built upon the traditional TSP with a customized
distance matrix that implicitly contains drivers’ routing behaviors for the
Amazon last-mile delivery. Compared to the existing TSP framework, which
minimizes the total vehicle travel distance, we modified the distance matrix
and generated optimal routes which minimized the total adjusted travel
distance.

##  3 Methodology

###  3.1 Data

We observe that most of the drivers tend to visit all stops in a zone before
going to the next zone. Hence, we divide the problem into two sub-problems.
The first is to identify the zone sequence, and the second is to recognize the
intra-zonal stop sequence.

The actual zone sequence is generated based on the order of each zone’s first
appearance. An example is shown in Figure 1. For stops without zone id (due to
missing data), we fill them with the zone ID of its (travel time-based)
nearest stop.

Three important properties are noticed while observing the zone sequences:

  * 1.

Most likely, the driver would finish a “major zone” first, then move to the
next “major zone”. A major zone is defined as the zone ID before the dot. For
example, the major zone for “A-2.2A” is “A-2”. For example, in Figure 1, the
driver first finishes major zone “A-2”, then “A-1”, finally “P-13”.

  * 2.

Within a specific major zone, two adjacent “inner zone” ids are most likely
have a “difference of one”. The “inner zone” is defined as the zone ID after
the dot. For example, the inner zone for “A-2.2A” is “2A”. The “difference of
one” is defined as follows. Given two inner zone IDs “XY” and “AB”, where X
and A are numbers and Y and B are characters, we have

$\displaystyle|X-A|+|\texttt{ord}(Y)-\texttt{ord}(B)|=1$ (1)

where $\texttt{ord}(\cdot)$ function returns an integer representing the
Unicode character. For example, “1A” and “1B” has a difference of one, so as
“1A” and “2A”. But “1A” and “2B” has a difference of two.

  * 3.

When a driver finishes a “major zone” and move to another, the two adjacent
major zone IDs are most likely to have a “difference of one”. For example, in
Figure 1, the driver first finishes major zone “A-2”, then “A-1”. Those two
major zone IDs have a difference of one.

Figure 1: Example of zone sequence. “INIT” indicates the delivery station

To validate these three properties, we calculate the frequency that these
rules hold in the data set. For all adjacent zone ID pairs, 87.67% of them
have the same major zone ID (Property 1). For all adjacent zone ID pairs
within a specific major zone, 82.49% of them have a “difference of one”
(Property 2). For all adjacent zone ID pairs with major zone ID changes,
96.17% of these changes lead to a “difference of one” between two major zone
IDs (Property 3). These statistics support the three properties, which implies
that the zone ID includes a lot of information for the sequence estimation.

Another information we use is the planned service time and package volumes.
Details on how these are utilized are shown in Section 3.3.

We also collected outside data sources from OpenStreetMap. Specifically, we
extract the number of traffic signals and highway ramps around every stop.
Unfortunately, this does not help to improve our model, thus is dropped from
our final submission.

For the model’s validation, we randomly separate the 6,112 routes into a
training data set (4,889 routes) and a testing data set (1,223 routes), though
our proposed solution does not require a training process.

###  3.2 Travelling Salesman Problem Formulation

With the observation that drivers visit all stops within the same zone first
and then move to the next zone, we solve a standard TSP with a modified travel
time matrix to generate zone sequence first and then solve multiple path-TSP
to identify intra-zonal stop sequence.

First, we provide the formulation of the standard TSP solved for generating
zone sequences. For a route instance with $n$ zones, the set of zones is
indexed by $[n]=\\{1,...,n\\}$ and the initial station location is indicated
by index $0$. Let $V$ represent the set of all locations that need to be
visited including the initial station, i.e., $V=\\{0,1,...,n\\}$. $t_{ij}$
denotes the travel time between any two locations, i.e., $\forall i\neq j\in
V$. The travel time between any two zones is calculated as the average travel
time between all possible pairs of stops between two zones. The decision
variable for this problem is $x_{ij}\in\\{0,1\\},\;\forall i,j\in V$.
$x_{ij}=1$ indicates that the driver will visit to the location $j$ after
visiting $i$. Then, the TSP problem can be formulated as:

$\displaystyle\min\quad$
$\displaystyle\sum_{i=0}^{n}\sum_{j=0}^{n}t_{ij}x_{ij}$ (2a) s.t.
$\displaystyle\sum_{i=0}^{n}x_{ij}=1\quad\forall j\in V$ (2b)
$\displaystyle\sum_{j=0}^{n}x_{ij}=1\quad\forall i\in V$ (2c)
$\displaystyle\sum_{i\in S}\sum_{j\notin S}x_{ij}\geq 1\quad\forall S\subset
V,S\neq\emptyset,V$ (2d) $\displaystyle\sum_{i\notin S}\sum_{j\in S}x_{ij}\geq
1\quad\forall S\subset V,S\neq\emptyset,V$ (2e) $\displaystyle
x_{ii}=0\quad\forall i\in V$ (2f) $\displaystyle
x_{ij}\in\\{0,1\\}\quad\forall i,j\in V$ (2g)

Where the objective (2a) minimizes the total travel time for the tour.
Constraints (2b) and (2c) make sure that each visited location has exactly one
predecessor and one successor in the optimal tour. Constraints (2d) and (2e)
are proposed to eliminate subtours in the optimal tour. Constraints (2f) avoid
self loops and constraints (2g) guarantee decision variables are binary.

The problem $(\ref{eq:Tour_TSP})$ is an Integer Linear Programming (ILP) with
exponential number of constraints due to constraints (2d) and (2e). To solve
this problem efficiently, we implemented both constraints (2d) and (2e) as
lazy constraints, indicating they are only added to the problem if subtours
are identified in the current optimal solution.

To account for the observations made in the zone sequence (Section 3.1), we
propose three heuristics to modify the travel time matrix, which is the input
for generating the optimal zone sequence.

  1. 1.

For travel time from the initial station to a zone $i$, if the zone is not
within either i) $h$ closest zones from the initial station regarding travel
times or ii) $h$ closest zones from the initial station regarding Euclidean
distances, we modify the travel time to $t_{0i}*\alpha$, where $\alpha$ and
$h$ are both parameters for the first proposed heuristic approach.

  2. 2.

For travel time between any two zones $i$ and $j$, if zone $i$ and zone $j$
are not from the same ”major zone”, we modify the travel time to
$t_{ij}*\beta$, where $\beta$ is the parameter for the second proposed
heuristic approach.

  3. 3.

For travel time between any two zones $i$ and $j$, if they are from the
identical ”major zone” and the difference between their zone ID after the dot
does not equal to 1, we modify the travel time to $t_{ij}*\gamma$, where
$\gamma$ is the parameter for the third proposed heuristic approach.

In the final submitted algorithm, we used the grid search approach to finalize
values for all four heuristic parameters: $h=9$, $\alpha=1.04$, $\beta=3.8$,
$\gamma=2.5$.

Solving the problem (2) with the modified travel time matrix leads to the
optimal zone sequence222Without loss of generality, we can assume the sequence
starts from the initial station indexed by $0$. $S^{*}=(0,s_{1},...,s_{n})$,
where $s_{i}$ indicates the $i$-th zone visited in the optimal sequence after
departing from the initial station. Then we solve the intra-zonal stop
sequence using path-based TSP. Given a set of locations $V$ need to be visited
and the starting location $v_{o}$ and the ending location $v_{d}$, we can
formulate the path-TSP problem as follows:

$\displaystyle\min\quad$
$\displaystyle\sum_{i=0}^{n}\sum_{j=0}^{n}t_{ij}x_{ij}$ (3a) s.t.
$\displaystyle\sum_{i=0}^{n}x_{ij}=1\quad\forall j\in
V\setminus\\{v_{o},v_{d}\\}$ (3b)
$\displaystyle\sum_{j=0}^{n}x_{ij}=1\quad\forall i\in
V\setminus\\{v_{o},v_{d}\\}$ (3c) $\displaystyle\sum_{j\in
V}x_{v_{o}j}=\sum_{i\in V}x_{iv_{d}}=1$ (3d) $\displaystyle\sum_{j\in
V}x_{v_{d}j}=\sum_{i\in V}x_{iv_{o}}=0$ (3e) $\displaystyle\sum_{i\in
S}\sum_{j\notin S}x_{ij}\geq 1\quad\forall S\subset V,S\neq\emptyset,V$ (3f)
$\displaystyle\sum_{i\notin S}\sum_{j\in S}x_{ij}\geq 1\quad\forall S\subset
V,S\neq\emptyset,V$ (3g) $\displaystyle x_{ii}=0\quad\forall i\in V$ (3h)
$\displaystyle x_{ij}\in\\{0,1\\}\quad\forall i,j\in V$ (3i)

The path-TSP problem (3) is similar to the standard TSP problem (2) except
that there will be no predecessors for the starting location $v_{o}$ and no
successors for the ending location $v_{d}$, indicating by constraints (3d) and
(3e). The complete sequence is generated according to Algorithm 1 based on
generated zone sequence, where a heuristic parameter $k=3$ is utilized in the
final implementation.

Algorithm 1 Complete sequence generation based on the generated zone sequence.
Input: optimal zone sequence $S^{*}=(0,s_{1},...,s_{n})$, heuristic parameter
$k$.

1:function CompletePathGeneration($S^{*}$)

2: $S^{*}_{complete}\leftarrow\\{0\\}$ $\triangleright$ Initialize the
complete sequence with the initial station

3: for $s_{i}=s_{1},...,s_{n}$ do

4: Find the previous visited zone $s_{i-1}$ and the next visited zone
$s_{i+1}$

5: Calculate the average travel time between any stop $v\in s_{i}$ to all
stops in zone $s_{i-1}$ and zone $s_{i+1}$

6: Find $k$ nearest stops in zone $s_{i}$ regarding to zone $s_{i-1}$ as the
set $M$

7: Find $k$ nearest stops in zone $s_{i}$ regarding to zone $s_{i+1}$ as the
set $N$

8: Solve $k^{2}$ path-TSP (3) between any pair of stops in $M\times N$.

9: Let the path $S^{*}_{i}$ with the minimum travel time as the optimal
sequence of zone $i$

10: Append the sequence $S^{*}_{i}$ to the complete sequence
$S^{*}_{complete}$

11: return $S^{*}_{complete}$

It is worth mentioning that all TSP instances are solved with the open-source
ILP solver GLPK implemented with programming language Julia [Bezanson et al.,
2017] and optimization package JuMP [Dunning et al., 2017]. After generating
the complete stop sequence $S^{*}_{complete}$, we enter the post-processing
stage to further improve sequence performances.

###  3.3 Post-Processing

After solving the stop sequence by TSP, we observe that most of the high-score
(i.e., low performance) routes are due to partially or fully reverse of the
sequence (i.e., a sequence A-B-C-D is erroneously estimated as D-C-B-A).
Hence, we propose a post-processing method to correct the erroneous estimation
due to reversal.

We observe two properties from the data set:

  * 1.

Most of the drivers tend to serve the business areas first. The potential
reason may be that it also takes a longer time to deliver packages in a
business building. Serving them first can make the total service time more
controllable at the end of the journey. Hence, we expect that the planned
service time at the first several stops is larger than that of the last
several stops.

  * 2.

Most of the drivers tend to deliver large-size packages first. This may be
because carrying large-size packages in the vehicle is not fuel-efficient.

Based on these properties, for every generated stop sequence by TSP, we check
whether we need to reverse it. Given a generated route $i$, let $p^{+}_{i}$
(resp. $p^{-}_{i}$) be the average planned service time of the first (resp.
last) $p\%$ stops in route $i$. We will reverse route $i$ if

$\displaystyle\frac{p^{-}_{i}}{p^{+}_{i}}\geq\theta,$ (4)

where $p$ and $\theta$ are hyperparameters representing the proportion of
stops and a threshold. We set $p=15$ and $\theta=1.22$ based on cross-
validation on the test set. Eq. 5 means that in a generated sequence if the
planned service time for the last several stops is too large, we may have the
reversal error and need to correct it by reverse the whole sequence.

After process by Eq. 5, we fixed all sequences that are already reversed. For
the remaining sequences, we further check whether they need to be reversed
based on package volumes. Specifically, given a generated route $i$, let
$v^{+}_{i}$ (resp. $v^{-}_{i}$) be the total package columns
(depth$\times$width$\times$height) of the first (resp. last) 15% stops in
route $i$. We will reverse route $i$ if

$\displaystyle\frac{v^{-}_{i}}{v^{+}_{i}}\geq\eta,$ (5)

where $\eta=3$ is used.

After post-processing, a sequence validity check is performed. Specifically,
we check whether the first stop of the estimated sequence is the delivery
station, and whether the estimated sequence has the same stop IDs as the
actual one. If either of these two criteria does not hold, we return a
sequence by simply sort the stops using zone IDs, which ensures that stops
with the same zone IDs are close to each other.

##  4 Results and Conclusions

###  4.1 Performance

Although the submitted formulation does not require model training, we have
separated the given training set into the training (4889) and test set (1223)
for self-evaluation of the machine learning models. Therefore, all self-
evaluation is done over the test set. To reduce the evaluation time, we
implemented the scoring function using Cython. Compared to the evaluation code
in Python provided by the challenge host team, our implementation evaluates
the same set of routes by using only one-third of the computation time.

Figure 2 shows the score distribution generated by our final algorithm. The
route performance score follows an exponential distribution and most routes
have a score below 0.1. The average route score is $0.0381$ for these 1223
testing routes. On the 13 routes in the given model apply set, the score was
$0.0375$.

Figure 2: Route score performances.

###  4.2 Discussion

Zone sequence dominates score. We observe that, if the zone sequence is
perfectly predicted, even if the stop IDs within a zone are shuffled, the
average route score can reach $0.0206$. Hence, most of our jobs focus on
predicting the zone sequence, instead of the stop sequence.

The three properties of zone IDs (see Section 3.1) may imply that drivers most
likely follow the planned route and seldom deviate. As the zone ID is used to
“help simplify the route planning process” (quoted from Slack Q&A), we believe
that Amazon plans the route in a way that the zone IDs exhibit clear patterns.
So the major challenge of this problem is to recover how Amazon plans the
routes. This explains why TSP works better than machine learning methods under
the given current information and sample size.

The reversal problem remains. Figure 3 shows an example of reverse prediction.
Since we are not able to increase the first-zone prediction accuracy beyond
35%, after post-processing, the reverse issues still exist. The post-
processing reduces our score on our test set from 0.391 to 0.381. However, if
we can have a 100% correction rate for the reversal problems (i.e., always use
the one with a smaller score), the score can reduce to 0.280, indicating that
some further correction methods are needed. Note that we have tried to use the
number of surrounding highway ramps as a new indicator, as well as using
machine learning to predict the first zone, but it does not increase the model
performance.

(a) Actual route

(b) Predicted route by TSP

Figure 3: Examples of reverse prediction

## References

  * Applegate et al. [2011] Applegate, D. L., Bixby, R. E., ChvÃ¡tal, V., & Cook, W. J. (2011).  The Traveling Salesman Problem.  Princeton University Press.  URL: https://doi.org/10.1515/9781400841103. doi:doi:10.1515/9781400841103.
  * Bello et al. [2019] Bello, I., Pham, H., Le, Q. V., Norouzi, M., & Bengio, S. (2019).  Neural combinatorial optimization with reinforcement learning.  5th International Conference on Learning Representations, ICLR 2017 - Workshop Track Proceedings, (pp. 1–15). arXiv:1611.09940.
  * Bezanson et al. [2017] Bezanson, J., Edelman, A., Karpinski, S., & Shah, V. B. (2017).  Julia: A fresh approach to numerical computing.  SIAM Review, 59, 65--98. URL: http://www.siam.org/journals/sirev/59-1/100067.html. doi:10.1137/141000671. arXiv:1411.1607.
  * Dantzig et al. [1954] Dantzig, G., Fulkerson, R., & Johnson, S. (1954).  Solution of a large-scale traveling-salesman problem.  Journal of the Operations Research Society of America, 2, 393--410. URL: http://www.jstor.org/stable/166695.
  * Dunning et al. [2017] Dunning, I., Huchette, J., & Lubin, M. (2017).  Jump: A modeling language for mathematical optimization.  SIAM Review, 59, 295--320. doi:10.1137/15M1020575.
  * Joshi et al. [2019] Joshi, C. K., Laurent, T., & Bresson, X. (2019).  On Learning Paradigms for the Travelling Salesman Problem, .  (pp. 1--9). URL: http://arxiv.org/abs/1910.07210. arXiv:1910.07210.
  * Ma et al. [2019] Ma, Q., Ge, S., He, D., Thaker, D., & Drori, I. (2019).  Combinatorial Optimization by Graph Pointer Networks and Hierarchical Reinforcement Learning, .  URL: http://arxiv.org/abs/1911.04936. arXiv:1911.04936.
  * Marchand et al. [2002] Marchand, H., Martin, A., Weismantel, R., & Wolsey, L. (2002).  Cutting planes in integer and mixed integer programming.  Discrete Applied Mathematics, 123, 397–446. doi:10.1016/S0166-218X(01)00348-1.
  * Sutskever et al. [2014] Sutskever, I., Vinyals, O., & Le, Q. V. (2014).  Sequence to sequence learning with neural networks.  In Proceedings of the 27th International Conference on Neural Information Processing Systems - Volume 2 NIPS’14 (p. 3104–3112).  Cambridge, MA, USA: MIT Press.
  * Vaswani et al. [2017] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., & Polosukhin, I. (2017).  Attention Is All You Need.  Advances in Neural Information Processing Systems, 2017-Decem, 5999--6009. URL: http://arxiv.org/abs/1706.03762. arXiv:1706.03762.
  * Vinyals et al. [2016] Vinyals, O., Bengio, S., & Kudlur, M. (2016).  Order matters: Sequence to sequence for sets.  4th International Conference on Learning Representations, ICLR 2016 - Conference Track Proceedings, (pp. 1--11). arXiv:1511.06391.
  * Vinyals et al. [2015] Vinyals, O., Meire, F., & Navdeep, J. (2015).  Pointer Networks.  Advances in Neural Information Processing Systems, (pp. 1--9).
  * Wang et al. [2020] Wang, S., Mo, B., & Zhao, J. (2020).  Deep neural networks for choice analysis: Architecture design with alternative-specific utility functions.  Transportation Research Part C: Emerging Technologies, 112, 234--251. URL: https://www.sciencedirect.com/science/article/pii/S0968090X19310381. doi:https://doi.org/10.1016/j.trc.2020.01.012.


11institutetext: ******************11institutetext: Department of Robotics and
Mechatronics Engineering, DGIST, Korea 22institutetext: Department of
Pathology, Asan Medical Center
33institutetext: Department of Pathology, University of Ulsan College of
Medicine
33email<EMAIL_ADDRESS><EMAIL_ADDRESS>
# Feature Re-calibration based Multiple Instance Learning for Whole Slide
Image Classification

Philip Chikontwe 11 Soo Jeong Nam 22 Heounjeong Go 2233 Meejeong Kim 22 Hyun
Jung Sung 22 Sang Hyun Park 11

###### Abstract

Whole slide image (WSI) classification is a fundamental task for the diagnosis
and treatment of diseases; but, curation of accurate labels is time-consuming
and limits the application of fully-supervised methods. To address this,
multiple instance learning (MIL) is a popular method that poses classification
as a weakly supervised learning task with slide-level labels only. While
current MIL methods apply variants of the attention mechanism to re-weight
instance features with stronger models, scant attention is paid to the
properties of the data distribution. In this work, we propose to re-calibrate
the distribution of a WSI bag (instances) by using the statistics of the max-
instance (critical) feature. We assume that in binary MIL, positive bags have
larger feature magnitudes than negatives, thus we can enforce the model to
maximize the discrepancy between bags with a metric feature loss that models
positive bags as out-of-distribution. To achieve this, unlike existing MIL
methods that use single-batch training modes, we propose balanced-batch
sampling to effectively use the feature loss i.e., (+/-) bags simultaneously.
Further, we employ a position encoding module (PEM) to model
spatial/morphological information, and perform pooling by multi-head self-
attention (PSMA) with a Transformer encoder. Experimental results on existing
benchmark datasets show our approach is effective and improves over state-of-
the-art MIL methods. https://github.com/PhilipChicco/FRMIL

##  1 Introduction

Histopathology image analysis (HIA) is an important task in modern medicine
and is the gold standard for cancer detection and treatment planning[17]. The
development of whole slide image (WSI) scanners has enabled the digitization
of tissue biopsies into gigapixel images and paved the way for the application
of machine learning techniques in the field of digital pathology[3, 22].
However, employing popular convolutional neural network (CNN) architectures
for varied tasks in HIA is non trivial and has several challenges, ranging
from the large size of WSIs and extreme high resolution to lack of precise
labeling and stain color variations[22]. This motivates the need for memory
efficient methods that mitigate the need for fine-grained labels and are
fairly interpretable. To address this, multiple instance learning (MIL)[31, 1]
is a popular formulation that considers diagnosis as a weakly supervised
learning problem[29].

Through the recent advances in deep learning[30, 16], MIL based
histopathology[8, 13, 27, 25, 10] analysis has achieved notable success[11,
24, 28, 18]. For instance, Li et al.[21] introduced non-local attention to re-
weight instances relative the highest scoring instance (critical) in a bag,
proving to be a simple yet effective approach. However, the critical instance
is only employed for implicit instance re-weighting and the method is
sensitive to both the choice of the instance feature encoder (i.e., pre-
trained ImageNet or self-supervised), and the scale of patches used. In MIL-
RNN[6], recurrent neural networks (RNN) are used to sequentially process
instance features, partially encoding position and context, but is limited in
the ability to capture long range dependences.

Thus, follow-up works[24, 21, 26, 23] built on the latter with more complex
attention-based variants using Transformer[30, 12] inspired architectures to
better model long range instance correlations via multi-head self-attention
(MSA) with positional information encoding. Along this line of thought,
TransMIL[26] highlights the importance of spatial positional encoding (PE) and
single-scale learning over the latter, but is relatively sensitive to the
depth of PE layers (i.e., x3) and does not explicitly pool all instances to a
single bag representation, instead uses a learnable class token for final bag-
level prediction. Thus, the use of Transformers with several MSA blocks can be
computationally prohibitive, and would be more desirable to have less over-
parameterized designs.

To address these challenges, we propose a Feature Re-calibration based MIL
framework (FRMIL), building upon prior MIL approaches[26, 21, 23] leveraging
MSA with Transformer encoders. Here, we argue that re-calibrating the
distribution of instance features can improve model performance towards better
generalization by using the properties of the data distribution directly. In
vision tasks such as few-shot learning, feature/distribution re-calibration is
used to enable better generalization when learning from limited samples by
transferring statistics from classes with sufficient examples[32]. However, in
the MIL scenario, instances are not always i.i.d [26], especially since
positive instances are often limited in a WSI i.e., ($\leq 10\%$). Thus, we
consider a simpler form that uses the max instance to shift the original
distribution towards better separability. Also, we consider MIL and anomaly
detection[7, 14, 20] as closely related tasks. For instance, Lee et al.[20]
leveraged MIL for weakly-supervised action localization by modeling
background/normal actions as out-of-distribution using uncertainty by
considering their inconsistency in a sequence with video-level labels only.

|  |
---|---|---
$44.96\%$ |  $83.72\%$ ($+39$) |  $89.10\%$ ($+44$)
(a) | (b) | (c)
Figure 1: Normalized density plots of the mean feature magnitudes on the
CAMELYON16[4] train-set, with test-set accuracy and improvements (red color).
(a) Original feature magnitudes. (b) Max-instance calibrated based features.
(c) Features learned by our FR-MIL model.

Inspired by these works, we hypothesize that features from positive and
negative bags (binary MIL) exhibit larger and smaller feature magnitudes
respectively, and this prior can be directly encoded into the learning
framework for better representation learning[20]. In Fig. 1, we show this
phenomena and our intuition to highlight how the standard MIL assumption of
having at-least one (+) instance in a bag can be used to make the distribution
more separable. Herein, we establish a simple non-parametric baseline that re-
calibrates features by subtracting the max instance per-bag, and then computes
the probability of a bag-label as the normalized minimum between the mean
magnitude and the estimated bag magnitude (see Sec. 2). Our evaluation shows
that the baseline performance is comparable to classic MIL operators (i.e.,
max/mean-pooling)[31].

To incorporate this idea in our framework, we explicitly re-calibrate features
with the aforementioned concept, and then feed the new features to a
positional encoding module (PEM)[26] followed by a single pooling multi-head
self-attention block (PMSA)[19] for bag classification. To effectively enforce
feature magnitude discrepancy, we propose a feature embedding loss that
maximizes the distance between positive and negative bag features, as well as
the standard cross-entropy losses. The main contributions of this work are as
follows: (i) We show that feature re-calibration using the max-critical
instance embedding is a simple yet powerful technique for MIL, (ii) We
introduce a feature magnitude loss to learn better instance/bag separation,
(iii) To obtain robust bag embeddings, we leverage a positional encoder and a
single self-attention block for instance aggregation, and (iv) Experimental
results on a public benchmark and inhouse-curated datasets demonstrate the
effectiveness of our method over state-of-the-art methods.

##  2 Methods

Overview. In this work, we consider a set of WSIs $\mathbf{X}=\\{X_{i}\\}$,
each associated with a slide-level label $Y_{i}=\\{0,1\\}$, and our goal
predict the slide labels using MIL (see. Fig. 2). We first extract instance
features $\mathbf{H}\in\mathbb{R}^{D}$ using a neural network
$\bf{F}_{\theta}$ i.e., $\bf{H}_{i}=\bf{F}_{\theta}(X_{i})$, where $\bf{F}$ is
either pre-trained on ImageNet or self-supervised learning[9, 15]. In FR-MIL,
we feed $\mathbf{H}$ to our max-instance selection module to obtain the
highest instance (critical) as well as it’s probability, then we re-calibrate
the features $\mathbf{H}$ with the max-instance to obtain $\mathbf{\hat{H}}$.
The position encoding module (PEM) creates a spatial representation of
$\mathbf{\hat{H}}$, applies a single group convolution $\mathbf{G}_{\theta}$
to obtain correlated features, and then concatenates with a learnable class
token $\mathbf{C}$. Finally, we perform MIL Pooling by Multi-head Self-
Attention (PSMA) using the max-instance as a query and the output of
$\mathbf{G}_{\theta}$ as key-value pairs to obtain the bag feature. FR-MIL is
trained to minimize the bag loss, max-instance loss, and feature magnitude
loss between positive and negative instance features. We detail each step
below.

Figure 2: Overview of the proposed FR-MIL framework.

Preliminaries: A Simple Baseline. Inspired by the work of Lee et al.[20] that
employ feature magnitudes for uncertainity estimation of background actions in
video sequences, we hypothesize that normal and positive bags should have
different magnitudes and can serve as a simple MIL baseline. Herein, given
instance features $\mathbf{H}^{c}_{i}=\\{h_{1},h_{2},\dots,h_{n}\\}$, where
$c$ denotes the WSI class; the mean feature magnitude per WSI can be obtained
as $\mu^{c}_{i}=\frac{1}{N}\sum||\mathbf{H}^{c}_{i}||^{2}_{2}$, where $N$
denotes the number of instances in a bag. To obtain the probability
$\mathbf{P}(y=1.)$ of a bag, we formalize our assumption as:

$\mathbf{P}(y=1.|\mu^{c}_{i})=\frac{\text{min}(\tau,\mu^{c}_{i})}{\tau},$ (1)

where $\tau$ is the pre-defined maximum feature magnitude determined on the
train-set only i.e., point at which the distributions first meet (see Fig.1).
Also, Eq. 1 is ensured to fall between 0 and 1, i.e.,
$0\leq\mathbf{P}(y=1.|\cdot)\leq 1$. In Fig. 1, we show the magnitudes on a
given train-set (Camelyon16[4]) as a density plot. Note that while both normal
and tumor slide curves appear to follow the Gaussian distribution (Fig. 1(a)),
separation is non-trivial due to the presence of more normal than tumor
instances ($\leq 10\%$), and reports a low accuracy ($44\%$) with $\tau=18.8$.
In Fig. 1(b), we show that re-calibrating the distribution by subtracting the
max-instance feature before computing the magnitudes creates better
separability. Formally,
$\hat{\mu}^{c}_{i}=\frac{1}{N}\sum||\mathbf{\hat{H}}^{c}_{i}||^{2}_{2}$, where
$\mathbf{\hat{H}}^{c}_{i}=\\{\mathbf{\hat{H}}-h^{c}_{\text{max}}\\}$ given
$h^{c}_{\text{max}}=\text{argmax}_{c}~{}\mathbf{\hat{H}}^{c}_{i}$. Notably,
re-calibration improves the test-accuracy by $+39$ with $\tau=8.2$. Finally,
Fig. 1(c) shows the learned distribution of FR-MIL when trained with a feature
magnitude loss $\mathcal{L}_{fm}$ and re-calibration, with more significant
improvements ($+44$), further validating our hypothesis.

Feature Re-calibration & Max-Instance Selection. Given the set of instance
features $\mathbf{H}$, our goal is to select the max-instance $h^{q}$ and it’s
associated score $A^{c}$ using an instance classifier
$\mathbf{f}^{m}_{\theta}$ in our Max-Pooling module (see. Fig. 2). Here,
$A^{c}=\rho(\mathbf{f}^{m}_{\theta}(\mathbf{H}))$ where $\rho$ denotes the
sigmoid function. Consequently, the sorted scores are used to index the max-
instance in $\mathbf{H}$ via an operator $g(\cdot)$, with $h^{q}$ later
employed for feature re-calibration, as well as instance feature aggregation
via PSMA for bag prediction. The max score $A^{c}$ is used to train the
instance classifier $\mathbf{f}^{m}_{\theta}$ using the loss
$\mathcal{L}_{max}$, in parallel with other modules in FR-MIL. Formally, re-
calibration of features can be modeled as

$\mathbf{\hat{H}}=\text{ReLU}(\mathbf{\hat{H}}-h^{q}),$ (2)

similar to the intuition highlighted by the simple baseline. To further
incorporate the concept of distribution re-calibration in our framework, we
draw connections to prior work[20] for anomaly detection i.e., assumes the
feature magnitudes of positive- and normal-bags are different, and can be
modeled via uncertainty. Therefore, to effectively model the ambiguous
normal/background features, the training procedure should employ both positive
and negative bags simultaneously instead of selecting normal features within a
single bag (single batch). Herein, counter to existing methods that use a
single bag for training, we employ a sampling strategy to produce balanced
bags per epoch i.e., we initialize a zero-tensor with the maximum bag size in
during training, and fill the relevant bag instance features. Note that by
‘balanced’, we imply 1-negative and 1-positive bag is sampled. Formally, to
enforce feature discrepancy we propose feature magnitude loss
$\mathcal{L}_{fm}$ as:

$\mathcal{L}_{fm}(\mathbf{\hat{H}}^{pos}_{i},\mathbf{\hat{H}}^{neg}_{i},\tau)=\frac{1}{N}\sum^{N}_{n=1}(\text{max}(0,\tau-||\mathbf{\hat{H}}^{pos}_{i}||)+||\mathbf{\hat{H}}^{neg}_{i}||),$
(3)

where $\mathbf{\hat{H}}^{pos}$, and $\mathbf{\hat{H}}^{neg}$ are the positive-
and negative bag instance features, and $\tau$ is the pre-defined margin,
respectively. While prior work[21] equally used max-pooling to select the max-
instance, note that non-local masked attention was proposed for bag feature
learning, whereas we use PSMA and propose feature re-calibration.

Positional Encoding Module (PEM). In the standard transformer[30, 12] design,
encoding spatial information has proved useful for recognition tasks. However,
it is non-trivial for WSIs due to varying sizes. In this work, we employ a
conditional position encoder (PEM)[23] that takes re-calibrated features
$\mathbf{\hat{H}}$, performs zero-padding to provide absolute position for a
convolution $\mathbf{G}_{\theta}$, and later concatenates the output with a
class token $\mathbf{C}$ initialized from the normal distribution. Here,
features $\mathbf{\hat{H}}$ are re-shaped into a 2D image by first computing
$\\{H,W\\}$ i.e., $H=\sqrt{N}=\sqrt{n}$, where $n$ is the number of instances
in a bag. Thus,

$\mathbf{\hat{H}}\in\mathbb{R}^{D}\rightarrow\mathbf{\hat{H}}\in\mathbb{R}^{B\times
C\times H\times W},$ (4)

where $B$ is the batch-size, $C=D$ are the instance feature dimensions, and
$\mathbf{G}_{\theta}$ is 2D convolution layer that performs group convolution
with kernel size $3\times 3$, and $1\times 1$ zero padding. Note that prior
work[23] used different sized convolutions in a pyramidal fashion. Instead, we
opt for a single layer to maintain computational feasibility. Finally, let
$\mathbf{\acute{H}}=\text{concat}(\mathbf{C}_{\theta},\mathbf{G}_{\theta}(\mathbf{\hat{H}}))$,
where $\mathbf{\acute{H}}\in\mathbb{R}^{(N+1)\times D}$ are the flattened
restored features i.e., in the case of a single bag.

MIL Pooling by Multi-head Self-Attention (PMSA). In order to pool instance
features $\mathbf{\acute{H}}$ to a single bag feature, we employ a single
multi-head Transformer encoder[19] that takes as input the max-instance
feature $h^{q}$ as a query and $\mathbf{\acute{H}}$ as key-value pairs i.e.,
$\text{PMSA}_{\theta}(h^{q},\mathbf{\acute{H}})$. The formulation proposed by
Lee et al.[19] for set-based tasks employs an attention function
$\varphi(\mathbf{Q},\mathbf{K},\mathbf{V})$ to measure similarity between a
query vector $\mathbf{Q}$ with key-value pairs
$\mathbf{K,V}\in\mathbb{R}^{d\times m}$ as:
$\varphi(\mathbf{Q},\mathbf{K},\mathbf{V})=\text{softmax}(\frac{\mathbf{Q}\mathbf{K}^{T}}{\sqrt{m}})\mathbf{V}$,
where $\\{d,m\\}$ is the instance feature dimension. This can be easily
extended to multi-head attention by first projecting vectors onto $k$
different dimensions. The encoder consists of feed-forward networks
$\\{\mathbf{f}^{q}_{\theta},\mathbf{f}^{k}_{\theta},\mathbf{f}^{v}_{\theta}\\}$,
where $\mathbf{f}^{o}_{\theta}$ is fed the output of $\varphi$ (prototype)
together with residual connections and optional Layer Normalization[2] (LN).
Formally, let
$\hat{\varphi}=\varphi(\mathbf{Q},\mathbf{K},\mathbf{V})+\mathbf{Q}$, then:

$z=\text{PSMA}(h^{q},\mathbf{\acute{H}},\mathbf{\acute{H}})=\text{LN}(\hat{\varphi}+\text{ReLU}(\mathbf{f}^{o}_{\theta}(\hat{\varphi}))),$
(5)

to produce a bag feature $z$, later fed to the bag classifier
$\mathbf{f}^{c}_{\theta}$ for WSI classification. Finally, FR-MIL is trained
to minimize the bag-, max-pooling and feature losses. Thus, the final
objective is:

$\mathcal{L}=\gamma_{1}\mathcal{L}_{bag}(\hat{y},y)+\gamma_{2}\mathcal{L}_{max}(A^{c},y)+\gamma_{3}L_{fm}(\mathbf{\hat{H}}^{pos}_{i},\mathbf{\hat{H}}^{neg}_{i},\tau),$
(6)

where $\\{\gamma_{i}\\}$ are balancing weights and
$\mathcal{L}_{\\{bag,max\\}}$ is the binary cross-entropy loss over the true
WSI labels $y$ given $\hat{y}=\mathbf{f}^{c}_{\theta}(z)$, respectively.

##  3 Experiments

Datasets. To demonstrate the effectiveness of FR-MIL, we conducted experiments
on the publicly available dataset CAMELYON16[4], and an in-house curated
dataset termed COLON-MSI i.e., colorectal (adenocarcinoma) cancer slides
involving microsatellite instable (MSI) molecular phenotypes[5]. CAMELYON16
dataset was proposed for metastatis detection in breast cancer, it consists of
271 training sets and 129 testing sets. After pre-processing, a total of 3.2
million patches at $\times$20 magnification, with an average of 8,800 patches
per bag, and a maximum of 30,000 patches per bag on the training set. On the
other hand, COLON-MSI consists of both microsatellite-stable (MSS) and
microsatellite-instable (MSI), and is thus a subtyping task. It consists of
625 images, split as follows: 360 training, 92 validation, and 173 testing
sets. Experts pathologists detected the presence of tumors with
Immunohistochemical analysis (IHC) and PCR-based amplification and
collectively agreed on the final slide-level label. Note that tumor ROIs are
not used in this work. After pre-proessing, a total of 3.5 million patches at
$\times$20 magnification, an average of 6,000 patches/bag, and a maximum of
8900 patches in the train-set.

Implementation Settings. In the pre-processing step, we extracted valid
patches of $256\times 256$ after tissue detection and discard patches with
$\leq 15\%$ tissue entropy. For the instance encoder $\mathbf{F}_{\theta}$, we
employed the SimCLR[9] ResNet18[16] encoder trained by Lee et al.[21] for the
CAMELYON16 dataset. On the COLON-MSI set, we used an ImageNet pre-trained
ResNet18. Thus, each instance feature is represented as
$\mathbf{H}_{i}\in\mathbb{R}^{n\times 512}$. FR-MIL is trained with balanced
batch sampling (B = 2), and learning rate of $1e-4$ with Adam optimizer for
100 epochs with $20\%$ dropout as regularization, and PSMA has heads $k=8$.
Hyper-parameters $\\{\gamma_{1,2,3}\\}=0.33$, with $\tau=8.48$ for CAMELYON16,
and $\tau=57.5$ on COLON-MSI, respectively.

Comparison Methods. We compare FR-MIL to traditional MIL methods max- and
mean-pooling[31], as well as existing state-of-the-art methods: ABMIL[18],
DSMIL[21], CLAM-SB[24], MIL-RNN[6], and TransMIL[26]. All compared methods are
trained for 200 epochs on COLON-MSI with similar settings.

Table 1:  Evaluation of the proposed method on CAMELYON16 (CM16) and COLON-MSI sets. Metrics accuracy (ACC) and area under the curve (AUC) were employed. $\dagger$ :_denotes scores reported in the paper using ResNet50 as $\mathbf{F}_{\theta}$_ with ImageNet features.  | CM16 |  | COLON-MSI
---|---|---|---
Method |  | ACC |  | AUC |  |  | ACC |  | AUC
Mean-pooling[31] |  | 0.7984 |  | 0.7620 |  |  | 0.624 |  | 0.830
Max-pooling[31] |  | 0.8295 |  | 0.8641 |  |  | 0.763 |  | 0.859
ABMIL[18] |  | 0.8450 |  | 0.8653 |  |  | 0.740 |  | 0.779
MIL-RNN[6] |  | 0.8062 |  | 0.8064 |  |  | 0.630 |  | 0.631
CLAM-SB[24] |  | 0.845 |  | 0.894 |  |  | 0.786 |  | 0.820
DSMIL[21] |  | 0.8682 |  | 0.8944 |  |  | 0.734 |  | 0.811
TransMIL[23] |  | 0.791 |  | 0.813 |  |  | 0.676 |  | 0.617
TransMIL$\dagger$[23] |  | 0.8837 |  | 0.9309 |  |  | - |  | -
FR-MIL (w/ $\mathcal{L}_{bag}$) |  | 0.8600 |  | 0.8990 |  |  | 0.809 |  | 0.880
FR-MIL (w/ $\mathcal{L}_{bag}+\mathcal{L}_{fm}$) |  | 0.8760 |  | 0.8990 |  |  | 0.775 |  | 0.842
FR-MIL (w/ $\mathcal{L}_{bag}+\mathcal{L}_{max}$) |  | 0.8840 |  | 0.8940 |  |  | 0.780 |  | 0.831
FR-MIL (w/ $\mathcal{L}_{bag}+\mathcal{L}_{max}+\mathcal{L}_{fm}$) |  | 0.8910 |  | 0.8950 |  |  | 0.809 |  | 0.901

##  4 Results and Discussion

Main Results. Table 1 presents the results of our approach against recent
methods. On the CAMELYON16 dataset, FR-MIL reports $+3\%$ (ACC) improvement
over DSMIL with comparable performance ($+1\%$) to the reported TransMIL
scores using a larger feature extractor. Given that only a small portion of a
each positive bag contains tumors, using the max-instance for bag pooling with
re-calibration is intuitively sound and shows better performance over other
methods. Moreover, though we employ PEM similar to TransMIL, the use of a
single PEM module on calibrated facilitates better correlation learning.

On the other hand, since the majority of slides contain relatively large tumor
regions (averagely $\geq 75\%$) in COLON-MSI, max- and mean-pooling show high
AUC but had inconsistent results (ACC). Overall, CLAM-SB reports the best ACC
among the compared methods i.e., $78.6\%$. Interestingly, TransMIL performed
poorly on this set, possibly due to over-parametrization, and the
morphological similarities between MSS and MSI instances. Similar observations
were drawn regarding MIL-RNN. Consequently, the proposed FR-MIL highlights the
importance of calibration in subtyping problems, reporting $+2\%$ (ACC) and
$+5\%$ (AUC) improvements achieving the best scores. See Fig. 3 for the
distribution of features.

Ablations. To validate the effectiveness of the proposed losses on learning,
we evaluated FR-MIL with/without certain losses (see. Table 1). First, on
CAMELYON16, we found $\mathcal{L}_{fm}$ was a crucial component to further
boost the AUC score. Overall, when both $\mathcal{L}_{max}$ and
$\mathcal{L}_{fm}$ were omitted, performance drops were noted. On the other
hand, on COLON-MSI, using a $\mathcal{L}_{bag}$ only had the best scores,
whereas the use of both $\mathcal{L}_{max}$ and $\mathcal{L}_{fm}$ resulted in
a significant reduction (AUC). However, employing all the losses resulted in
more stable scores.

|  |
---|---|---
(a) | (b) | (c)
Figure 3: Density plots of the mean feature magnitudes on the COLON-MSI train-
set. (a) Original feature magnitudes. (b) Max-instance re-calibration based
features. (c) Features learned by our FR-MIL model.

##  5 Conclusion

In this work, we presented a MIL framework for Whole Slide Image
classification that leverages Feature Re-calibration, applicable to both
binary and sub-typing tasks. We show that: (i) by leveraging feature magnitude
discrepancy between positive and negative bags as a probabilistic measure; a
simple baseline is comparable in performance to classic MIL operators, (ii)
explicitly re-calibrating the data distribution with max-instances during
training by drawing connections to the standard MIL assumption is simple yet
effective, and (iii) the use of a metric feature loss to encourage better
feature separation in (+/-) bags improves both Accuracy and AUC over state-of-
the-art methods. Further exploring the utility of this approach in multi-scale
setting, or designing an adaptable margin (mean magnitude) estimator will be
topics of future research.

Acknowledgments This work was supported by the DGIST R&D program of the
Ministry of Science and ICT of KOREA (21-DPIC-08), Smart HealthCare Program
funded by the Korean National Police Agency (220222M01), and IITP grant funded
by the Korean government (MSIT) (No.2021-0-02068, Artificial Intelligence
Innovation Hub).

## References

  * [1] Amores, J.: Multiple instance classification: Review, taxonomy and comparative study. Artificial intelligence 201, 81–105 (2013)
  * [2] Ba, J.L., Kiros, J.R., Hinton, G.E.: Layer normalization. arXiv preprint arXiv:1607.06450 (2016)
  * [3] Banerji, S., Mitra, S.: Deep learning in histopathology: A review. Wiley Interdisciplinary Reviews: Data Mining and Knowledge Discovery 12(1), e1439 (2022)
  * [4] Bejnordi, B.E., Veta, M., Van Diest, P.J., Van Ginneken, B., Karssemeijer, N., Litjens, G., Van Der Laak, J.A., Hermsen, M., Manson, Q.F., Balkenhol, M., et al.: Diagnostic assessment of deep learning algorithms for detection of lymph node metastases in women with breast cancer. Jama 318(22), 2199–2210 (2017)
  * [5] Boland, C.R., Goel, A.: Microsatellite instability in colorectal cancer. Gastroenterology 138(6), 2073–2087 (2010)
  * [6] Campanella, G., Hanna, M.G., Geneslaw, L., Miraflor, A., Werneck Krauss Silva, V., Busam, K.J., Brogi, E., Reuter, V.E., Klimstra, D.S., Fuchs, T.J.: Clinical-grade computational pathology using weakly supervised deep learning on whole slide images. Nature medicine 25(8), 1301–1309 (2019)
  * [7] Chalapathy, R., Chawla, S.: Deep learning for anomaly detection: A survey. arXiv preprint arXiv:1901.03407 (2019)
  * [8] Chen, H., Wang, K., Zhu, Y., Yan, J., Ji, Y., Li, J., Xie, D., Huang, J., Cheng, S., Yao, J.: From pixel to whole slide: Automatic detection of microvascular invasion in hepatocellular carcinoma on histopathological image via cascaded networks. In: MICCAI. pp. 196–205. Springer (2021)
  * [9] Chen, T., Kornblith, S., Norouzi, M., Hinton, G.: A simple framework for contrastive learning of visual representations. In: ICML. pp. 1597–1607. PMLR (2020)
  * [10] Chikontwe, P., Kim, M., Nam, S.J., Go, H., Park, S.H.: Multiple instance learning with center embeddings for histopathology classification. In: MICCAI. pp. 519–528. Springer (2020)
  * [11] Dimitriou, N., Arandjelović, O., Caie, P.D.: Deep learning for whole slide image analysis: an overview. Frontiers in medicine p. 264 (2019)
  * [12] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., Houlsby, N.: An image is worth 16x16 words: Transformers for image recognition at scale. In: ICLR (2021)
  * [13] Fan, L., Sowmya, A., Meijering, E., Song, Y.: Learning visual features by colorization for slide-consistent survival prediction from whole slide images. In: MICCAI. pp. 592–601. Springer (2021)
  * [14] Feng, J.C., Hong, F.T., Zheng, W.S.: Mist: Multiple instance self-training framework for video anomaly detection. In: CVPR. pp. 14009–14018 (2021)
  * [15] Grill, J.B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, Z., Gheshlaghi Azar, M., et al.: Bootstrap your own latent-a new approach to self-supervised learning. NeurIPS 33, 21271–21284 (2020)
  * [16] He, K., Zhang, X., Ren, S., Sun, J.: Deep residual learning for image recognition. In: CVPR. pp. 770–778 (2016)
  * [17] He, L., Long, L.R., Antani, S., Thoma, G.R.: Histology image analysis for carcinoma detection and grading. Computer methods and programs in biomedicine 107(3), 538–556 (2012)
  * [18] Ilse, M., Tomczak, J., Welling, M.: Attention-based deep multiple instance learning. In: ICML. pp. 2127–2136. PMLR (2018)
  * [19] Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., Teh, Y.W.: Set transformer: A framework for attention-based permutation-invariant neural networks. In: ICML. pp. 3744–3753. PMLR (2019)
  * [20] Lee, P., Wang, J., Lu, Y., Byun, H.: Weakly-supervised temporal action localization by uncertainty modeling. In: AAAI. vol. 2 (2021)
  * [21] Li, B., Li, Y., Eliceiri, K.W.: Dual-stream multiple instance learning network for whole slide image classification with self-supervised contrastive learning. In: CVPR. pp. 14318–14328 (2021)
  * [22] Li, C., Li, X., Rahaman, M., Li, X., Sun, H., Zhang, H., Zhang, Y., Li, X., Wu, J., Yao, Y., et al.: A comprehensive review of computer-aided whole-slide image analysis: from datasets to feature extraction, segmentation, classification, and detection approaches. arXiv preprint arXiv:2102.10553 (2021)
  * [23] Li, H., Yang, F., Zhao, Y., Xing, X., Zhang, J., Gao, M., Huang, J., Wang, L., Yao, J.: Dt-mil: Deformable transformer for multi-instance learning on histopathological image. In: MICCAI. pp. 206–216. Springer (2021)
  * [24] Lu, M.Y., Williamson, D.F., Chen, T.Y., Chen, R.J., Barbieri, M., Mahmood, F.: Data-efficient and weakly supervised computational pathology on whole-slide images. Nature biomedical engineering 5(6), 555–570 (2021)
  * [25] Rymarczyk, D., Borowa, A., Tabor, J., Zielinski, B.: Kernel self-attention for weakly-supervised image classification using deep multiple instance learning. In: IEEE Winter. Conf. Application. Comput. Vis. pp. 1721–1730 (2021)
  * [26] Shao, Z., Bian, H., Chen, Y., Wang, Y., Zhang, J., Ji, X., et al.: Transmil: Transformer based correlated multiple instance learning for whole slide image classification. NeurIPS 34 (2021)
  * [27] Sharma, Y., Shrivastava, A., Ehsan, L., Moskaluk, C.A., Syed, S., Brown, D.: Cluster-to-conquer: A framework for end-to-end multi-instance learning for whole slide image classification. In: Medical Imaging with Deep Learning. pp. 682–698. PMLR (2021)
  * [28] Shi, X., Xing, F., Xie, Y., Zhang, Z., Cui, L., Yang, L.: Loss-based attention for deep multiple instance learning. In: AAAI. vol. 34, pp. 5742–5749 (2020)
  * [29] Srinidhi, C.L., Ciga, O., Martel, A.L.: Deep neural network models for computational histopathology: A survey. Medical Image Analysis 67, 101813 (2021)
  * [30] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł., Polosukhin, I.: Attention is all you need. NeurIPS 30 (2017)
  * [31] Wang, X., Yan, Y., Tang, P., Bai, X., Liu, W.: Revisiting multiple instance neural networks. Pattern Recognition 74, 15–24 (2018)
  * [32] Yang, S., Liu, L., Xu, M.: Free lunch for few-shot learning: Distribution calibration. In: ICLR (2020)