diff --git a/apps/graph/values/storage_function_parameter_controller.cpp b/apps/graph/values/storage_function_parameter_controller.cpp
index 667ab51a0..737cafada 100644
--- a/apps/graph/values/storage_function_parameter_controller.cpp
+++ b/apps/graph/values/storage_function_parameter_controller.cpp
@@ -19,7 +19,9 @@ bool StorageFunctionParameterController::handleEvent(Ion::Events::Event event) {
     switch (selectedRow()) {
       case 0:
       {
-        function()->setDisplayDerivative(!function()->displayDerivative());
+        bool isDisplayingDerivative = function()->displayDerivative();
+        function()->setDisplayDerivative(!isDisplayingDerivative);
+        m_valuesController->selectCellAtLocation(isDisplayingDerivative ? m_selectedFunctionColumn : m_selectedFunctionColumn + 1, m_valuesController->selectedRow());
         m_selectableTableView.reloadData();
         return true;
       }
@@ -57,9 +59,7 @@ int StorageFunctionParameterController::reusableCellCount() {
 
 void StorageFunctionParameterController::viewWillAppear() {
   StorageValuesFunctionParameterController::viewWillAppear();
-  if (function()->displayDerivative()) {
-    m_valuesController->selectCellAtLocation(m_valuesController->selectedColumn()+1, m_valuesController->selectedRow());
-  }
+  m_selectedFunctionColumn = m_valuesController->selectedColumn();
 }
 
 void StorageFunctionParameterController::willDisplayCellForIndex(HighlightCell * cell, int index) {
diff --git a/apps/graph/values/storage_function_parameter_controller.h b/apps/graph/values/storage_function_parameter_controller.h
index 205320927..c629c73b6 100644
--- a/apps/graph/values/storage_function_parameter_controller.h
+++ b/apps/graph/values/storage_function_parameter_controller.h
@@ -27,6 +27,8 @@ private:
 #endif
   MessageTableCellWithSwitch m_displayDerivativeColumn;
   StorageValuesController * m_valuesController;
+  // Index of the column corresponding to the function in the values controller
+  int m_selectedFunctionColumn;
 };
 
 }