diff --git a/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementation.java b/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementation.java
index c5211b7d1ea688..755b43e645b9bb 100644
--- a/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementation.java
+++ b/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementation.java
@@ -37,7 +37,8 @@ public class FlatUIImplementation extends UIImplementation {
   public static FlatUIImplementation createInstance(
       ReactApplicationContext reactContext,
       List<ViewManager> viewManagers,
-      EventDispatcher eventDispatcher) {
+      EventDispatcher eventDispatcher,
+      boolean memoryImprovementEnabled) {
 
     RCTImageViewManager rctImageViewManager = findRCTImageManager(viewManagers);
     if (rctImageViewManager != null) {
@@ -61,7 +62,8 @@ public static FlatUIImplementation createInstance(
       rctImageViewManager,
       viewManagerRegistry,
       operationsQueue,
-      eventDispatcher
+      eventDispatcher,
+      memoryImprovementEnabled
     );
   }
 
@@ -73,17 +75,20 @@ public static FlatUIImplementation createInstance(
   private final ReactApplicationContext mReactContext;
   private @Nullable RCTImageViewManager mRCTImageViewManager;
   private final StateBuilder mStateBuilder;
+  private final boolean mMemoryImprovementEnabled;
 
   private FlatUIImplementation(
       ReactApplicationContext reactContext,
       @Nullable RCTImageViewManager rctImageViewManager,
       ViewManagerRegistry viewManagers,
       FlatUIViewOperationQueue operationsQueue,
-      EventDispatcher eventDispatcher) {
+      EventDispatcher eventDispatcher,
+      boolean memoryImprovementEnabled) {
     super(reactContext, viewManagers, operationsQueue, eventDispatcher);
     mReactContext = reactContext;
     mRCTImageViewManager = rctImageViewManager;
     mStateBuilder = new StateBuilder(operationsQueue);
+    mMemoryImprovementEnabled = memoryImprovementEnabled;
   }
 
   @Override
@@ -521,6 +526,9 @@ protected void applyUpdatesRecursive(
 
   @Override
   public void removeRootView(int rootViewTag) {
+    if (mMemoryImprovementEnabled) {
+      removeRootShadowNode(rootViewTag);
+    }
     mStateBuilder.removeRootView(rootViewTag);
   }
 
diff --git a/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementationProvider.java b/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementationProvider.java
index d3921faedf07f6..3a28f5df6ac78a 100644
--- a/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementationProvider.java
+++ b/ReactAndroid/src/main/java/com/facebook/react/flat/FlatUIImplementationProvider.java
@@ -20,11 +20,26 @@
  * UIImplementationProvider that creates instances of {@link FlatUIImplementation}.
  */
 public final class FlatUIImplementationProvider extends UIImplementationProvider {
+
+  private final boolean mMemoryImprovementEnabled;
+
+  public FlatUIImplementationProvider() {
+    mMemoryImprovementEnabled = true;
+  }
+
+  public FlatUIImplementationProvider(boolean memoryImprovementEnabled) {
+    mMemoryImprovementEnabled = memoryImprovementEnabled;
+  }
+
   @Override
   public FlatUIImplementation createUIImplementation(
       ReactApplicationContext reactContext,
       List<ViewManager> viewManagers,
       EventDispatcher eventDispatcher) {
-    return FlatUIImplementation.createInstance(reactContext, viewManagers, eventDispatcher);
+    return FlatUIImplementation.createInstance(
+      reactContext,
+      viewManagers,
+      eventDispatcher,
+      mMemoryImprovementEnabled);
   }
 }
diff --git a/ReactAndroid/src/main/java/com/facebook/react/uimanager/UIImplementation.java b/ReactAndroid/src/main/java/com/facebook/react/uimanager/UIImplementation.java
index f176d07a1a4935..7db28fd36a2153 100644
--- a/ReactAndroid/src/main/java/com/facebook/react/uimanager/UIImplementation.java
+++ b/ReactAndroid/src/main/java/com/facebook/react/uimanager/UIImplementation.java
@@ -133,10 +133,17 @@ public void registerRootView(
    * Unregisters a root node with a given tag.
    */
   public void removeRootView(int rootViewTag) {
-    mShadowNodeRegistry.removeRootNode(rootViewTag);
+    removeRootShadowNode(rootViewTag);
     mOperationsQueue.enqueueRemoveRootView(rootViewTag);
   }
 
+  /**
+   * Unregisters a root node with a given tag from the shadow node registry
+   */
+  public void removeRootShadowNode(int rootViewTag) {
+    mShadowNodeRegistry.removeRootNode(rootViewTag);
+  }
+
   /**
    * Invoked when native view that corresponds to a root node, or acts as a root view (ie. Modals)
    * has its size changed.