diff --git a/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelView.java b/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelView.java
index 60462042b..d80c87b90 100644
--- a/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelView.java
+++ b/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelView.java
@@ -155,11 +155,6 @@ public void doCreatePartControl(Composite parent) {
         // Drill down
         fDrillDownAdapter = new DrillDownAdapter(fTreeViewer);
         
-        // Set drill down to home when disposing parent (not when the tree is disposed as the content provider is null at that point)
-        parent.addDisposeListener(e -> {
-            setDrillDownHome();
-        });
-        
         // Listen to Double-click and press Return Action
         fTreeViewer.addDoubleClickListener(new IDoubleClickListener() {
             @Override
@@ -214,12 +209,6 @@ public void init(IViewSite site, IMemento memento) throws PartInitException {
     
     @Override
     public void saveState(IMemento memento) {
-        // saveState() is called periodically by Eclipse when auto-saving the workbench so only reset the drill-down when closing the Workbench.
-        // This period is set in org.eclipse.ui.internal.IPreferenceConstants#WORKBENCH_SAVE_INTERVAL and the default is 5 minutes.
-        if(PlatformUI.getWorkbench().isClosing()) {
-            setDrillDownHome();
-        }
-        
         // Save expanded tree state
         TreeStateHelper.INSTANCE.saveStateToMemento(fTreeViewer, memento);
     }
diff --git a/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelViewer.java b/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelViewer.java
index da193e613..e0d835354 100644
--- a/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelViewer.java
+++ b/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeModelViewer.java
@@ -373,6 +373,13 @@ private String getAncestorFolderRenderText(IArchimateModelObject object) {
     
     // ========================= Model Providers =====================================
     
+    
+    private Object[] rootExpandedTreeElements;
+    
+    Object[] getRootExpandedElements() {
+        return rootExpandedTreeElements;
+    }
+    
     /**
      *  Content Provider
      */
@@ -380,6 +387,7 @@ private class ModelTreeViewerContentProvider implements ITreeContentProvider {
         
         @Override
         public void inputChanged(Viewer v, Object oldInput, Object newInput) {
+            rootExpandedTreeElements = oldInput == IEditorModelManager.INSTANCE ? getVisibleExpandedElements() : null;
         }
         
         @Override
diff --git a/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeStateHelper.java b/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeStateHelper.java
index ea498c6e6..5917f8073 100644
--- a/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeStateHelper.java
+++ b/com.archimatetool.editor/src/com/archimatetool/editor/views/tree/TreeStateHelper.java
@@ -13,7 +13,6 @@
 import java.util.Map.Entry;
 
 import org.eclipse.emf.ecore.EObject;
-import org.eclipse.jface.viewers.TreeViewer;
 import org.eclipse.ui.IMemento;
 
 import com.archimatetool.editor.model.IEditorModelManager;
@@ -44,9 +43,6 @@ public class TreeStateHelper {
     // Expanded tree elements
     private List<Object> expandedElements;
     
-    // State has been restored from memento the first time the Tree is created
-    private boolean restoredFromMemento;
-    
     private TreeStateHelper() {}
 
     /**
@@ -56,12 +52,10 @@ private TreeStateHelper() {}
      * This is called from {@link TreeModelView#init(org.eclipse.ui.IViewSite, IMemento)}, but we only want to do this once.
      */
     void setMemento(IMemento memento) {
-        if(restoredFromMemento || memento == null) {
+        if(memento == null) {
             return;
         }
         
-        restoredFromMemento = true;
-        
         IMemento expandedMem = memento.getChild(MEMENTO_EXPANDED);
         if(expandedMem != null) {
             expandedElements = new ArrayList<>();
@@ -101,16 +95,7 @@ void setMemento(IMemento memento) {
      * Restore expanded elements on TreeView part creation.
      * This is called from {@link TreeModelView#doCreatePartControl(org.eclipse.swt.widgets.Composite)
      */
-    void restoreExpandedTreeElements(TreeViewer viewer) {
-        // Store expanded tree elements if the TreeViewer is closed so they can be restored when it's re-opened
-        // We could restore from the memento but this is more efficient and the drill-down is reset when closing the Tree.
-        viewer.getTree().addDisposeListener(e -> {
-            expandedElements = new ArrayList<>();
-            for(Object element : viewer.getVisibleExpandedElements()) {
-                expandedElements.add(element);
-            }
-        });
-
+    void restoreExpandedTreeElements(TreeModelViewer viewer) {
         if(expandedElements != null) {
             for(Object element : expandedElements) {
                 viewer.expandToLevel(element, 1);
@@ -123,10 +108,12 @@ void restoreExpandedTreeElements(TreeViewer viewer) {
      * Save expanded state of tree elements to a memento.
      * This is called from {@link TreeModelView#saveState(IMemento)}
      */
-    void saveStateToMemento(TreeViewer viewer, IMemento memento) {
+    void saveStateToMemento(TreeModelViewer viewer, IMemento memento) {
         Map<File, String> map = new HashMap<>();
         
-        for(Object object : viewer.getVisibleExpandedElements()) {
+        Object[] elements = viewer.getRootExpandedElements() == null ? viewer.getVisibleExpandedElements() : viewer.getRootExpandedElements();
+        
+        for(Object object : elements) {
             if(object instanceof IArchimateModelObject modelObject) {
                 // Only store if model has been saved to file
                 File file = modelObject.getArchimateModel().getFile();