diff --git a/smartlearner/__init__.py b/smartlearner/__init__.py
index 5159ac3..d6347f4 100644
--- a/smartlearner/__init__.py
+++ b/smartlearner/__init__.py
@@ -1,5 +1,10 @@
 from .interfaces.model import Model
 from .interfaces.dataset import Dataset
+from .interfaces.loss import Loss
+from .interfaces.task import Task
+
 from .trainer import Trainer
 
 from .tasks import tasks
+from .tasks import views
+from .tasks import stopping_criteria
diff --git a/smartlearner/batch_schedulers/__init__.py b/smartlearner/batch_schedulers/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/smartlearner/batch_schedulers/minibatch.py b/smartlearner/batch_schedulers/minibatch.py
new file mode 100644
index 0000000..54ea692
--- /dev/null
+++ b/smartlearner/batch_schedulers/minibatch.py
@@ -0,0 +1,46 @@
+import theano
+import numpy as np
+
+from ..interfaces.batch_scheduler import BatchScheduler
+
+
+class MiniBatchScheduler(BatchScheduler):
+    def __init__(self, dataset, batch_size):
+        super(MiniBatchScheduler, self).__init__(dataset)
+        self._shared_batch_size = theano.shared(np.array(0, dtype='i4'))
+        self.batch_size = batch_size
+        self.shared_batch_count = theano.shared(np.array(0, dtype='i4'))
+
+    @property
+    def batch_size(self):
+        return self._shared_batch_size.get_value()
+
+    @batch_size.setter
+    def batch_size(self, value):
+        self._shared_batch_size.set_value(np.array(value, dtype='i4'))
+        self.nb_updates_per_epoch = int(np.ceil(len(self.dataset)/self.batch_size))
+
+    @property
+    def updates(self):
+        return {}  # No updates
+
+    @property
+    def givens(self):
+        start = self.shared_batch_count * self._shared_batch_size
+        end = (self.shared_batch_count + 1) * self._shared_batch_size
+
+        if self.dataset.targets is not None:
+            return {self.dataset.symb_inputs: self.dataset.inputs[start:end],
+                    self.dataset.symb_targets: self.dataset.targets[start:end]}
+        else:
+            return {self.dataset.symb_inputs: self.dataset.inputs[start:end]}
+
+    def __iter__(self):
+        for batch_count in range(self.nb_updates_per_epoch):
+            self.shared_batch_count.set_value(batch_count)
+            yield batch_count + 1
+
+
+class FullBatchScheduler(MiniBatchScheduler):
+    def __init__(self, dataset):
+        super(FullBatchScheduler, self).__init__(dataset, batch_size=len(self.dataset))
diff --git a/smartlearner/update_rules/__init__.py b/smartlearner/direction_modifiers/__init__.py
similarity index 76%
rename from smartlearner/update_rules/__init__.py
rename to smartlearner/direction_modifiers/__init__.py
index f1121ac..33afb25 100644
--- a/smartlearner/update_rules/__init__.py
+++ b/smartlearner/direction_modifiers/__init__.py
@@ -1,3 +1,2 @@
-from .update_rule import UpdateRule
 from .decreasing_learning_rate import DecreasingLearningRate
 from .constant_learning_rate import ConstantLearningRate
diff --git a/smartlearner/update_rules/constant_learning_rate.py b/smartlearner/direction_modifiers/constant_learning_rate.py
similarity index 100%
rename from smartlearner/update_rules/constant_learning_rate.py
rename to smartlearner/direction_modifiers/constant_learning_rate.py
diff --git a/smartlearner/direction_modifiers/decreasing_learning_rate.py b/smartlearner/direction_modifiers/decreasing_learning_rate.py
new file mode 100644
index 0000000..eaca014
--- /dev/null
+++ b/smartlearner/direction_modifiers/decreasing_learning_rate.py
@@ -0,0 +1,44 @@
+from collections import OrderedDict
+
+import numpy as np
+
+from ..utils import sharedX
+from ..interfaces.direction_modifier import DirectionModifier
+
+
+class DecreasingLearningRate(DirectionModifier):
+    def __init__(self, lr, dc=0.):
+        """
+        Implements a decreasing learning rate update rule.
+
+        Parameters
+        ----------
+        lr: float
+            Learning rate.
+        dc: float in [0,1) (optional)
+            Decreasing constant (decay). Default: 0.
+        """
+        if dc < 0. or dc >= 1:
+            raise ValueError("`dc` ({}) must be between 0 (inclusive) and 1 (exclusive)!".format(dc))
+
+        super(DecreasingLearningRate, self).__init__()
+        self.lr = lr
+        self.dc = dc
+        self._updates = OrderedDict()
+
+    def _get_updates(self):
+        return self._updates
+
+    def apply(self, directions):
+        new_directions = OrderedDict()
+
+        for param, gparam in directions.items():
+            lr = sharedX(self.lr * np.ones_like(param.get_value()), name='lr_' + param.name)
+
+            if self.dc != 0.:
+                # Decrease the learning rate by a factor of `dc` after each update.
+                self._updates[lr] = self.dc * lr
+
+            new_directions[param] = lr * gparam
+
+        return new_directions
diff --git a/smartlearner/interfaces/__init__.py b/smartlearner/interfaces/__init__.py
index e69de29..889fe01 100644
--- a/smartlearner/interfaces/__init__.py
+++ b/smartlearner/interfaces/__init__.py
@@ -0,0 +1,7 @@
+from .batch_scheduler import BatchScheduler
+from .dataset import Dataset
+from .model import Model
+from .loss import Loss
+from .task import Task
+from .direction_modifier import DirectionModifier
+from .direction_modifier import ParamModifier
diff --git a/smartlearner/batch_scheduler.py b/smartlearner/interfaces/batch_scheduler.py
similarity index 84%
rename from smartlearner/batch_scheduler.py
rename to smartlearner/interfaces/batch_scheduler.py
index 99f33ec..2324f97 100644
--- a/smartlearner/batch_scheduler.py
+++ b/smartlearner/interfaces/batch_scheduler.py
@@ -1,4 +1,4 @@
-from abc import ABCMeta, abstractmethod
+from abc import ABCMeta, abstractmethod, abstractproperty
 
 import numpy as np
 import theano
@@ -11,8 +11,16 @@ class BatchScheduler(object):
         self.dataset = dataset
 
     @property
+    def tasks(self):
+        return []
+
+    @abstractproperty
+    def updates(self):
+        raise NotImplementedError("Subclass of 'BatchScheduler' must implement property 'updates'.")
+
+    @abstractproperty
     def givens(self):
-        raise NotImplementedError("Subclass of 'BatchScheduler' must implement 'givens'.")
+        raise NotImplementedError("Subclass of 'BatchScheduler' must implement property 'givens'.")
 
     @abstractmethod
     def __iter__(self):
@@ -40,6 +48,10 @@ class MiniBatchScheduler(BatchScheduler):
         self._shared_batch_size.set_value(np.array(value, dtype='i4'))
         self.nb_updates_per_epoch = int(np.ceil(len(self.dataset)/self.batch_size))
 
+    @property
+    def updates(self):
+        return {}  # No updates
+
     @property
     def givens(self):
         start = self.shared_batch_count * self._shared_batch_size
@@ -59,4 +71,4 @@ class MiniBatchScheduler(BatchScheduler):
 
 class FullBatchScheduler(MiniBatchScheduler):
     def __init__(self, dataset):
-        super(FullBatchScheduler, self).__init__(dataset, batch_size=len(self.dataset))
\ No newline at end of file
+        super(FullBatchScheduler, self).__init__(dataset, batch_size=len(self.dataset))
diff --git a/smartlearner/interfaces/direction_modifier.py b/smartlearner/interfaces/direction_modifier.py
new file mode 100644
index 0000000..9823a42
--- /dev/null
+++ b/smartlearner/interfaces/direction_modifier.py
@@ -0,0 +1,41 @@
+from abc import ABCMeta, abstractmethod
+
+
+class DirectionModifier(object):
+    __metaclass__ = ABCMeta
+
+    @property
+    def tasks(self):
+        return []
+
+    @property
+    def updates(self):
+        return self._get_updates()
+
+    @abstractmethod
+    def _get_updates(self):
+        raise NotImplementedError("Subclass of 'DirectionModifier' must implement '_get_updates()'.")
+
+    @abstractmethod
+    def apply(self, directions):
+        raise NotImplementedError("Subclass of 'DirectionModifier' must implement 'apply(directions)'.")
+
+
+class ParamModifier(object):
+    __metaclass__ = ABCMeta
+
+    @property
+    def tasks(self):
+        return []
+
+    @property
+    def updates(self):
+        return self._get_updates()
+
+    @abstractmethod
+    def _get_updates(self):
+        raise NotImplementedError("Subclass of 'ParamModifier' must implement '_get_updates()'.")
+
+    @abstractmethod
+    def apply(self, params):
+        raise NotImplementedError("Subclass of 'ParamModifier' must implement 'apply(params)'.")
diff --git a/smartlearner/interfaces/loss.py b/smartlearner/interfaces/loss.py
index e9823e9..cd1742f 100644
--- a/smartlearner/interfaces/loss.py
+++ b/smartlearner/interfaces/loss.py
@@ -1,24 +1,53 @@
+from collections import OrderedDict
+
 from theano import tensor as T
 
+from abc import ABCMeta, abstractmethod
+
 
 class Loss(object):
+    __metaclass__ = ABCMeta
+
     def __init__(self, model, dataset):
         self.model = model
         self.dataset = dataset
-        self.target = dataset.symb_targets
         self.consider_constant = []  # Part of the computational graph to be considered as a constant.
+        self._tasks = []
+        self._gradients = None
+
+        # Build the graph for the loss.
+        model_output = self.model.get_output(self.dataset.symb_inputs)
+        self._loss = self._compute_loss(model_output)
+
+    @abstractmethod
+    def _get_updates(self):
+        raise NotImplementedError("Subclass of 'Loss' must implement '_get_updates()'.")
 
-    def get_graph_output(self):
-        output, updates = self.model.get_model_output(self.dataset.symb_inputs)
-        return self._loss_function(output), updates
+    @abstractmethod
+    def _compute_loss(self, model_output):
+        raise NotImplementedError("Subclass of 'Loss' must implement '_compute_loss(model_output)'.")
 
-    def get_gradients(self):
-        cost, updates = self.get_graph_output()
-        gparams = T.grad(cost=cost,
+    @property
+    def gradients(self):
+        if self._gradients is None:
+            self._gradients = self._get_gradients()
+
+        return self._gradients
+
+    @property
+    def tasks(self):
+        return self.model.tasks + self._tasks
+
+    @property
+    def updates(self):
+        updates = OrderedDict()
+        updates.update(self.model.updates)
+        updates.update(self._get_updates())
+        return updates
+
+    def _get_gradients(self):
+        gparams = T.grad(cost=self._loss,
                          wrt=self.model.parameters,
                          consider_constant=self.consider_constant)
-        gradients = dict(zip(self.model.parameters, gparams))
-        return gradients, updates
-
-    def _loss_function(self, model_output):
-        raise NotImplementedError("Subclass of 'Loss' must implement '_loss_function(model_output)'.")
+        self._gradients = dict(zip(self.model.parameters, gparams))
+        return self.gradients
diff --git a/smartlearner/interfaces/model.py b/smartlearner/interfaces/model.py
index f72e79b..bad865b 100644
--- a/smartlearner/interfaces/model.py
+++ b/smartlearner/interfaces/model.py
@@ -1,5 +1,3 @@
-import theano.tensor as T
-from collections import OrderedDict
 from abc import ABCMeta, abstractmethod, abstractproperty
 
 
@@ -14,9 +12,18 @@ class abstractclassmethod(classmethod):
 class Model(object):
     __metaclass__ = ABCMeta
 
-    def get_model_output(self, inputs):
+    @property
+    def tasks(self):
+        return []
+
+    @abstractmethod
+    def get_output(self, inputs):
         raise NotImplementedError("Subclass of 'Model' must define a model output (a theano graph)")
 
+    @abstractproperty
+    def updates(self):
+        raise NotImplementedError("Subclass of 'Model' must implement property 'updates'.")
+
     @abstractproperty
     def parameters(self):
         raise NotImplementedError("Subclass of 'Model' must implement property 'parameters'.")
diff --git a/smartlearner/interfaces/optimizer.py b/smartlearner/interfaces/optimizer.py
new file mode 100644
index 0000000..13da9d1
--- /dev/null
+++ b/smartlearner/interfaces/optimizer.py
@@ -0,0 +1,81 @@
+from collections import OrderedDict
+
+from abc import ABCMeta, abstractmethod
+
+
+class Optimizer(object):
+    __metaclass__ = ABCMeta
+
+    def __init__(self, loss):
+        self.loss = loss
+        self._tasks = []
+
+        self._direction_modifiers = []
+        self._param_modifiers = []
+        self._directions = None
+
+    def append_direction_modifier(self, direction_modifier):
+        self._direction_modifiers.append(direction_modifier)
+
+    def append_param_modifier(self, param_modifier):
+        self._param_modifiers.append(param_modifier)
+
+    @abstractmethod
+    def _get_directions(self):
+        raise NotImplementedError("Subclass of 'Optimizer' must implement '_get_directions()'.")
+
+    @abstractmethod
+    def _get_updates(self):
+        raise NotImplementedError("Subclass of 'Optimizer' must implement private property '_updates'.")
+
+    @property
+    def directions(self):
+        if self._directions is None:
+            self._directions = self._get_directions()
+
+        return self._directions
+
+    @property
+    def tasks(self):
+        tasks = []
+        tasks.extend(self.loss.tasks)
+
+        for direction_modifier in self._direction_modifiers:
+            tasks.extend(direction_modifier.tasks)
+
+        for param_modifier in self._param_modifiers:
+            tasks.extend(param_modifier.tasks)
+
+        tasks.extend(self._tasks)
+        return tasks
+
+    @property
+    def updates(self):
+        updates = OrderedDict()
+
+        directions = self.directions
+        updates.update(self.loss.updates)  # Gather updates from the loss.
+        updates.update(self._get_updates())  # Gather updates from the optimizer.
+
+        # Apply directions modifiers and gather updates from these modifiers.
+        updates.update(self._apply_modifiers(self._direction_modifiers, directions))
+
+        # Update parameters
+        params_updates = OrderedDict()
+        for param, direction in directions.items():
+            params_updates[param] = param + direction
+        updates.update(params_updates)
+
+        # Apply parameters modifiers and gather updates from these modifiers.
+        updates.update(self._apply_modifiers(self._param_modifiers, params_updates))
+
+        return updates
+
+    def _apply_modifiers(self, list_modifiers, objects_to_modify):
+        updates = OrderedDict()
+        for modifier in list_modifiers:
+            modified_objects = modifier.apply(objects_to_modify)
+            objects_to_modify.update(modified_objects)
+            updates.update(modifier.updates)
+
+        return updates
diff --git a/smartlearner/losses/classification_losses.py b/smartlearner/losses/classification_losses.py
index b5bfc74..738a899 100644
--- a/smartlearner/losses/classification_losses.py
+++ b/smartlearner/losses/classification_losses.py
@@ -4,13 +4,19 @@ import theano.tensor as T
 
 
 class NegativeLogLikelihood(Loss):
-    def _loss_function(self, model_output):
+    def _get_updates(self):
+        return {}  # There is no updates for NegativeLogLikelihood.
+
+    def _compute_loss(self, model_output):
         nll = -T.log(model_output)
-        indices = T.cast(self.target[:, 0], dtype="int32")  # Targets are floats.
-        selected_nll = nll[T.arange(self.target.shape[0]), indices]
+        indices = T.cast(self.dataset.symb_targets[:, 0], dtype="int32")  # Targets are floats.
+        selected_nll = nll[T.arange(self.dataset.symb_targets.shape[0]), indices]
         return T.mean(selected_nll)
 
 
 class CategoricalCrossEntropy(Loss):
-    def _loss_function(self, model_output):
-        return T.mean(T.nnet.categorical_crossentropy(model_output, self.target))
+    def _get_updates(self):
+        return {}  # There is no updates for CategoricalCrossEntropy.
+
+    def _compute_loss(self, model_output):
+        return T.mean(T.nnet.categorical_crossentropy(model_output, self.dataset.symb_targets))
diff --git a/smartlearner/losses/distribution_losses.py b/smartlearner/losses/distribution_losses.py
index c20506b..97a99c7 100644
--- a/smartlearner/losses/distribution_losses.py
+++ b/smartlearner/losses/distribution_losses.py
@@ -4,5 +4,8 @@ import theano.tensor as T
 
 
 class BinaryCrossEntropy(Loss):
-    def _loss_function(self, model_output):
-        return T.mean(T.nnet.binary_crossentropy(model_output, self.target))
+    def _get_updates(self):
+        return {}  # There is no updates for BinaryCrossEntropy.
+
+    def _compute_loss(self, model_output):
+        return T.mean(T.nnet.binary_crossentropy(model_output, self.dataset.symb_targets))
diff --git a/smartlearner/losses/reconstruction_losses.py b/smartlearner/losses/reconstruction_losses.py
index 59f3e0e..19f5199 100644
--- a/smartlearner/losses/reconstruction_losses.py
+++ b/smartlearner/losses/reconstruction_losses.py
@@ -4,10 +4,16 @@ import theano.tensor as T
 
 
 class L2Distance(Loss):
-    def _loss_function(self, model_output):
-        return T.mean((model_output - self.target)**2)
+    def _get_updates(self):
+        return {}  # There is no updates for L2Distance.
+
+    def _compute_loss(self, model_output):
+        return T.mean((model_output - self.dataset.symb_targets)**2)
 
 
 class L1Distance(Loss):
-    def _loss_function(self, model_output):
-        return T.mean(abs(model_output - self.target))
+    def _get_updates(self):
+        return {}  # There is no updates for L1Distance.
+
+    def _compute_loss(self, model_output):
+        return T.mean(abs(model_output - self.dataset.symb_targets))
diff --git a/smartlearner/optimizers/__init__.py b/smartlearner/optimizers/__init__.py
index 3ddc3b4..d2baa34 100644
--- a/smartlearner/optimizers/__init__.py
+++ b/smartlearner/optimizers/__init__.py
@@ -1,2 +1,2 @@
-from .optimizer import Optimizer
 from .sgd import SGD
+from .adagrad import AdaGrad
diff --git a/smartlearner/optimizers/adagrad.py b/smartlearner/optimizers/adagrad.py
new file mode 100644
index 0000000..c132e9d
--- /dev/null
+++ b/smartlearner/optimizers/adagrad.py
@@ -0,0 +1,58 @@
+from collections import OrderedDict
+import theano.tensor as T
+
+from . import SGD
+from ..utils import sharedX
+
+
+class AdaGrad(SGD):
+    """ Implements the AdaGrad optimizer [Duchi11]_.
+
+    References
+    ----------
+    .. [Duchi11] Duchi, J., Hazan, E., & Singer, Y., "Adaptive Subgradient
+                 Methods for Online Learning and Stochastic Optimization",
+                 Journal of Machine Learning Research, vol. 12, pp. 2121-2159,
+                 2011.
+    """
+    def __init__(self, loss, lr, eps=1e-6):
+        """
+        Parameters
+        ----------
+        loss: `smartlearner.interfaces.loss.Loss` object
+            Loss function from which to obtain the gradients.
+        lr: float
+            Initial learning rate.
+        eps: float (optional)
+            Epsilon needed to avoid division by zero.
+        """
+        super().__init__(loss)
+        self.lr = lr
+        self.eps = eps
+        self.parameters = {}
+        self._updates = OrderedDict()
+
+    def _get_updates(self):
+        return self._updates
+
+    def _get_directions(self):
+        """ Produces descending directions. """
+        directions = OrderedDict()
+
+        for i, (param, direction) in enumerate(super()._get_directions().items()):
+            # sum_squared_grad := \sum g_t^2
+            param_name = param.name if param.name is not None else str(i)
+            sum_squared_grad = sharedX(param.get_value() * 0., name='sum_squared_grad_' + param_name)
+            self.parameters[sum_squared_grad.name] = sum_squared_grad
+
+            # Accumulate gradient
+            new_sum_squared_grad = sum_squared_grad + T.sqr(direction)
+
+            # Compute update
+            root_sum_squared = T.sqrt(new_sum_squared_grad + self.eps)
+
+            # Apply update
+            self._updates[sum_squared_grad] = new_sum_squared_grad
+            directions[param] = (self.lr/root_sum_squared) * direction
+
+        return directions
diff --git a/smartlearner/optimizers/optimizer.py b/smartlearner/optimizers/optimizer.py
deleted file mode 100644
index 954a241..0000000
--- a/smartlearner/optimizers/optimizer.py
+++ /dev/null
@@ -1,49 +0,0 @@
-from collections import OrderedDict
-
-from abc import ABCMeta, abstractmethod
-
-
-class Optimizer(object):
-    __metaclass__ = ABCMeta
-
-    def __init__(self, loss):
-        self.loss = loss
-
-        self._update_rules = []
-        self._param_modifiers = []
-
-    def append_update_rule(self, update_rule):
-        self._update_rules.append(update_rule)
-
-    def append_param_modifier(self, param_modifier):
-        self._param_modifiers.append(param_modifier)
-
-    @abstractmethod
-    def _get_directions(self):
-        raise NotImplementedError("Subclass of 'Optimizer' must implement '_get_directions()'.")
-
-    def gather_updates(self):
-        updates = OrderedDict()
-
-        self.directions, updates_from_get_directions = self._get_directions()
-        updates.update(updates_from_get_directions)
-
-        updates.update(self._apply_updates(self._update_rules, self.directions))
-
-        # Update parameters
-        params_updates = OrderedDict()
-        for param, gparam in self.directions.items():
-            params_updates[param] = param + self.directions[param]
-        updates.update(params_updates)
-
-        updates.update(self._apply_updates(self._param_modifiers, params_updates))
-
-        return updates
-
-    def _apply_updates(self, list_updates, object_to_update):
-        update_dict = OrderedDict()
-        for update in list_updates:
-            modified_object, updates_to_add = update.apply(object_to_update)
-            object_to_update.update(modified_object)
-            update_dict.update(updates_to_add)
-        return update_dict
\ No newline at end of file
diff --git a/smartlearner/optimizers/sgd.py b/smartlearner/optimizers/sgd.py
index 2130265..8f73365 100644
--- a/smartlearner/optimizers/sgd.py
+++ b/smartlearner/optimizers/sgd.py
@@ -1,16 +1,19 @@
-import numpy as np
-from . import Optimizer
+from collections import OrderedDict
+
+from ..interfaces.optimizer import Optimizer
 
 
 class SGD(Optimizer):
     def __init__(self, loss):
-        super(SGD, self).__init__(loss)
+        super().__init__(loss)
 
-    def _get_directions(self):
-        self.gradients, updates_from_get_gradients = self.loss.get_gradients()
+    def _get_updates(self):
+        return {}  # There is no updates for SGD.
 
-        # Take the opposite of the gradient.
-        for param, gparam in self.gradients.items():
-            self.gradients[param] = -gparam
+    def _get_directions(self):
+        # Take the opposite of the gradients as directions.
+        directions = OrderedDict()
+        for param, gradient in self.loss.gradients.items():
+            directions[param] = -gradient
 
-        return self.gradients, updates_from_get_gradients
+        return directions
diff --git a/smartlearner/tasks/tasks.py b/smartlearner/tasks/tasks.py
index cd6021f..75ba0c4 100644
--- a/smartlearner/tasks/tasks.py
+++ b/smartlearner/tasks/tasks.py
@@ -6,6 +6,16 @@ from time import time
 from smartlearner.interfaces.task import Task, RecurrentTask
 
 
+class MonitorVariable(Task):
+    def __init__(self, var):
+        super().__init__()
+        self.var = self.track_variable(var)
+
+    @property
+    def value(self):
+        return self.var.get_value()
+
+
 class PrintVariable(RecurrentTask):
     def __init__(self, msg, *variables, **recurrent_options):
         # TODO: docstring should include **recurrent_options.
diff --git a/smartlearner/trainer.py b/smartlearner/trainer.py
index c25770f..d9cfda0 100644
--- a/smartlearner/trainer.py
+++ b/smartlearner/trainer.py
@@ -1,9 +1,8 @@
 from collections import OrderedDict
 
 import theano
-from time import time
 from .status import Status
-from smartlearner.tasks.stopping_criteria import TrainingExit
+from .tasks.stopping_criteria import TrainingExit
 
 
 class Trainer(object):
@@ -11,8 +10,16 @@ class Trainer(object):
         self.status = status if status is not None else Status(self)
         self._optimizer = optimizer
         self._batch_scheduler = batch_scheduler
-        self._updates = OrderedDict()
+
+        # Gather updates from the optimizer and the batch scheduler.
+        self._graph_updates = OrderedDict()
+        self._graph_updates.update(self._optimizer.updates)
+        self._graph_updates.update(self._batch_scheduler.updates)
+
+        # Gather tasks from the optimizer and the batch scheduler.
         self._tasks = []
+        self._tasks.extend(self._optimizer.tasks)
+        self._tasks.extend(self._batch_scheduler.tasks)
 
     def train(self):
         self._pre_learning()
@@ -20,16 +27,17 @@ class Trainer(object):
         self._post_learning()
 
     def append_task(self, task):
-        self._updates.update(task.updates)
         self._tasks.append(task)
 
     def _build_theano_graph(self):
-        updates = self._optimizer.gather_updates()
-        updates.update(self._updates)
+        # Get updates from tasks.
+        for task in self._tasks:
+            self._graph_updates.update(task.updates)
+
         self._learn = theano.function([],
-                       updates=updates,
-                       givens=self._batch_scheduler.givens,
-                       name="learn")
+                                      updates=self._graph_updates,
+                                      givens=self._batch_scheduler.givens,
+                                      name="learn")
 
     def _pre_learning(self):
         self._build_theano_graph()
diff --git a/smartlearner/update_rules/decreasing_learning_rate.py b/smartlearner/update_rules/decreasing_learning_rate.py
deleted file mode 100644
index f721726..0000000
--- a/smartlearner/update_rules/decreasing_learning_rate.py
+++ /dev/null
@@ -1,40 +0,0 @@
-from collections import OrderedDict
-
-import numpy as np
-
-from ..utils import sharedX
-from . import UpdateRule
-
-
-class DecreasingLearningRate(UpdateRule):
-    def __init__(self, lr, dc=0.):
-        """
-        Implements a decreasing learning rate update rule.
-
-        Parameters
-        ----------
-        lr: float
-            learning rate
-        dc: float
-            decreasing constant (decay)
-        """
-        super(DecreasingLearningRate, self).__init__()
-        assert dc <= 1.
-        assert dc >= 0.
-        self.lr = lr
-        self.dc = dc
-
-    def apply(self, gradients):
-        updates = OrderedDict()
-        new_gradients = OrderedDict()
-
-        for param, gparam in gradients.items():
-            lr = sharedX(self.lr * np.ones_like(param.get_value()), name='lr_' + param.name)
-
-            if self.dc != 0.:
-                # Decrease the learning rate by a factor of `dc` after each update.
-                updates[lr] = self.dc * lr
-
-            new_gradients[param] = lr * gparam
-
-        return new_gradients, updates
diff --git a/smartlearner/update_rules/update_rule.py b/smartlearner/update_rules/update_rule.py
deleted file mode 100644
index e5933ad..0000000
--- a/smartlearner/update_rules/update_rule.py
+++ /dev/null
@@ -1,9 +0,0 @@
-from abc import ABCMeta, abstractmethod
-
-
-class UpdateRule(object):
-    __metaclass__ = ABCMeta
-
-    @abstractmethod
-    def apply(self, gradients):
-        raise NotImplementedError("Subclass of 'UpdateRule' must implement 'apply(gradients)'.")

Refactor how the updates are handled
I suggest a global updates gatherer (either a module function or a singleton class if you feel fancy) that each parts would have access to. This way, everyone would be in charge of registering its updates.

Another solution is a chain call to gather the updates. This would have to be added in the interfaces of each abstract classes (`Model`, `Optimizer`, `Loss`?, `Tasks`?).
diff --git a/smartlearner/optimizers/tests/tests_optimizers.py b/smartlearner/optimizers/tests/tests_optimizers.py
new file mode 100644
index 0000000..7735fb9
--- /dev/null
+++ b/smartlearner/optimizers/tests/tests_optimizers.py
@@ -0,0 +1,110 @@
+import numpy as np
+
+import theano
+import theano.tensor as T
+
+from smartlearner.tasks import stopping_criteria
+from smartlearner.utils import sharedX
+
+from smartlearner import Trainer
+from smartlearner.optimizers import SGD, AdaGrad
+
+from smartlearner.tasks import tasks
+
+from numpy.testing import assert_array_almost_equal
+
+from smartlearner.testing import DummyLoss, DummyBatchScheduler
+
+floatX = theano.config.floatX
+
+
+class DummyLossWithGradient(DummyLoss):
+    def __init__(self, cost, param):
+        super().__init__()
+        self.cost = cost
+        self.param = param
+
+    def _get_gradients(self):
+        gparam = T.grad(cost=self.cost, wrt=self.param)
+        return {self.param: gparam}
+
+
+def test_sgd():
+    # Create simple Nd gaussian functions to optimize. These functions are
+    # (perfectly) well-conditioned so it should take only one gradient step
+    # to converge using 1/L, where L is the largest eigenvalue of the hessian.
+    max_epoch = 2
+    for N in range(1, 5):
+        center = np.arange(1, N+1)[None, :].astype(floatX)
+        param = sharedX(np.zeros((1, N)))
+        cost = T.sum(0.5*T.dot(T.dot((param-center), T.eye(N)), (param-center).T))
+        loss = DummyLossWithGradient(cost, param)
+
+        trainer = Trainer(SGD(loss), DummyBatchScheduler())
+        trainer.append_task(stopping_criteria.MaxEpochStopping(max_epoch))
+
+        # Monitor the gradient of `loss` w.r.t. to `param`.
+        gparam = tasks.MonitorVariable(loss.gradients[param])
+        trainer.append_task(gparam)
+        trainer.train()
+
+        # Since the problem is well-conditionned and we use an optimal gradient step 1/L,
+        # two epochs should be enough for `param` to be around `center` and the gradients near 0.
+        assert_array_almost_equal(param.get_value(), center)
+        assert_array_almost_equal(gparam.value, 0.)
+
+    # Create an Nd gaussian function to optimize. This function is not
+    # well-conditioned and there exists no perfect gradient step to converge in
+    # only one iteration.
+    #cost = T.sum(N*0.5*T.dot(T.dot((param-center), np.diag(1./np.arange(1, N+1))), ((param-center).T)))
+    max_epoch = 80
+    N = 4
+    center = 5*np.ones((1, N)).astype(floatX)
+    param = sharedX(np.zeros((1, N)))
+    cost = T.sum(0.5*T.dot(T.dot((param-center), np.diag(1./np.arange(1, N+1))), (param-center).T))
+    loss = DummyLossWithGradient(cost, param)
+
+    trainer = Trainer(SGD(loss), DummyBatchScheduler())
+    trainer.append_task(stopping_criteria.MaxEpochStopping(max_epoch))
+    #trainer.append_task(tasks.PrintVariable("Loss param   : {}", param))
+    #trainer.append_task(tasks.PrintVariable("Loss gradient: {}", loss.gradients[param]))
+
+    # Monitor the gradient of `loss` w.r.t. to `param`.
+    gparam = tasks.MonitorVariable(loss.gradients[param])
+    trainer.append_task(gparam)
+    trainer.train()
+
+    # Since the problem is well-conditionned and we use an optimal gradient step 1/L,
+    # two epochs should be enough for `param` to be around `center` and the gradients near 0.
+    assert_array_almost_equal(param.get_value(), center, decimal=6)
+    assert_array_almost_equal(gparam.value, 0.)
+
+
+def test_adagrad():
+    max_epoch = 15
+
+    # Create an Nd gaussian functions to optimize. These functions are not
+    # well-conditioned and there exists no perfect gradient step to converge in
+    # only one iteration.
+    for N in range(1, 5):
+        center = 5*np.ones((1, N)).astype(floatX)
+        param = sharedX(np.zeros((1, N)))
+        cost = T.sum(0.5*T.dot(T.dot((param-center), np.diag(1./np.arange(1, N+1))), ((param-center).T)))
+        loss = DummyLossWithGradient(cost, param)
+
+        # Even with a really high gradient step, AdaGrad can still converge.
+        # Actually, it is faster than using the optimal gradient step with SGD.
+        optimizer = AdaGrad(loss, lr=100, eps=1e-1)
+        trainer = Trainer(optimizer, DummyBatchScheduler())
+        trainer.append_task(stopping_criteria.MaxEpochStopping(max_epoch))
+        #trainer.append_task(tasks.PrintVariable("Loss param   : {}", param))
+        #trainer.append_task(tasks.PrintVariable("Loss gradient: {}", loss.gradients[param]))
+
+        # Monitor the gradient of `loss` w.r.t. to `param`.
+        gparam = tasks.MonitorVariable(loss.gradients[param])
+        trainer.append_task(gparam)
+        trainer.train()
+
+        # After 30 epochs, param should be around the center and gradients near 0.
+        assert_array_almost_equal(param.get_value(), center)
+        assert_array_almost_equal(gparam.value, 0.)
diff --git a/smartlearner/testing.py b/smartlearner/testing.py
index 4036c3a..5edf1ce 100644
--- a/smartlearner/testing.py
+++ b/smartlearner/testing.py
@@ -1,10 +1,10 @@
 import numpy as np
 
-from smartlearner.interfaces.dataset import Dataset
-from smartlearner.interfaces.model import Model
-from smartlearner.interfaces.loss import Loss
-from smartlearner.optimizers.optimizer import Optimizer
-from smartlearner.batch_scheduler import BatchScheduler
+from .interfaces.dataset import Dataset
+from .interfaces.model import Model
+from .interfaces.loss import Loss
+from .interfaces.optimizer import Optimizer
+from .interfaces.batch_scheduler import BatchScheduler
 
 
 class DummyDataset(Dataset):
@@ -21,8 +21,12 @@ class DummyModel(Model):
     def parameters(self):
         return self._parameters
 
-    def get_model_output(self, inputs):
-        pass
+    @property
+    def updates(self):
+        return {}
+
+    def get_output(self, inputs):
+        return inputs
 
     def save(self, path):
         pass
@@ -35,16 +39,22 @@ class DummyLoss(Loss):
     def __init__(self):
         super(DummyLoss, self).__init__(DummyModel(), DummyDataset())
 
-    def _loss_function(self, model_output):
-        pass
+    def _compute_loss(self, model_output):
+        return model_output
+
+    def _get_updates(self):
+        return {}
 
 
 class DummyOptimizer(Optimizer):
     def __init__(self):
         super(DummyOptimizer, self).__init__(loss=DummyLoss())
 
+    def _get_updates(self):
+        return {}
+
     def _get_directions(self):
-        return {}, {}
+        return {}
 
 
 class DummyBatchScheduler(BatchScheduler):
@@ -55,5 +65,9 @@ class DummyBatchScheduler(BatchScheduler):
     def givens(self):
         return {}
 
+    @property
+    def updates(self):
+        return {}
+
     def __iter__(self):
         return iter(range(1))

pip install -e .
parse_log_pytest
theano numpy scipy nose pyparsing pip flake8 six pep8 pyflakes
nose nose-cov pytest
3.9
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
cov-core==1.15.0
coverage==7.8.0
exceptiongroup==1.2.2
flake8 @ file:///croot/flake8_1726157165993/work
iniconfig==2.1.0
Mako @ file:///croot/mako_1665472421453/work
MarkupSafe @ file:///croot/markupsafe_1738584038848/work
mccabe @ file:///opt/conda/conda-bld/mccabe_1644221741721/work
nose @ file:///opt/conda/conda-bld/nose_1642704612149/work
nose-cov==1.6
numpy @ file:///opt/conda/conda-bld/numpy_and_numpy_base_1652801679809/work
packaging==24.2
pep8==1.7.1
pluggy==1.5.0
pycodestyle @ file:///croot/pycodestyle_1726150303809/work
pyflakes @ file:///croot/pyflakes_1708962956225/work
pygpu==0.7.6
pyparsing @ file:///croot/pyparsing_1731445506121/work
pytest==8.3.5
scipy @ file:///opt/conda/conda-bld/scipy_1661390393401/work
six @ file:///tmp/build/80754af9/six_1644875935023/work
-e git+https://github.com/SMART-Lab/smartlearner.git@0f11484800712ac92d6027754ab0a026193fa985#egg=smartlearner
Theano==1.0.5
tomli==2.2.1

name: smartlearner
channels:
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - _sysroot_linux-64_curr_repodata_hack=3=haa98f57_10
  - binutils_impl_linux-64=2.35.1=h27ae35d_9
  - binutils_linux-64=2.35.1=h454624a_30
  - blas=1.0=openblas
  - ca-certificates=2025.2.25=h06a4308_0
  - fftw=3.3.9=h5eee18b_2
  - flake8=7.1.1=py39h06a4308_0
  - gcc_impl_linux-64=7.5.0=h7105cf2_17
  - gcc_linux-64=7.5.0=h8f34230_30
  - gxx_impl_linux-64=7.5.0=h0a5bf11_17
  - gxx_linux-64=7.5.0=hffc177d_30
  - kernel-headers_linux-64=3.10.0=h57e8cba_10
  - ld_impl_linux-64=2.35.1=h7274673_9
  - libffi=3.4.4=h6a678d5_1
  - libgcc-devel_linux-64=7.5.0=hbbeae57_17
  - libgcc-ng=11.2.0=h1234567_1
  - libgfortran-ng=11.2.0=h00389a5_1
  - libgfortran5=11.2.0=h1234567_1
  - libgomp=11.2.0=h1234567_1
  - libgpuarray=0.7.6=h7f8727e_1
  - libopenblas=0.3.21=h043d6bf_0
  - libstdcxx-devel_linux-64=7.5.0=hf0c5c8d_17
  - libstdcxx-ng=11.2.0=h1234567_1
  - mako=1.2.3=py39h06a4308_0
  - markupsafe=3.0.2=py39h5eee18b_0
  - mccabe=0.7.0=pyhd3eb1b0_0
  - ncurses=6.4=h6a678d5_0
  - nose=1.3.7=pyhd3eb1b0_1008
  - numpy=1.22.3=py39h7a5d4dd_0
  - numpy-base=1.22.3=py39hb8be1f0_0
  - openssl=3.0.16=h5eee18b_0
  - pep8=1.7.1=py39h06a4308_1
  - pip=25.0=py39h06a4308_0
  - pycodestyle=2.12.1=py39h06a4308_0
  - pyflakes=3.2.0=py39h06a4308_0
  - pygpu=0.7.6=py39hce1f21e_1
  - pyparsing=3.2.0=py39h06a4308_0
  - python=3.9.21=he870216_1
  - readline=8.2=h5eee18b_0
  - scipy=1.7.3=py39hf838250_2
  - setuptools=75.8.0=py39h06a4308_0
  - six=1.16.0=pyhd3eb1b0_1
  - sqlite=3.45.3=h5eee18b_0
  - sysroot_linux-64=2.17=h57e8cba_10
  - theano=1.0.5=py39h295c915_1
  - tk=8.6.14=h39e8969_0
  - tzdata=2025a=h04d1e81_0
  - wheel=0.45.1=py39h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - pip:
      - cov-core==1.15.0
      - coverage==7.8.0
      - exceptiongroup==1.2.2
      - iniconfig==2.1.0
      - nose-cov==1.6
      - packaging==24.2
      - pluggy==1.5.0
      - pytest==8.3.5
      - tomli==2.2.1
prefix: /opt/conda/envs/smartlearner

smartlearner/optimizers/tests/tests_optimizers.py::test_sgd smartlearner/optimizers/tests/tests_optimizers.py::test_adagrad
smartlearner/batch_scheduler.py smartlearner/interfaces/optimizer.py smartlearner/batch_schedulers/__init__.py smartlearner/update_rules/constant_learning_rate.py smartlearner/update_rules/update_rule.py smartlearner/tasks/tasks.py smartlearner/interfaces/__init__.py smartlearner/direction_modifiers/decreasing_learning_rate.py smartlearner/optimizers/__init__.py smartlearner/update_rules/decreasing_learning_rate.py smartlearner/losses/distribution_losses.py smartlearner/optimizers/sgd.py smartlearner/interfaces/direction_modifier.py smartlearner/interfaces/loss.py smartlearner/__init__.py smartlearner/optimizers/adagrad.py smartlearner/interfaces/model.py smartlearner/update_rules/__init__.py smartlearner/trainer.py smartlearner/optimizers/optimizer.py smartlearner/losses/classification_losses.py smartlearner/batch_schedulers/minibatch.py smartlearner/losses/reconstruction_losses.py
smartlearner/interfaces/optimizer.py smartlearner/batch_schedulers/__init__.py smartlearner/update_rules/update_rule.py smartlearner/tasks/tasks.py smartlearner/interfaces/__init__.py smartlearner/direction_modifiers/decreasing_learning_rate.py smartlearner/optimizers/__init__.py smartlearner/update_rules/decreasing_learning_rate.py smartlearner/losses/distribution_losses.py smartlearner/optimizers/sgd.py smartlearner/interfaces/batch_scheduler.py smartlearner/interfaces/direction_modifier.py smartlearner/direction_modifiers/__init__.py smartlearner/interfaces/loss.py smartlearner/__init__.py smartlearner/optimizers/adagrad.py smartlearner/interfaces/model.py smartlearner/trainer.py smartlearner/optimizers/optimizer.py smartlearner/losses/classification_losses.py smartlearner/batch_schedulers/minibatch.py smartlearner/losses/reconstruction_losses.py smartlearner/direction_modifiers/constant_learning_rate.py