diff -u b/term_reference_tree.js b/term_reference_tree.js
--- b/term_reference_tree.js
+++ b/term_reference_tree.js
@@ -1,6 +1,6 @@
 (function($) {
 
-Drupal.ajax.prototype.commands.term_reference_tree_append_children = function(ajax, response, status) {
+Drupal.ajax.prototype.commands.term_reference_tree_append_children = function (ajax, response, status) {
   $(ajax.wrapper).siblings('.term-reference-tree-button').removeClass('term-reference-tree-collapsed');
   $newElement = $('<div></div>').html(response.data);
   Drupal.attachBehaviors($newElement);
diff -u b/term_reference_tree.module b/term_reference_tree.module
--- b/term_reference_tree.module
+++ b/term_reference_tree.module
@@ -65,51 +65,62 @@
   );
 }
 
-
 /**
  * This function returns a taxonomy term hierarchy in a nested array.
  *
- * @param $tid
+ * @param int $tid
  *   The ID of the root term.
- * @param $vid
+ * @param int $vid
  *   The vocabulary ID to restrict the child search.
+ * @param array $allowed
+ *   (Optional) An array of allowed tids.
+ * @param array $default
+ *   (Optional) An array of tids to set as default.
+ * @param array $opened
+ *   (Optional) An array of tids currently opened with ajax.
+ * @param bool $use_ajax
+ *   (Optional) TRUE if using ajax mode, default to FALSE.
+ * @param int $max_depth
+ *   (Optional) The max depth, default to NULL (unlimited).
  *
- * @return
+ * @return array
  *   A nested array of the term's child objects.
  */
-function _term_reference_tree_get_term_hierarchy($terms, $allowed, $default = array(), $opened = array(), $use_ajax = FALSE) {
-  foreach ($terms as $tid => &$term) {
-    // Remove filtered items.
-    if (!empty($allowed) and !array_key_exists($term->tid, $allowed)) {
-      unset($terms[$tid]);
-    }
-    else {
-      // By default set children to FALSE.
-      $term->children = FALSE;
-      $term->children_selected = FALSE;
-
-      // If there are any children elements then check if we need to add them
-      // for rendering or leave them for ajax.
-      if ($children = _term_reference_tree_get_children($term->tid, $term->vid)) {
-        $term->children = TRUE;
-        // Selected items need to be displayed from the begining.
-        if (in_array($term->tid, $default) || in_array($term->tid, $opened) || !$use_ajax) {
-          if ($children = _term_reference_tree_get_term_hierarchy($children, $allowed, $default, $opened, $use_ajax)) {
-
-            // Check if at least one children is selected.
-            foreach ($children as $child) {
-              if (isset($default[$child->tid])) {
-                $term->children_selected = TRUE;
-                break;
-              }
-            }
-
-            if ($term->children_selected || in_array($term->tid, $opened) || !$use_ajax) {
-              $term->children = $children;
-            }
-          }
+function _term_reference_tree_get_term_hierarchy($tid, $vid, $allowed = array(), $default = array(), $opened = array(), $use_ajax = FALSE, $max_depth = NULL) {
+  $tree = _term_reference_tree_taxonomy_get_tree($vid);
+  if (!empty($allowed)) {
+    $tree['terms'] = array_intersect_key($tree['terms'], $allowed);
+  }
+
+  $default_parents = _term_reference_tree_taxonomy_get_parents_all($vid, $default);
+  return _term_reference_tree_get_term_hierarchy_recursive($tid, $tree, $default_parents, $opened, $use_ajax, $max_depth);
+}
+
+/**
+ * Recursive helper function for _term_reference_tree_get_term_hierarchy().
+ */
+function _term_reference_tree_get_term_hierarchy_recursive($parent_tid, $tree, $default, $opened, $use_ajax, $max_depth, $depth = 1) {
+  $terms = array();
+
+  if (isset($tree['children'][$parent_tid])) {
+    foreach($tree['children'][$parent_tid] as $child_tid) {
+      $term = $tree['terms'][$child_tid];
+      $max_depth_reached = isset($max_depth) && $depth >= $max_depth;
+      $term->has_children = isset($tree['children'][$term->tid]) && !$max_depth_reached;
+
+      if ($term->has_children) {
+        // Process children if:
+        // - we don't use ajax,
+        // - OR if this term is opened via ajax,
+        // - OR if one of the children of this term is in $default.
+        if (!$use_ajax || in_array($term->tid, $opened) || in_array($term->tid, $default)) {
+          $term->children = _term_reference_tree_get_term_hierarchy_recursive($term->tid, $tree, $default, $opened, $use_ajax, $max_depth, ++$depth);
         }
       }
+
+      // Resolve label and add the term to the final array.
+      _term_reference_tree_taxonomy_resolve_label($term);
+      $terms[$term->tid] = $term;
     }
   }
 
@@ -117,41 +128,99 @@
 }
 
 /**
- * This function is like taxonomy_get_children, except it doesn't load the entire term.
+ * Gets all terms and their parent and children hierarchy for a given vid.
+ *
+ * @param int $vid
+ *   The taxonomy vocabulary ID.
  *
- * @param $tid
- *   The ID of the term whose children you want to get.
- * @param $vid
- *   The vocabulary ID.
- *
- * @return
- *   An array of taxonomy terms, each in the form array('tid' => $tid, 'name' => $name)
- */
-function _term_reference_tree_get_children($tid, $vid) {
-  // DO NOT LOAD TAXONOMY TERMS HERE
-  // Taxonomy terms take a lot of time and memory to load, and this can be
-  // very bad on large vocabularies.  Instead, we load the term as necessary
-  // in cases where it's needed (such as using tokens or when the locale
-  // module is enabled).
-  $query = db_select('taxonomy_term_data', 't');
-  $query->join('taxonomy_term_hierarchy', 'h', 't.tid = h.tid');
-  $query->join('taxonomy_vocabulary', 'v', 'v.vid = t.vid');
-  $query->fields('t', array('tid', 'name', 'vid'));
-  $query->addField('v', 'machine_name', 'vocabulary_machine_name');
-  $query->condition('t.vid', $vid);
-  $query->condition('h.parent', $tid);
-  $query->addTag('term_access');
-  $query->addTag('translatable');
-  $query->orderBy('t.weight');
-  $query->orderBy('t.name');
-  $results = $query->execute();
+ * @return array
+ *   An array containing all terms, children and parent hierarchy.
+ */
+function _term_reference_tree_taxonomy_get_tree($vid) {
+  $tree = &drupal_static(__FUNCTION__, array());
 
-  $terms = array();
-  while ($term = $results->fetchObject()) {
-    $terms[$term->tid] = $term;
+  if (!isset($tree[$vid])) {
+    $query = db_select('taxonomy_term_data', 't');
+    $query->join('taxonomy_term_hierarchy', 'h', 'h.tid = t.tid');
+    $query->join('taxonomy_vocabulary', 'v', 't.vid = v.vid');
+    $query->addField('v', 'machine_name', 'vocabulary_machine_name');
+    $result = $query->addTag('translatable')
+      ->addTag('taxonomy_term_access')
+      ->fields('t')
+      ->fields('h', array('parent'))
+      ->condition('t.vid', $vid)
+      ->orderBy('t.weight')
+      ->orderBy('t.name')
+      ->execute();
+
+    foreach ($result as $term) {
+      $tree[$vid]['children'][$term->parent][] = $term->tid;
+      $tree[$vid]['parents'][$term->tid][] = $term->parent;
+      $tree[$vid]['terms'][$term->tid] = $term;
+    }
   }
 
-  return $terms;
+  return $tree[$vid];
+}
+
+/**
+ * Gets all parent tids for a given vid and children tids.
+ *
+ * @param int $vid
+ *   The taxonomy vocabulary ID.
+ * @param array $children_tids
+ *   An array of tids for which to retrieve the parents.
+ * @param array $hierarchy
+ *   The "parents" hierarchy from _term_reference_tree_taxonomy_get_tree().
+ *
+ * @return array
+ *   All parent tids of the given children tids.
+ */
+function _term_reference_tree_taxonomy_get_parents_all($vid, $children_tids) {
+  $tree = _term_reference_tree_taxonomy_get_tree($vid);
+  return _term_reference_tree_taxonomy_get_parents_all_recursive($vid, $children_tids, $tree);
+}
+
+/**
+ * Helper function for _term_reference_tree_taxonomy_get_parents_all().
+ */
+function _term_reference_tree_taxonomy_get_parents_all_recursive($vid, $children_tids, $tree) {
+  $parents = array();
+  foreach ($children_tids as $child_tid) {
+    if (isset($tree['parents'][$child_tid])) {
+      foreach ($tree['parents'][$child_tid] as $parent_tid) {
+        $parents[$parent_tid] = $parent_tid;
+      }
+      $parents += _term_reference_tree_taxonomy_get_parents_all_recursive($vid, $tree['parents'][$child_tid], $tree);
+    }
+  }
+  return $parents;
+}
+
+/**
+ * Resolves the label to use in the widget for a given taxonomy term.
+ *
+ * @param object $term
+ *   The taxonomy term object.
+ */
+function _term_reference_tree_taxonomy_resolve_label($term) {
+  // Ajax calls the process function twice so we set a property of the term
+  // in order no to resolve it twice.
+  if (!isset($term->term_reference_tree_label_resolved) || !$term->term_reference_tree_label_resolved) {
+    // If a callback is implemented to get the label name use it.
+    if ($callback = variable_get('term_reference_tree_taxonomy_resolve_label_callback')) {
+      if (function_exists($callback)) {
+        $term->name = call_user_func($callback, $term);
+      }
+    }
+    // Otherwise fallback to entity_label().
+    elseif (module_exists('locale')) {
+      $term = taxonomy_term_load($term->tid);
+      $term->name = entity_label('taxonomy_term', $term);
+    }
+
+    $term->term_reference_tree_label_resolved = TRUE;
+  }
 }
 
 function _term_reference_tree_get_parent($tid) {
@@ -186,30 +254,0 @@
-
-/**
- * Return an array of options.
- *
- * This function converts a list of taxonomy terms to a key/value list of options.
- *
- * @param $terms
- *   An array of taxonomy term IDs.
- * @param $allowed
- *   An array containing the terms allowed by the filter view
- * @param $filter
- *   A string defining the view to filter by (only used to detect whether view
- *   filtering is enabled
- *
- * @return
- *   A key/value array of taxonomy terms (name => id)
- */
-function _term_reference_tree_get_options(&$terms, &$allowed, $filter) {
-  $options = array();
-
-  if (is_array($terms) && count($terms) > 0) {
-    foreach ($terms as $term) {
-      if (!$filter || (is_array($allowed) && $allowed[$term->tid])) {
-        $options[$term->tid] = entity_label('taxonomy_term', $term);
-        $options += _term_reference_tree_get_options($term->children, $allowed, $filter);
-      }
-    }
-  }
-  return $options;
-}
diff -u b/term_reference_tree.widget.inc b/term_reference_tree.widget.inc
--- b/term_reference_tree.widget.inc
+++ b/term_reference_tree.widget.inc
@@ -117,6 +117,7 @@
       $form['tokens_list'] = array(
         '#theme' => 'token_tree',
         '#token_types' => array('term'),
+        '#dialog' => TRUE,
       );
     }
     else {
@@ -282,65 +283,42 @@
  */
 function term_reference_tree_process_checkbox_tree($element, $form_state) {
   if (is_array($form_state)) {
-    if (!empty($element['#max_choices']) && $element['#max_choices'] != '-1') {
+    $allowed = !empty($element['#filter_view']) ? _term_reference_tree_get_allowed_values($element['#filter_view']) : '';
+    $default = !empty($element['#default_value']) ? $element['#default_value'] : array();
+    $opened = !empty($form_state['children']) ? $form_state['children'] : array();
+    $max_depth = !empty($element['#max_depth']) ? $element['#max_depth'] : NULL;
+    $element['#options_tree'] = _term_reference_tree_get_term_hierarchy($element['#parent_tid'], $element['#vocabulary']->vid, $allowed, $default, $opened, $element['#use_ajax'], $max_depth);
+
+    $max_choices = !empty($element['#max_choices']) ? $element['#max_choices'] : 1;
+    if ($max_choices != '-1') {
       drupal_add_js(array('term_reference_tree' => array('trees' => array($element['#id'] => array('max_choices' => $element['#max_choices'])))), 'setting');
     }
 
-    $allowed = '';
-    if ($element['#filter_view'] != '') {
-      $allowed = _term_reference_tree_get_allowed_values($element['#filter_view']);
-    }
-
-    $value = !empty($element['#default_value']) ? $element['#default_value'] : array();
-    if (empty($element['#options'])) {
-      // Get default terms and their parents.
-      $default = array();
-      foreach ($value as $tid) {
-        $parents = taxonomy_get_parents_all($tid);
-        foreach ($parents as $parent) {
-          $default[$parent->tid] = $parent->tid;
-        }
-      }
-      $opened = array();
-      if (isset($form_state['children'])) {
-        foreach ($form_state['children'] as $tid) {
-          $opened[$tid] = $tid;
-        }
-      }
-
-      // Get term tree.
-      $children = _term_reference_tree_get_children($element['#parent_tid'], $element['#vocabulary']->vid);
-      $element['#options_tree'] = _term_reference_tree_get_term_hierarchy($children, $allowed, $default, $opened, $element['#use_ajax']);
-
-      $required = $element['#required'];
-      if ($element['#max_choices'] == 1 && !$required) {
-        array_unshift($element['#options_tree'], (object) array(
-          'tid' => '',
-          'name' => 'N/A',
-          'depth' => 0,
-          'vocabulary_machine_name' => $element['#vocabulary']->machine_name,
-        ));
-      }
-      $element['#options'] = _term_reference_tree_get_options($element['#options_tree'], $allowed, $element['#filter_view']);
-    }
-
-    $terms = !empty($element['#options_tree']) ? $element['#options_tree'] : array();
-    $max_choices = !empty($element['#max_choices']) ? $element['#max_choices'] : 1;
-    if (array_key_exists('#select_parents', $element) && $element['#select_parents']) {
-      $element['#attributes']['class'][] = 'select-parents';
+    if ($max_choices == 1 && !$element['#required']) {
+      array_unshift($element['#options_tree'], (object) array(
+        'tid' => '',
+        'name' => t('N/A'),
+        'depth' => 0,
+        'vocabulary_machine_name' => $element['#vocabulary']->machine_name,
+        'has_children' => FALSE,
+      ));
     }
 
     if ($max_choices != 1) {
       $element['#tree'] = TRUE;
     }
 
+    if (!empty($element['#select_parents'])) {
+      $element['#attributes']['class'][] = 'select-parents';
+    }
+
     $tree = new stdClass();
-    $tree->children = $terms;
-    $element[] = _term_reference_tree_build_level($element, $tree, $form_state, $value, $max_choices, array(), 1);
+    $tree->children = $element['#options_tree'];
+    $tree->has_children = TRUE;
+    $element[] = _term_reference_tree_build_level($element, $tree, $form_state, $default, $max_choices, array(), 1);
 
     // Add a track list element?
-    $track_list = !empty($element['#track_list']) && $element['#track_list'];
-    if ($track_list) {
+    if (!empty($element['#track_list'])) {
       $element[] = array(
         '#type' => 'checkbox_tree_track_list',
         '#max_choices' => $max_choices,
@@ -692,39 +670,20 @@
  *   possibly a checkbox_tree_level element as well.
  */
 function _term_reference_tree_build_item($element, $term, $form_state, $value, $max_choices, $parent_tids, $parent, $depth) {
-  $start_minimized = FALSE;
-  if (array_key_exists('#start_minimized', $element)) {
-    $start_minimized = $element['#start_minimized'];
-  }
-
-  $leaves_only = FALSE;
-  if (array_key_exists('#leaves_only', $element)) {
-    $leaves_only = $element['#leaves_only'];
-  }
-
-  $t = NULL;
-  if (module_exists('locale') && !empty($term->tid) && FALSE) {
-    $t = taxonomy_term_load($term->tid);
-    $term_name = entity_label('taxonomy_term', $t);
-  }
-  else {
-    $term_name = $term->name;
-  }
   $container = array(
     '#type' => 'checkbox_tree_item',
     '#max_choices' => $max_choices,
-    '#leaves_only' => $leaves_only,
-    '#term_name' => $term_name,
+    '#leaves_only' => isset($element['#leaves_only']) ? $element['#leaves_only'] : FALSE,
+    '#term_name' => $term->name,
     '#level_start_minimized' => FALSE,
     '#depth' => $depth,
     '#has_children' => FALSE,
   );
 
-  if (!$element['#leaves_only'] || !$term->children) {
-    $name = "edit-" . str_replace('_', '-', $element['#field_name']);
+  if (!$element['#leaves_only'] || empty($term->children)) {
     $e = array(
       '#type' => ($max_choices == 1) ? 'radio' : 'checkbox',
-      '#title' => $term_name,
+      '#title' => $term->name,
       '#on_value' => $term->tid,
       '#off_value' => 0,
       '#return_value' => $term->tid,
@@ -734,10 +693,9 @@
       '#ajax' => isset($element['#ajax']) ? $element['#ajax'] : NULL,
     );
 
-    if ($element['#token_display'] != '' && module_exists('token')) {
-      if (!$t) {
-        $t = taxonomy_term_load($term->tid);
-      }
+    if (!empty($element['#token_display']) && module_exists('token')) {
+      // @TODO Load term and resolve tokens in the process part?
+      $t = taxonomy_term_load($term->tid);
       $e['#title'] = token_replace($element['#token_display'], array('term' => $t), array('clear' => TRUE));
     }
 
@@ -750,13 +708,13 @@
   else {
     $e = array(
       '#type' => 'checkbox_tree_label',
-      '#value' => $term_name,
+      '#value' => $term->name,
     );
   }
 
   $container[$term->tid] = $e;
 
-  if (($depth + 1 <= $element['#max_depth'] || !$element['#max_depth']) && $term->children) {
+  if ($term->has_children) {
     $container['#has_children'] = TRUE;
     $parents = $parent_tids;
     $parents[] = $term->tid;
@@ -789,51 +747,39 @@
  *   A completed checkbox_tree_level element.
  */
 function _term_reference_tree_build_level($element, $term, $form_state, $value, $max_choices, $parent_tids, $depth) {
-  $start_minimized = FALSE;
-  if (array_key_exists('#start_minimized', $element)) {
-    $start_minimized = $element['#start_minimized'];
-  }
-
-  $leaves_only = FALSE;
-  if (array_key_exists('#leaves_only', $element)) {
-    $leaves_only = $element['#leaves_only'];
-  }
-
   $container = array(
     '#max_choices' => $max_choices,
-    '#leaves_only' => $leaves_only,
-    '#start_minimized' => $start_minimized,
-    '#level_start_minimized' => $depth > 1 && $element['#start_minimized'] && !($term->children_selected),
+    '#leaves_only' => isset($element['#leaves_only']) ? $element['#leaves_only'] : FALSE,
+    '#start_minimized' => isset($element['#start_minimized']) ? $element['#start_minimized'] : FALSE,
+    '#level_start_minimized' => $depth > 1 &&  $element['#start_minimized'] && empty($term->children),
     '#depth' => $depth,
   );
 
-  // Adds checkbox_tree_level when children exists
-  // or ajax wrapper when children are not loaded.
-  if ($term->children) {
-    if (is_array($term->children)) {
-      $container['#type'] = 'checkbox_tree_level';
-      foreach ($term->children as $child) {
-        $container[$child->tid] = _term_reference_tree_build_item($element, $child, $form_state, $value, $max_choices, $parent_tids, $container, $depth);
-      }
-    }
-    else {
-      $id = drupal_html_id('term-reference-tree-widget-level-ajax-' . $term->tid);
-      $container['#type'] = 'container';
-      $container['#attributes']['id'] = $id;
-      $container['#attributes']['class'][] = 'term-reference-tree-widget-level-ajax';
-      $container['load_children_' . $term->tid] = array(
-        '#type' => 'button',
-        '#attributes' => array('class' => array('term-reference-tree-widget-level-ajax-button')),
-        '#value' => 'load children ' . $term->tid,
-        '#limit_validation_errors' => array(),
-        '#term' => $term,
-        '#ajax' => array(
-          'wrapper' => $id,
-          'callback' => 'term_reference_tree_children_callback',
-        ),
-      );
+  // Adds checkbox_tree_level when children exists or ajax wrapper when children
+  // are not loaded (has_children = TRUE but children is empty).
+  if (!empty($term->children)) {
+    $container['#type'] = 'checkbox_tree_level';
+    foreach ($term->children as $child) {
+      $container[$child->tid] = _term_reference_tree_build_item($element, $child, $form_state, $value, $max_choices, $parent_tids, $container, $depth);
     }
   }
+  else {
+    $id = drupal_html_id('term-reference-tree-widget-level-ajax-' . $term->tid);
+    $container['#type'] = 'container';
+    $container['#attributes']['id'] = $id;
+    $container['#attributes']['class'][] = 'term-reference-tree-widget-level-ajax';
+    $container['load_children_' . $term->tid] = array(
+      '#type' => 'button',
+      '#attributes' => array('class' => array('term-reference-tree-widget-level-ajax-button')),
+      '#value' => 'load children ' . $term->tid,
+      '#limit_validation_errors' => array(),
+      '#term' => $term,
+      '#ajax' => array(
+        'wrapper' => $id,
+        'callback' => 'term_reference_tree_children_callback',
+      ),
+    );
+  }
 
   return $container;
 }
@@ -842,9 +788,8 @@
  * Ajax callback to load children for a list.
  */
 function term_reference_tree_children_callback($form, $form_state) {
-  $commands = array();
-  $parents = array_slice($form_state['triggering_element']['#parents'], 0, -1);
-  $element = drupal_array_get_nested_value($form, $parents);
+  $button = $form_state['triggering_element'];
+  $element = drupal_array_get_nested_value($form, array_slice($button['#array_parents'], 0, -1));
   $commands[] = array(
     'command' => 'term_reference_tree_append_children',
     'data' => drupal_render($element),