diff --git a/av1/encoder/partition_model_weights.h b/av1/encoder/partition_model_weights.h
index ac2890c..5baf852 100644
--- a/av1/encoder/partition_model_weights.h
+++ b/av1/encoder/partition_model_weights.h
@@ -2442,266 +2442,6 @@
 #undef LABEL_SIZE
 
 // Below are the models used for simple_motion_search_based_split
-// The model for the faster model
-static const float av1_simple_motion_search_based_split_thresh_128 = 2.0f;
-static const float av1_simple_motion_search_based_split_thresh_64 = 2.0f;
-static const float av1_simple_motion_search_based_split_thresh_32 = 2.0f;
-static const float av1_simple_motion_search_based_split_thresh_16 = 2.0f;
-static const float av1_simple_motion_search_based_split_thresh_8 = 2.0f;
-
-// BLOCK_128X128
-#define NUM_HIDDEN_LAYERS_128 1
-#define NUM_FEATURES_128 6
-#define NUM_LAYER_0_UNITS_128 16
-#define NUM_LOGITS_128 1
-
-static const float av1_simple_motion_search_based_split_layer_0_kernel_128[] = {
-  -0.807346f,  0.242298f,   12.9862f,   -1.19161f,  5.21734f,    -1.1363f,
-  -2.39127f,   0.930915f,   -2.44285f,  -2.42966f,  5.73476f,    0.0506879f,
-  -0.234878f,  -0.317875f,  0.361322f,  0.431648f,  -0.39105f,   -0.110225f,
-  -2.46236f,   0.979713f,   -10.5596f,  -7.76653f,  -3.06518f,   2.42554f,
-  0.0492961f,  -0.467176f,  0.130746f,  0.494527f,  -0.0336645f, 0.501755f,
-  0.176486f,   -0.869541f,  7.77757f,   6.81303f,   6.00771f,    7.35696f,
-  0.150731f,   -0.307017f,  -0.437639f, -0.082924f, 0.379107f,   0.452278f,
-  -0.0143836f, -0.183691f,  -0.604698f, -9.2681f,   -2.06087f,   11.0256f,
-  0.0487599f,  -0.249168f,  -0.180407f, 0.304772f,  0.218642f,   -0.406073f,
-  -0.0289919f, -0.794381f,  5.45092f,   5.38374f,   3.25745f,    5.32903f,
-  1.12718f,    -0.0215478f, 2.78552f,   4.8951f,    -0.959671f,  0.694264f,
-  -0.0611219f, -0.331937f,  0.258252f,  -0.495331f, -0.285923f,  0.294713f,
-  -0.119947f,  0.0753204f,  10.2021f,   -5.82147f,  -12.0137f,   3.0365f,
-  0.366697f,   0.142683f,   -3.29731f,  -5.76651f,  -5.62578f,   10.9462f,
-  -0.325459f,  0.092602f,   -0.868027f, -0.691768f, -0.292017f,  -0.00841203f,
-  0.702545f,   -0.612227f,  -7.68881f,  9.52225f,   -1.18581f,   -2.56762f
-};
-
-static const float av1_simple_motion_search_based_split_logits_kernel_128[] = {
-  0.364895f,    0.577553f,  0.115758f,  -0.999496f, 0.124885f, 3.23193f,
-  -0.00386642f, 0.970794f,  0.136637f,  -4.28052f,  -1.49234f, 0.370436f,
-  0.576981f,    -0.469656f, -0.124071f, 1.07669f
-};
-
-static const float av1_simple_motion_search_based_split_layer_0_bias_128[] = {
-  1.32916f,    0.817212f,  0.0f,       -0.921066f, 0.0f,      3.57649f,
-  -0.0204517f, 2.97286f,   0.0f,       5.49957f,   -8.14518f, 0.0f,
-  1.30826f,    -0.349536f, -0.638933f, 5.4496f
-};
-
-static const float av1_simple_motion_search_based_split_logits_bias_128[] = {
-  0.683442f
-};
-
-static const NN_CONFIG av1_simple_motion_search_based_split_nn_config_128 = {
-  NUM_FEATURES_128,
-  NUM_LOGITS_128,
-  NUM_HIDDEN_LAYERS_128,
-  {
-      NUM_LAYER_0_UNITS_128,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_kernel_128,
-      av1_simple_motion_search_based_split_logits_kernel_128,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_bias_128,
-      av1_simple_motion_search_based_split_logits_bias_128,
-  },
-};
-
-#undef NUM_HIDDEN_LAYERS_128
-#undef NUM_FEATURES_128
-#undef NUM_LAYER_0_UNITS_128
-#undef NUM_LOGITS_128
-
-// BLOCK_64X64
-#define NUM_HIDDEN_LAYERS_64 1
-#define NUM_FEATURES_64 6
-#define NUM_LAYER_0_UNITS_64 16
-#define NUM_LOGITS_64 1
-
-static const float av1_simple_motion_search_based_split_layer_0_kernel_64[] = {
-  0.0345945f,  -0.394064f,  0.0919978f, 0.270358f,  -0.384502f, -0.504608f,
-  -0.25759f,   0.155981f,   2.62567f,   -10.7204f,  -0.709802f, 8.15948f,
-  0.589866f,   -0.445645f,  -1.68232f,  10.0061f,   -3.17671f,  4.87259f,
-  -0.448886f,  -0.205568f,  -0.462388f, 0.385001f,  -0.451687f, 0.49602f,
-  -0.256708f,  0.803322f,   3.25594f,   0.38541f,   -1.83867f,  -2.15132f,
-  0.936059f,   -0.203056f,  -5.92959f,  -6.24554f,  -6.68631f,  -6.85977f,
-  -0.0407565f, -0.258902f,  0.195053f,  -0.366515f, 0.339543f,  -0.433017f,
-  -2.67026f,   0.385457f,   1.86683f,   1.9501f,    0.0381398f, 1.086f,
-  -0.153729f,  0.173772f,   -42.9029f,  -36.8934f,  -2.892f,    -0.0540691f,
-  0.77469f,    -0.380145f,  2.2689f,    -9.53332f,  1.15712f,   2.86601f,
-  -0.437036f,  0.247132f,   -8.51058f,  -3.62972f,  -8.99449f,  -0.638738f,
-  0.0609263f,  -0.0614603f, 5.42307f,   5.35926f,   5.27437f,   5.26599f,
-  -0.0729677f, 0.0306104f,  -7.77867f,  5.03598f,   -8.17832f,  5.85461f,
-  -0.253269f,  0.164582f,   -4.49713f,  3.83265f,   9.04851f,   -2.85668f,
-  1.22618f,    0.166904f,   -1.51975f,  -4.01576f,  -1.44374f,  -2.22147f,
-  -0.217072f,  -0.0984913f, -0.265515f, 0.360021f,  0.0779512f, 0.361516f
-};
-
-static const float av1_simple_motion_search_based_split_logits_kernel_64[] = {
-  0.470821f, 0.474747f, -0.571292f, 0.403221f,  0.628966f,  -0.617029f,
-  0.501105f, 0.499962f, -1.5451f,   -0.473518f, -0.730568f, -5.55817f,
-  0.776761f, 0.42569f,  0.311925f,  0.469968f
-};
-
-static const float av1_simple_motion_search_based_split_layer_0_bias_64[] = {
-  -0.134085f, 0.0758715f, 1.10419f,  0.0f,       -5.75737f, 1.65494f,
-  0.0f,       3.44047f,   0.394852f, 3.43858f,   3.65871f,  -4.84987f,
-  1.21207f,   -1.7705f,   -5.46469f, -0.0889634f
-};
-
-static const float av1_simple_motion_search_based_split_logits_bias_64[] = {
-  -0.479491f
-};
-
-static const NN_CONFIG av1_simple_motion_search_based_split_nn_config_64 = {
-  NUM_FEATURES_64,
-  NUM_LOGITS_64,
-  NUM_HIDDEN_LAYERS_64,
-  {
-      NUM_LAYER_0_UNITS_64,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_kernel_64,
-      av1_simple_motion_search_based_split_logits_kernel_64,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_bias_64,
-      av1_simple_motion_search_based_split_logits_bias_64,
-  },
-};
-
-#undef NUM_HIDDEN_LAYERS_64
-#undef NUM_FEATURES_64
-#undef NUM_LAYER_0_UNITS_64
-#undef NUM_LOGITS_64
-
-// BLOCK_32X32
-#define NUM_HIDDEN_LAYERS_32 1
-#define NUM_FEATURES_32 6
-#define NUM_LAYER_0_UNITS_32 16
-#define NUM_LOGITS_32 1
-
-static const float av1_simple_motion_search_based_split_layer_0_kernel_32[] = {
-  -1.61796f,   0.0585128f,  1.57904f,   1.52703f,   0.367779f, 0.220434f,
-  1.66652f,    -1.77782f,   6.41118f,   4.16976f,   4.97299f,  4.84111f,
-  -0.0956536f, -0.163284f,  -0.143662f, 0.129329f,  0.449659f, -0.528844f,
-  -1.00067f,   1.17203f,    -4.26777f,  -4.78521f,  8.45658f,  -3.49498f,
-  -1.78386f,   0.111488f,   4.176f,     6.31911f,   -10.5369f, 6.26983f,
-  -1.32233f,   1.22999f,    -4.1666f,   -10.0359f,  -4.14779f, -10.4695f,
-  1.83011f,    -0.333152f,  -9.87986f,  -8.11992f,  -8.2775f,  -7.79918f,
-  -0.101404f,  0.00401393f, 8.89046f,   -7.32186f,  -6.59597f, 9.66257f,
-  -1.1492f,    1.23067f,    -3.6341f,   6.59275f,   -3.2373f,  -3.42564f,
-  0.371736f,   -0.140902f,  -2.75715f,  5.92487f,   -7.9185f,  9.13743f,
-  -3.52698f,   -0.191044f,  5.96691f,   6.26327f,   4.36378f,  5.69354f,
-  -0.608845f,  -0.191236f,  -0.482191f, -0.180474f, -3.8838f,  -3.92934f,
-  -1.03191f,   0.994568f,   7.95516f,   -4.0035f,   -2.86266f, -4.96105f,
-  1.75022f,    0.125058f,   -1.52159f,  -3.59304f,  -2.82634f, -2.49556f,
-  -2.05557f,   -0.222577f,  3.7608f,    5.50475f,   2.7046f,   5.25952f,
-  -1.91327f,   -0.0356497f, 1.47611f,   1.27499f,   -1.76108f, -0.578954f
-};
-
-static const float av1_simple_motion_search_based_split_logits_kernel_32[] = {
-  -0.220382f, -0.693902f, 0.424827f, 0.379952f, -0.413791f, -0.326785f,
-  -0.455086f, 0.242402f,  0.307986f, 0.175746f, 0.498901f,  -0.628053f,
-  0.285447f,  0.230052f,  0.415151f, -0.842946f
-};
-
-static const float av1_simple_motion_search_based_split_layer_0_bias_32[] = {
-  -1.80751f, 6.40356f,   -0.0512058f, -4.59163f, -0.369933f, -0.195755f,
-  -0.16648f, -0.599755f, -5.35975f,   -1.21349f, 2.48414f,   1.07096f,
-  -3.66684f, -6.17761f,  4.2159f,     -1.05286f
-};
-
-static const float av1_simple_motion_search_based_split_logits_bias_32[] = {
-  -2.58676f
-};
-
-static const NN_CONFIG av1_simple_motion_search_based_split_nn_config_32 = {
-  NUM_FEATURES_32,
-  NUM_LOGITS_32,
-  NUM_HIDDEN_LAYERS_32,
-  {
-      NUM_LAYER_0_UNITS_32,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_kernel_32,
-      av1_simple_motion_search_based_split_logits_kernel_32,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_bias_32,
-      av1_simple_motion_search_based_split_logits_bias_32,
-  },
-};
-
-#undef NUM_HIDDEN_LAYERS_32
-#undef NUM_FEATURES_32
-#undef NUM_LAYER_0_UNITS_32
-#undef NUM_LOGITS_32
-
-// BLOCK_16X16
-#define NUM_HIDDEN_LAYERS_16 1
-#define NUM_FEATURES_16 6
-#define NUM_LAYER_0_UNITS_16 16
-#define NUM_LOGITS_16 1
-
-static const float av1_simple_motion_search_based_split_layer_0_kernel_16[] = {
-  -0.611497f,  -0.0422086f, -0.555957f,   -0.632451f, -0.144179f, -0.152722f,
-  -0.330265f,  -0.419866f,  0.287343f,    0.385295f,  -0.424486f, 0.424281f,
-  2.27442f,    -2.47933f,   5.24731f,     4.33827f,   4.73215f,   3.41909f,
-  1.16058f,    -0.364505f,  0.12207f,     -0.287749f, 0.0509783f, -0.0200119f,
-  1.52907f,    -1.1905f,    -2.56978f,    -3.00186f,  -3.56084f,  -3.89276f,
-  0.00365657f, 1.57125f,    -4.421f,      -2.48803f,  -2.51531f,  -4.28646f,
-  2.52248f,    -1.03377f,   -1.09607f,    -1.44633f,  -1.58736f,  -1.25927f,
-  -1.45841f,   -0.566619f,  -0.246166f,   -0.182289f, -0.238156f, 0.177991f,
-  0.0112509f,  -0.17677f,   -0.485877f,   0.0812852f, 0.104975f,  0.222793f,
-  -0.372858f,  -0.48624f,   -0.00870389f, -0.385019f, 0.405842f,  0.288523f,
-  0.167374f,   -0.204208f,  -8.74148f,    -8.59267f,  -8.42492f,  -8.3778f,
-  -5.57063f,   -0.406818f,  -0.873199f,   -0.896224f, -0.701479f, -0.985736f,
-  -0.625956f,  -0.0446202f, -0.509987f,   -0.321804f, -0.470759f, -0.248556f,
-  -0.369436f,  -0.160828f,  0.0591148f,   0.405218f,  0.142584f,  -0.130106f,
-  0.125321f,   0.0888179f,  7.34822f,     -6.71488f,  -7.06592f,  6.33224f,
-  0.0333619f,  -0.377782f,  0.160767f,    -0.128169f, -0.484818f, -0.311973f
-};
-
-static const float av1_simple_motion_search_based_split_logits_kernel_16[] = {
-  -0.132207f,   0.15176f,   -0.680086f, 0.605921f, -0.43294f,  0.485811f,
-  -0.306286f,   0.551368f,  0.413904f,  0.548748f, -0.437391f, 0.560778f,
-  -0.00685266f, -0.558657f, 0.122127f,  0.260165f
-};
-
-static const float av1_simple_motion_search_based_split_layer_0_bias_16[] = {
-  -0.200928f, -0.074132f, 8.69963f,    -9.00807f,  9.08983f, -6.83586f,
-  -3.89329f,  10.4881f,   -0.0670618f, 0.0f,       9.21614f, 8.41773f,
-  -0.145851f, 0.0f,       -1.43038f,   -0.0460311f
-};
-
-static const float av1_simple_motion_search_based_split_logits_bias_16[] = {
-  -4.19885f
-};
-
-static const NN_CONFIG av1_simple_motion_search_based_split_nn_config_16 = {
-  NUM_FEATURES_16,
-  NUM_LOGITS_16,
-  NUM_HIDDEN_LAYERS_16,
-  {
-      NUM_LAYER_0_UNITS_16,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_kernel_16,
-      av1_simple_motion_search_based_split_logits_kernel_16,
-  },
-  {
-      av1_simple_motion_search_based_split_layer_0_bias_16,
-      av1_simple_motion_search_based_split_logits_bias_16,
-  },
-};
-
-#undef NUM_HIDDEN_LAYERS_16
-#undef NUM_FEATURES_16
-#undef NUM_LAYER_0_UNITS_16
-#undef NUM_LOGITS_16
-
-// Slower higher quality split model
 // Threshold from BLOCK_128X128 to BLOCK_8X88
 static const float av1_simple_motion_search_split_midres_thresh[5] = {
   5.839480f,  // p = 0.997098
@@ -2713,7 +2453,7 @@
 
 static const float av1_simple_motion_search_split_midres_no_thresh[5] = {
   -3.38168078f,  // p = 0.032872917
-  -4.08610739f,  // p = 0.14393017
+  -4.08610739f,  // p = 0.016526795
   -1.7830237f,   // p = 0.15270848
   -100.0f,       // p = 0
   -100.0f,       // p = 0
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c
index 3d2a947..50aa770 100644
--- a/av1/encoder/partition_strategy.c
+++ b/av1/encoder/partition_strategy.c
@@ -102,9 +102,6 @@
     int mi_col, BLOCK_SIZE bsize, const int *const refs, int num_refs,
     int use_subpixel, int save_mv_code, unsigned int *best_sse,
     unsigned int *best_var) {
-  // TODO(chiyotsai@google.com): The calculation of variance currently uses
-  // bsize, so we might take area outside of the image into account. We need to
-  // modify the SIMD functions to fix this later.
   const AV1_COMMON *const cm = &cpi->common;
   int best_ref = -1;
 
@@ -160,19 +157,21 @@
   return best_ref;
 }
 
-// Performs fullpixel simple_motion_search with LAST_FRAME and ALTREF_FRAME on
-// each subblock and extract the variance and sse of residues. Then store the
-// var and sse from each partition subblock to features. The DC qindex is also
-// stored in features.
-// Here features is assumed to be a length 19 array.
-// After this function is called, we will store the following to features:
-// features[0:17] = var and sse from subblocks
-// features[18] = DC q_index
+// Collects features using simple_motion_search and store them in features. The
+// features are also cached in PC_TREE. By default, the features collected are
+// the sse and var from the subblocks flagged by features_to_get. Furthermore,
+// if features is not NULL, then 7 more features are appended to the end of
+// features:
+//  - log(1.0 + dc_q ** 2)
+//  - whether an above macroblock exists
+//  - width of above macroblock
+//  - height of above macroblock
+//  - whether a left marcoblock exists
+//  - width of left macroblock
+//  - height of left macroblock
 static void simple_motion_search_prune_part_features(
     AV1_COMP *const cpi, MACROBLOCK *x, PC_TREE *pc_tree, int mi_row,
     int mi_col, BLOCK_SIZE bsize, float *features, int features_to_get) {
-  // TODO(chiyotsai@google.com): Cache the result of the motion search from the
-  // larger bsize.
   const int w_mi = mi_size_wide[bsize];
   const int h_mi = mi_size_high[bsize];
   assert(mi_size_wide[bsize] == mi_size_high[bsize]);
diff --git a/av1/encoder/partition_strategy.h b/av1/encoder/partition_strategy.h
index fbe832d..1b002cd 100644
--- a/av1/encoder/partition_strategy.h
+++ b/av1/encoder/partition_strategy.h
@@ -75,7 +75,7 @@
                                                    int *do_square_split);
 
 // Get the features for selecting the max and min partition size. Currently this
-// performs simple_motion_search on 16X16 subblocks of the currnet superblock,
+// performs simple_motion_search on 16X16 subblocks of the current superblock,
 // and then extract the statistics of sse and motion vectors as features.
 void av1_get_max_min_partition_features(AV1_COMP *const cpi, MACROBLOCK *x,
                                         int mi_row, int mi_col,
diff --git a/av1/encoder/speed_features.c b/av1/encoder/speed_features.c
index faeb7af..82ffd1d 100644
--- a/av1/encoder/speed_features.c
+++ b/av1/encoder/speed_features.c
@@ -155,9 +155,6 @@
     } else if (is_480p_or_larger) {
       sf->use_square_partition_only_threshold = BLOCK_32X32;
     } else {
-      // TODO(chiyotsai@google.com): Setting the threshold to BLOCK_16X16 incurs
-      // a large loss (about 0.584%). Try increasing the threshold on boosted
-      // frame and see if it improves the performance.
       sf->use_square_partition_only_threshold = BLOCK_32X32;
     }
 
diff --git a/av1/encoder/speed_features.h b/av1/encoder/speed_features.h
index ef86923..1eab76d 100644
--- a/av1/encoder/speed_features.h
+++ b/av1/encoder/speed_features.h
@@ -703,13 +703,13 @@
   int inter_mode_rd_model_estimation_adaptive;
 
   // Use very reduced set of inter mode checks and fast non-rd mode cost
-  // estimation Only enabled when use_nonrd_pick_mode is != 0
+  // estimation. Only enabled when use_nonrd_pick_mode is != 0
   int use_fast_nonrd_pick_mode;
 
   // Reuse inter prediction in fast non-rd mode.
   int reuse_inter_pred_nonrd;
 
-  // Perform croase ME before calculating variance in variance-based partition
+  // Perform coarse ME before calculating variance in variance-based partition
   int estimate_motion_for_var_based_partition;
 } SPEED_FEATURES;