extend multiobjective docs

Author: stes

cebra/attribution/jacobian_attribution.py

@@ -72,10 +72,9 @@ def get_attribution_map(
         hybrid_solver: If ``True``, handle multi-objective models differently.
 
     Returns:
-        A tuple containing:
-            - jf: The Jacobian matrix of shape (num_samples, output_dim, input_dim)
-            - jhatg: The pseudo-inverse of the Jacobian matrix
-        The result is effectively a ``(num_inputs, num_features)`` attribution map.
+        A tuple containing the Jacobian matrix of shape (num_samples, output_dim, input_dim)
+        and the pseudo-inverse of the Jacobian matrix.
+
     """
     assert aggregate in ["mean", "sum", "max"]
 

cebra/solver/multiobjective.py

@@ -19,7 +19,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-"""Multiobjective contrastive learning."""
+"""Multiobjective contrastive learning.
+
+Starting in CEBRA 0.6.0, we have added support for subspace contrastive learning.
+This is a method for training models that are able to learn multiple subspaces of the
+feature space simultaneously.
+
+Subspace contrastive learning requires to use specialized models and criterions.
+This module specifies a test of classes required for training CEBRA models with multiple objectives.
+The objectives are defined by the wrapper class :py:class:`cebra.models.multicriterions.MultiCriterions`.
+
+Two solvers are currently implemented:
+
+- :py:class:`cebra.solver.multiobjective.ContrastiveMultiobjectiveSolverxCEBRA`
+- :py:class:`cebra.solver.multiobjective.SupervisedMultiobjectiveSolverxCEBRA`
+
+See Also:
+    :py:class:`cebra.solver.multiobjective.SupervisedMultiobjectiveSolverxCEBRA`
+    :py:class:`cebra.solver.multiobjective.MultiObjectiveConfig`
+    :py:class:`cebra.models.multicriterions.MultiCriterions`
+"""
 
 import logging
 import time
@@ -43,6 +62,8 @@
 class MultiObjectiveConfig:
     """Configuration class for setting up multi-objective learning with Cebra.
 
+
+
     Args:
         loader: Data loader used for configurations.
     """
@@ -458,7 +479,11 @@ def transform(self, inputs: torch.Tensor) -> torch.Tensor:
 @register("supervised-solver-xcebra")
 @dataclasses.dataclass
 class SupervisedMultiobjectiveSolverxCEBRA(MultiobjectiveSolverBase):
-    """Supervised neural network training with MSE loss"""
+    """Supervised neural network training using the MSE loss.
+
+    This solver can be used as a baseline variant instead of the contrastive solver,
+    :py:class:`cebra.solver.multiobjective.ContrastiveMultiobjectiveSolverxCEBRA`.
+    """
 
     _variant_name = "supervised-solver-xcebra"
 
@@ -477,6 +502,15 @@ def _inference(self, batch):
 @register("multiobjective-solver")
 @dataclasses.dataclass
 class ContrastiveMultiobjectiveSolverxCEBRA(MultiobjectiveSolverBase):
+    """Multi-objective solver for CEBRA.
+
+    This solver is used for training CEBRA models with multiple objectives.
+
+    See Also:
+        :py:class:`cebra.solver.multiobjective.SupervisedMultiobjectiveSolverxCEBRA`
+        :py:class:`cebra.solver.multiobjective.MultiObjectiveConfig`
+        :py:class:`cebra.models.multicriterions.MultiCriterions`
+    """
 
     _variant_name = "contrastive-solver-xcebra"
 

docs/source/api/xcebra/multiobjective.rst

@@ -2,20 +2,14 @@
 Multi-objective models
 ======================
 
-Starting in CEBRA 0.6.0, we have added support for subspace contrastive learning.
-This is a method for training models that are able to learn multiple subspaces of the
-feature space simultaneously.
-
-Subspace contrastive learning requires to use specialized models and criterions.
-
-.. automodule:: cebra.models.multicriterions
+.. automodule:: cebra.solver.multiobjective
    :members:
    :show-inheritance:
 
-.. automodule:: cebra.models.multiobjective
+.. automodule:: cebra.models.multicriterions
    :members:
    :show-inheritance:
 
-.. automodule:: cebra.solver.multiobjective
+.. automodule:: cebra.models.multiobjective
    :members:
    :show-inheritance: