diff --git a/native/jni/src/suggest/core/layout/proximity_info_params.cpp b/native/jni/src/suggest/core/layout/proximity_info_params.cpp
index a70dd7e..68bb0ae 100644
--- a/native/jni/src/suggest/core/layout/proximity_info_params.cpp
+++ b/native/jni/src/suggest/core/layout/proximity_info_params.cpp
@@ -24,9 +24,6 @@
 const float ProximityInfoParams::VERTICAL_SWEET_SPOT_SCALE_G = 0.5f;
 
 /* Per method constants */
-// Used by ProximityInfoStateUtils::initGeometricDistanceInfos()
-const float ProximityInfoParams::NEAR_KEY_NORMALIZED_SQUARED_THRESHOLD = 4.0f;
-
 // Used by ProximityInfoStateUtils::updateNearKeysDistances()
 const float ProximityInfoParams::NEAR_KEY_THRESHOLD_FOR_DISTANCE = 2.0f;
 
@@ -50,7 +47,7 @@
 const int ProximityInfoParams::LAST_POINT_SKIP_DISTANCE_SCALE = 4;
 
 // Used by ProximityInfoStateUtils::updateAlignPointProbabilities()
-const float ProximityInfoParams::MIN_PROBABILITY = 0.000001f;
+const float ProximityInfoParams::MIN_PROBABILITY = 0.000005f;
 const float ProximityInfoParams::MAX_SKIP_PROBABILITY = 0.95f;
 const float ProximityInfoParams::SKIP_FIRST_POINT_PROBABILITY = 0.01f;
 const float ProximityInfoParams::SKIP_LAST_POINT_PROBABILITY = 0.1f;
diff --git a/native/jni/src/suggest/core/layout/proximity_info_params.h b/native/jni/src/suggest/core/layout/proximity_info_params.h
index b8e9f5d..d9515c8 100644
--- a/native/jni/src/suggest/core/layout/proximity_info_params.h
+++ b/native/jni/src/suggest/core/layout/proximity_info_params.h
@@ -28,9 +28,6 @@
     static const float VERTICAL_SWEET_SPOT_SCALE;
     static const float VERTICAL_SWEET_SPOT_SCALE_G;
 
-    // Used by ProximityInfoStateUtils::initGeometricDistanceInfos()
-    static const float NEAR_KEY_NORMALIZED_SQUARED_THRESHOLD;
-
     // Used by ProximityInfoStateUtils::updateNearKeysDistances()
     static const float NEAR_KEY_THRESHOLD_FOR_DISTANCE;
 
diff --git a/native/jni/src/suggest/core/layout/proximity_info_state.cpp b/native/jni/src/suggest/core/layout/proximity_info_state.cpp
index b75c2ef..e585f90 100644
--- a/native/jni/src/suggest/core/layout/proximity_info_state.cpp
+++ b/native/jni/src/suggest/core/layout/proximity_info_state.cpp
@@ -91,7 +91,6 @@
         mSampledInputIndice.clear();
         mSampledLengthCache.clear();
         mSampledNormalizedSquaredLengthCache.clear();
-        mSampledNearKeySets.clear();
         mSampledSearchKeySets.clear();
         mSpeedRates.clear();
         mBeelineSpeedPercentiles.clear();
@@ -126,18 +125,17 @@
     if (mSampledInputSize > 0) {
         ProximityInfoStateUtils::initGeometricDistanceInfos(mProximityInfo, mSampledInputSize,
                 lastSavedInputSize, isGeometric, &mSampledInputXs, &mSampledInputYs,
-                &mSampledNearKeySets, &mSampledNormalizedSquaredLengthCache);
+                &mSampledNormalizedSquaredLengthCache);
         if (isGeometric) {
             // updates probabilities of skipping or mapping each key for all points.
             ProximityInfoStateUtils::updateAlignPointProbabilities(
                     mMaxPointToKeyLength, mProximityInfo->getMostCommonKeyWidth(),
                     mProximityInfo->getKeyCount(), lastSavedInputSize, mSampledInputSize,
                     &mSampledInputXs, &mSampledInputYs, &mSpeedRates, &mSampledLengthCache,
-                    &mSampledNormalizedSquaredLengthCache, &mSampledNearKeySets,
-                    mProximityInfo, &mCharProbabilities);
+                    &mSampledNormalizedSquaredLengthCache, mProximityInfo, &mCharProbabilities);
             ProximityInfoStateUtils::updateSampledSearchKeySets(mProximityInfo,
                     mSampledInputSize, lastSavedInputSize, &mSampledLengthCache,
-                    &mSampledNearKeySets, &mSampledSearchKeySets,
+                    &mCharProbabilities, &mSampledSearchKeySets,
                     &mSampledSearchKeyVectors);
             mMostProbableStringProbability = ProximityInfoStateUtils::getMostProbableString(
                     mProximityInfo, mSampledInputSize, &mCharProbabilities, mMostProbableString);
diff --git a/native/jni/src/suggest/core/layout/proximity_info_state.h b/native/jni/src/suggest/core/layout/proximity_info_state.h
index e253d95..d66121b 100644
--- a/native/jni/src/suggest/core/layout/proximity_info_state.h
+++ b/native/jni/src/suggest/core/layout/proximity_info_state.h
@@ -50,9 +50,9 @@
               mSampledInputXs(), mSampledInputYs(), mSampledTimes(), mSampledInputIndice(),
               mSampledLengthCache(), mBeelineSpeedPercentiles(),
               mSampledNormalizedSquaredLengthCache(), mSpeedRates(), mDirections(),
-              mCharProbabilities(), mSampledNearKeySets(), mSampledSearchKeySets(),
-              mSampledSearchKeyVectors(), mTouchPositionCorrectionEnabled(false),
-              mSampledInputSize(0), mMostProbableStringProbability(0.0f) {
+              mCharProbabilities(), mSampledSearchKeySets(), mSampledSearchKeyVectors(),
+              mTouchPositionCorrectionEnabled(false), mSampledInputSize(0),
+              mMostProbableStringProbability(0.0f) {
         memset(mInputProximities, 0, sizeof(mInputProximities));
         memset(mPrimaryInputWord, 0, sizeof(mPrimaryInputWord));
         memset(mMostProbableString, 0, sizeof(mMostProbableString));
@@ -216,10 +216,6 @@
     std::vector<float> mDirections;
     // probabilities of skipping or mapping to a key for each point.
     std::vector<hash_map_compat<int, float> > mCharProbabilities;
-    // The vector for the key code set which holds nearby keys for each sampled input point
-    // 1. Used to calculate the probability of the key
-    // 2. Used to calculate mSampledSearchKeySets
-    std::vector<ProximityInfoStateUtils::NearKeycodesSet> mSampledNearKeySets;
     // The vector for the key code set which holds nearby keys of some trailing sampled input points
     // for each sampled input point. These nearby keys contain the next characters which can be in
     // the dictionary. Specifically, currently we are looking for keys nearby trailing sampled
diff --git a/native/jni/src/suggest/core/layout/proximity_info_state_utils.cpp b/native/jni/src/suggest/core/layout/proximity_info_state_utils.cpp
index 638297e..72bb68f 100644
--- a/native/jni/src/suggest/core/layout/proximity_info_state_utils.cpp
+++ b/native/jni/src/suggest/core/layout/proximity_info_state_utils.cpp
@@ -188,13 +188,10 @@
         const int lastSavedInputSize, const bool isGeometric,
         const std::vector<int> *const sampledInputXs,
         const std::vector<int> *const sampledInputYs,
-        std::vector<NearKeycodesSet> *sampledNearKeySets,
         std::vector<float> *sampledNormalizedSquaredLengthCache) {
-    sampledNearKeySets->resize(sampledInputSize);
     const int keyCount = proximityInfo->getKeyCount();
     sampledNormalizedSquaredLengthCache->resize(sampledInputSize * keyCount);
     for (int i = lastSavedInputSize; i < sampledInputSize; ++i) {
-        (*sampledNearKeySets)[i].reset();
         for (int k = 0; k < keyCount; ++k) {
             const int index = i * keyCount + k;
             const int x = (*sampledInputXs)[i];
@@ -203,10 +200,6 @@
                     proximityInfo->getNormalizedSquaredDistanceFromCenterFloatG(
                             k, x, y, isGeometric);
             (*sampledNormalizedSquaredLengthCache)[index] = normalizedSquaredDistance;
-            if (normalizedSquaredDistance
-                    < ProximityInfoParams::NEAR_KEY_NORMALIZED_SQUARED_THRESHOLD) {
-                (*sampledNearKeySets)[i][k] = true;
-            }
         }
     }
 }
@@ -626,7 +619,6 @@
         const std::vector<float> *const sampledSpeedRates,
         const std::vector<int> *const sampledLengthCache,
         const std::vector<float> *const sampledNormalizedSquaredLengthCache,
-        std::vector<NearKeycodesSet> *sampledNearKeySets,
         const ProximityInfo *const proximityInfo,
         std::vector<hash_map_compat<int, float> > *charProbabilities) {
     charProbabilities->resize(sampledInputSize);
@@ -643,12 +635,10 @@
 
         float nearestKeyDistance = static_cast<float>(MAX_VALUE_FOR_WEIGHTING);
         for (int j = 0; j < keyCount; ++j) {
-            if ((*sampledNearKeySets)[i].test(j)) {
-                const float distance = getPointToKeyByIdLength(
-                        maxPointToKeyLength, sampledNormalizedSquaredLengthCache, keyCount, i, j);
-                if (distance < nearestKeyDistance) {
-                    nearestKeyDistance = distance;
-                }
+            const float distance = getPointToKeyByIdLength(
+                    maxPointToKeyLength, sampledNormalizedSquaredLengthCache, keyCount, i, j);
+            if (distance < nearestKeyDistance) {
+                nearestKeyDistance = distance;
             }
         }
 
@@ -744,27 +734,23 @@
         // Summing up probability densities of all near keys.
         float sumOfProbabilityDensities = 0.0f;
         for (int j = 0; j < keyCount; ++j) {
-            if ((*sampledNearKeySets)[i].test(j)) {
-                sumOfProbabilityDensities += distribution.getProbabilityDensity(
-                        proximityInfo->getKeyCenterXOfKeyIdG(j,
-                                NOT_A_COORDINATE /* referencePointX */, true /* isGeometric */),
-                        proximityInfo->getKeyCenterYOfKeyIdG(j,
-                                NOT_A_COORDINATE /* referencePointY */, true /* isGeometric */));
-            }
+            sumOfProbabilityDensities += distribution.getProbabilityDensity(
+                    proximityInfo->getKeyCenterXOfKeyIdG(j,
+                            NOT_A_COORDINATE /* referencePointX */, true /* isGeometric */),
+                    proximityInfo->getKeyCenterYOfKeyIdG(j,
+                            NOT_A_COORDINATE /* referencePointY */, true /* isGeometric */));
         }
 
         // Split the probability of an input point to keys that are close to the input point.
         for (int j = 0; j < keyCount; ++j) {
-            if ((*sampledNearKeySets)[i].test(j)) {
-                const float probabilityDensity = distribution.getProbabilityDensity(
-                        proximityInfo->getKeyCenterXOfKeyIdG(j,
-                                NOT_A_COORDINATE /* referencePointX */, true /* isGeometric */),
-                        proximityInfo->getKeyCenterYOfKeyIdG(j,
-                                NOT_A_COORDINATE /* referencePointY */, true /* isGeometric */));
-                const float probability = inputCharProbability * probabilityDensity
-                        / sumOfProbabilityDensities;
-                (*charProbabilities)[i][j] = probability;
-            }
+            const float probabilityDensity = distribution.getProbabilityDensity(
+                    proximityInfo->getKeyCenterXOfKeyIdG(j,
+                            NOT_A_COORDINATE /* referencePointX */, true /* isGeometric */),
+                    proximityInfo->getKeyCenterYOfKeyIdG(j,
+                            NOT_A_COORDINATE /* referencePointY */, true /* isGeometric */));
+            const float probability = inputCharProbability * probabilityDensity
+                    / sumOfProbabilityDensities;
+            (*charProbabilities)[i][j] = probability;
         }
     }
 
@@ -820,10 +806,9 @@
         for (int j = 0; j < keyCount; ++j) {
             hash_map_compat<int, float>::iterator it = (*charProbabilities)[i].find(j);
             if (it == (*charProbabilities)[i].end()){
-                (*sampledNearKeySets)[i].reset(j);
+                continue;
             } else if(it->second < ProximityInfoParams::MIN_PROBABILITY) {
                 // Erases from near keys vector because it has very low probability.
-                (*sampledNearKeySets)[i].reset(j);
                 (*charProbabilities)[i].erase(j);
             } else {
                 it->second = -logf(it->second);
@@ -835,9 +820,8 @@
 
 /* static */ void ProximityInfoStateUtils::updateSampledSearchKeySets(
         const ProximityInfo *const proximityInfo, const int sampledInputSize,
-        const int lastSavedInputSize,
-        const std::vector<int> *const sampledLengthCache,
-        const std::vector<NearKeycodesSet> *const sampledNearKeySets,
+        const int lastSavedInputSize, const std::vector<int> *const sampledLengthCache,
+        const std::vector<hash_map_compat<int, float> > *const charProbabilities,
         std::vector<NearKeycodesSet> *sampledSearchKeySets,
         std::vector<std::vector<int> > *sampledSearchKeyVectors) {
     sampledSearchKeySets->resize(sampledInputSize);
@@ -854,7 +838,12 @@
             if ((*sampledLengthCache)[j] - (*sampledLengthCache)[i] >= readForwordLength) {
                 break;
             }
-            (*sampledSearchKeySets)[i] |= (*sampledNearKeySets)[j];
+            for(const auto& charProbability : charProbabilities->at(j)) {
+                if (charProbability.first == NOT_AN_INDEX) {
+                    continue;
+                }
+                (*sampledSearchKeySets)[i].set(charProbability.first);
+            }
         }
     }
     const int keyCount = proximityInfo->getKeyCount();
diff --git a/native/jni/src/suggest/core/layout/proximity_info_state_utils.h b/native/jni/src/suggest/core/layout/proximity_info_state_utils.h
index 5d7a9c5..7aa20c3 100644
--- a/native/jni/src/suggest/core/layout/proximity_info_state_utils.h
+++ b/native/jni/src/suggest/core/layout/proximity_info_state_utils.h
@@ -71,13 +71,12 @@
             const std::vector<float> *const sampledSpeedRates,
             const std::vector<int> *const sampledLengthCache,
             const std::vector<float> *const sampledNormalizedSquaredLengthCache,
-            std::vector<NearKeycodesSet> *sampledNearKeySets,
             const ProximityInfo *const proximityInfo,
             std::vector<hash_map_compat<int, float> > *charProbabilities);
     static void updateSampledSearchKeySets(const ProximityInfo *const proximityInfo,
             const int sampledInputSize, const int lastSavedInputSize,
             const std::vector<int> *const sampledLengthCache,
-            const std::vector<NearKeycodesSet> *const sampledNearKeySets,
+            const std::vector<hash_map_compat<int, float> > *const charProbabilities,
             std::vector<NearKeycodesSet> *sampledSearchKeySets,
             std::vector<std::vector<int> > *sampledSearchKeyVectors);
     static float getPointToKeyByIdLength(const float maxPointToKeyLength,
@@ -87,7 +86,6 @@
             const int sampledInputSize, const int lastSavedInputSize, const bool isGeometric,
             const std::vector<int> *const sampledInputXs,
             const std::vector<int> *const sampledInputYs,
-            std::vector<NearKeycodesSet> *sampledNearKeySets,
             std::vector<float> *sampledNormalizedSquaredLengthCache);
     static void initPrimaryInputWord(const int inputSize, const int *const inputProximities,
             int *primaryInputWord);