sklears_multioutput/
multitask.rs

1//! Multi-Task Neural Networks with Shared Representation Learning
2//!
3//! This module implements multi-task learning where multiple related tasks share
4//! common representations in lower layers while having task-specific layers for final predictions.
5//! This approach allows for better generalization and improved performance when tasks are related.
6#![allow(non_snake_case)] // Standard ML notation: X for feature matrices, K for kernels
7
8// Use SciRS2-Core for arrays and random number generation (SciRS2 Policy)
9use scirs2_core::ndarray::{Array1, Array2, ArrayView2};
10use scirs2_core::random::thread_rng;
11use sklears_core::{
12    error::{Result as SklResult, SklearsError},
13    traits::{Estimator, Fit, Predict, Untrained},
14    types::Float,
15};
16use std::collections::HashMap;
17
18use crate::activation::ActivationFunction;
19use crate::loss::LossFunction;
20
21/// Task balancing strategies for multi-task learning
22#[derive(Debug, Clone, PartialEq)]
23pub enum TaskBalancing {
24    /// Equal weights for all tasks
25    Equal,
26    /// Custom weights for each task
27    Weighted,
28    /// Adaptive weighting based on task difficulty
29    Adaptive,
30    /// Gradient balancing
31    GradientBalancing,
32}
33
34/// Multi-Task Neural Network with Shared Representation Learning
35///
36/// This neural network implements multi-task learning where multiple related tasks share
37/// common representations in lower layers while having task-specific layers for final predictions.
38/// This approach allows for better generalization and improved performance when tasks are related.
39///
40/// # Architecture
41///
42/// The network consists of:
43/// - Shared layers: Learn common representations across all tasks
44/// - Task-specific layers: Learn task-specific transformations
45/// - Multiple outputs: One output per task
46///
47/// # Examples
48///
49/// ```
50/// use sklears_multioutput::multitask::{MultiTaskNeuralNetwork, TaskBalancing};
51/// use sklears_multioutput::activation::ActivationFunction;
52/// use sklears_core::traits::{Predict, Fit};
53/// // Use SciRS2-Core for arrays and random number generation (SciRS2 Policy)
54/// use scirs2_core::ndarray::array;
55/// use std::collections::HashMap;
56///
57/// let X = array![[1.0, 2.0], [2.0, 3.0], [3.0, 1.0], [4.0, 4.0]];
58/// let mut tasks = HashMap::new();
59/// tasks.insert("task1".to_string(), array![[0.5], [1.0], [1.5], [2.0]]); // Regression task
60/// tasks.insert("task2".to_string(), array![[1.0], [0.0], [1.0], [0.0]]); // Classification task
61///
62/// let mt_net = MultiTaskNeuralNetwork::new()
63///     .shared_layers(vec![20, 10])
64///     .task_specific_layers(vec![5])
65///     .task_outputs(&[("task1", 1), ("task2", 1)])
66///     .shared_activation(ActivationFunction::ReLU)
67///     .learning_rate(0.01)
68///     .max_iter(1000)
69///     .task_weights(&[("task1", 1.0), ("task2", 0.8)])
70///     .random_state(Some(42));
71/// ```
72#[derive(Debug, Clone)]
73pub struct MultiTaskNeuralNetwork<S = Untrained> {
74    state: S,
75    /// Sizes of shared representation layers
76    shared_layer_sizes: Vec<usize>,
77    /// Sizes of task-specific layers
78    task_specific_layer_sizes: Vec<usize>,
79    /// Task names and their output dimensions
80    task_outputs: HashMap<String, usize>,
81    /// Task loss functions
82    task_loss_functions: HashMap<String, LossFunction>,
83    /// Task weights for multi-task loss computation
84    task_weights: HashMap<String, Float>,
85    /// Activation function for shared layers
86    shared_activation: ActivationFunction,
87    /// Activation function for task-specific layers
88    task_activation: ActivationFunction,
89    /// Output activation functions per task
90    output_activations: HashMap<String, ActivationFunction>,
91    /// Learning rate
92    learning_rate: Float,
93    /// Maximum number of iterations
94    max_iter: usize,
95    /// Convergence tolerance
96    tolerance: Float,
97    /// Random state for reproducibility
98    random_state: Option<u64>,
99    /// L2 regularization strength
100    alpha: Float,
101    /// Batch size for training
102    batch_size: Option<usize>,
103    /// Early stopping
104    early_stopping: bool,
105    /// Validation fraction for early stopping
106    validation_fraction: Float,
107    /// Task balancing strategy
108    task_balancing: TaskBalancing,
109}
110
111/// Trained state for MultiTaskNeuralNetwork
112#[derive(Debug, Clone)]
113pub struct MultiTaskNeuralNetworkTrained {
114    #[allow(dead_code)]
115    /// Weights for shared layers
116    shared_weights: Vec<Array2<Float>>,
117    #[allow(dead_code)]
118    /// Biases for shared layers
119    shared_biases: Vec<Array1<Float>>,
120    #[allow(dead_code)]
121    /// Task-specific weights per task
122    task_weights: HashMap<String, Vec<Array2<Float>>>,
123    #[allow(dead_code)]
124    /// Task-specific biases per task
125    task_biases: HashMap<String, Vec<Array1<Float>>>,
126    #[allow(dead_code)]
127    /// Output layer weights per task
128    output_weights: HashMap<String, Array2<Float>>,
129    #[allow(dead_code)]
130    /// Output layer biases per task
131    output_biases: HashMap<String, Array1<Float>>,
132    /// Number of input features
133    n_features: usize,
134    /// Task configurations
135    task_outputs: HashMap<String, usize>,
136    #[allow(dead_code)]
137    /// Network architecture
138    shared_layer_sizes: Vec<usize>,
139    #[allow(dead_code)]
140    task_specific_layer_sizes: Vec<usize>,
141    #[allow(dead_code)]
142    shared_activation: ActivationFunction,
143    #[allow(dead_code)]
144    task_activation: ActivationFunction,
145    #[allow(dead_code)]
146    output_activations: HashMap<String, ActivationFunction>,
147    /// Training history per task
148    task_loss_curves: HashMap<String, Vec<Float>>,
149    /// Combined loss curve
150    combined_loss_curve: Vec<Float>,
151    /// Number of iterations performed
152    n_iter: usize,
153}
154
155impl MultiTaskNeuralNetwork<Untrained> {
156    /// Create a new MultiTaskNeuralNetwork
157    pub fn new() -> Self {
158        Self {
159            state: Untrained,
160            shared_layer_sizes: vec![100],
161            task_specific_layer_sizes: vec![50],
162            task_outputs: HashMap::new(),
163            task_loss_functions: HashMap::new(),
164            task_weights: HashMap::new(),
165            shared_activation: ActivationFunction::ReLU,
166            task_activation: ActivationFunction::ReLU,
167            output_activations: HashMap::new(),
168            learning_rate: 0.001,
169            max_iter: 1000,
170            tolerance: 1e-6,
171            random_state: None,
172            alpha: 0.0001,
173            batch_size: None,
174            early_stopping: false,
175            validation_fraction: 0.1,
176            task_balancing: TaskBalancing::Equal,
177        }
178    }
179
180    /// Set the sizes of shared representation layers
181    pub fn shared_layers(mut self, sizes: Vec<usize>) -> Self {
182        self.shared_layer_sizes = sizes;
183        self
184    }
185
186    /// Set the sizes of task-specific layers
187    pub fn task_specific_layers(mut self, sizes: Vec<usize>) -> Self {
188        self.task_specific_layer_sizes = sizes;
189        self
190    }
191
192    /// Configure task outputs
193    pub fn task_outputs(mut self, tasks: &[(&str, usize)]) -> Self {
194        for (task_name, output_size) in tasks {
195            self.task_outputs
196                .insert(task_name.to_string(), *output_size);
197            // Set default configurations
198            self.task_loss_functions.insert(
199                task_name.to_string(),
200                if *output_size == 1 {
201                    LossFunction::MeanSquaredError
202                } else {
203                    LossFunction::CrossEntropy
204                },
205            );
206            self.task_weights.insert(task_name.to_string(), 1.0);
207            self.output_activations.insert(
208                task_name.to_string(),
209                if *output_size == 1 {
210                    ActivationFunction::Linear
211                } else {
212                    ActivationFunction::Softmax
213                },
214            );
215        }
216        self
217    }
218
219    /// Set loss functions for specific tasks
220    pub fn task_loss_functions(mut self, loss_functions: &[(&str, LossFunction)]) -> Self {
221        for (task_name, loss_fn) in loss_functions {
222            self.task_loss_functions
223                .insert(task_name.to_string(), *loss_fn);
224        }
225        self
226    }
227
228    /// Set task weights for multi-task loss computation
229    pub fn task_weights(mut self, weights: &[(&str, Float)]) -> Self {
230        for (task_name, weight) in weights {
231            self.task_weights.insert(task_name.to_string(), *weight);
232        }
233        self
234    }
235
236    /// Set activation function for shared layers
237    pub fn shared_activation(mut self, activation: ActivationFunction) -> Self {
238        self.shared_activation = activation;
239        self
240    }
241
242    /// Set activation function for task-specific layers
243    pub fn task_activation(mut self, activation: ActivationFunction) -> Self {
244        self.task_activation = activation;
245        self
246    }
247
248    /// Set output activation functions for specific tasks
249    pub fn output_activations(mut self, activations: &[(&str, ActivationFunction)]) -> Self {
250        for (task_name, activation) in activations {
251            self.output_activations
252                .insert(task_name.to_string(), *activation);
253        }
254        self
255    }
256
257    /// Set learning rate
258    pub fn learning_rate(mut self, lr: Float) -> Self {
259        self.learning_rate = lr;
260        self
261    }
262
263    /// Set maximum number of iterations
264    pub fn max_iter(mut self, max_iter: usize) -> Self {
265        self.max_iter = max_iter;
266        self
267    }
268
269    /// Set convergence tolerance
270    pub fn tolerance(mut self, tolerance: Float) -> Self {
271        self.tolerance = tolerance;
272        self
273    }
274
275    /// Set random state for reproducibility
276    pub fn random_state(mut self, seed: Option<u64>) -> Self {
277        self.random_state = seed;
278        self
279    }
280
281    /// Set L2 regularization strength
282    pub fn alpha(mut self, alpha: Float) -> Self {
283        self.alpha = alpha;
284        self
285    }
286
287    /// Set batch size for training
288    pub fn batch_size(mut self, batch_size: Option<usize>) -> Self {
289        self.batch_size = batch_size;
290        self
291    }
292
293    /// Enable/disable early stopping
294    pub fn early_stopping(mut self, early_stopping: bool) -> Self {
295        self.early_stopping = early_stopping;
296        self
297    }
298
299    /// Set validation fraction for early stopping
300    pub fn validation_fraction(mut self, fraction: Float) -> Self {
301        self.validation_fraction = fraction;
302        self
303    }
304
305    /// Set task balancing strategy
306    pub fn task_balancing(mut self, strategy: TaskBalancing) -> Self {
307        self.task_balancing = strategy;
308        self
309    }
310}
311
312impl Default for MultiTaskNeuralNetwork<Untrained> {
313    fn default() -> Self {
314        Self::new()
315    }
316}
317
318impl Estimator for MultiTaskNeuralNetwork<Untrained> {
319    type Config = ();
320    type Error = SklearsError;
321    type Float = Float;
322
323    fn config(&self) -> &Self::Config {
324        &()
325    }
326}
327
328// Implementation of Fit trait with simplified training logic
329impl Fit<ArrayView2<'_, Float>, HashMap<String, Array2<Float>>>
330    for MultiTaskNeuralNetwork<Untrained>
331{
332    type Fitted = MultiTaskNeuralNetwork<MultiTaskNeuralNetworkTrained>;
333
334    fn fit(
335        self,
336        x: &ArrayView2<Float>,
337        y: &HashMap<String, Array2<Float>>,
338    ) -> SklResult<Self::Fitted> {
339        if x.nrows() == 0 || x.ncols() == 0 {
340            return Err(SklearsError::InvalidInput("Empty input data".to_string()));
341        }
342
343        if y.is_empty() {
344            return Err(SklearsError::InvalidInput("No tasks provided".to_string()));
345        }
346
347        // Validate that all tasks have consistent sample counts
348        let n_samples = x.nrows();
349        for (task_name, task_targets) in y {
350            if task_targets.nrows() != n_samples {
351                return Err(SklearsError::ShapeMismatch {
352                    expected: format!("{}", n_samples),
353                    actual: format!("{}", task_targets.nrows()),
354                });
355            }
356            if !self.task_outputs.contains_key(task_name) {
357                return Err(SklearsError::InvalidInput(format!(
358                    "Unknown task: {}",
359                    task_name
360                )));
361            }
362        }
363
364        let n_features = x.ncols();
365        let _rng = thread_rng();
366
367        // Initialize network parameters (simplified)
368        let shared_weights = vec![Array2::<Float>::zeros((n_features, 50))];
369        let shared_biases = vec![Array1::<Float>::zeros(50)];
370        let mut task_weights = HashMap::new();
371        let mut task_biases = HashMap::new();
372        let mut output_weights = HashMap::new();
373        let mut output_biases = HashMap::new();
374
375        for (task_name, &output_size) in &self.task_outputs {
376            task_weights.insert(task_name.clone(), vec![Array2::<Float>::zeros((50, 25))]);
377            task_biases.insert(task_name.clone(), vec![Array1::<Float>::zeros(25)]);
378            output_weights.insert(task_name.clone(), Array2::<Float>::zeros((25, output_size)));
379            output_biases.insert(task_name.clone(), Array1::<Float>::zeros(output_size));
380        }
381
382        // Simplified training loop
383        let mut task_loss_curves = HashMap::new();
384        let combined_loss_curve = vec![0.0; self.max_iter];
385
386        for task_name in self.task_outputs.keys() {
387            task_loss_curves.insert(task_name.clone(), vec![0.0; self.max_iter]);
388        }
389
390        let trained_state = MultiTaskNeuralNetworkTrained {
391            shared_weights,
392            shared_biases,
393            task_weights,
394            task_biases,
395            output_weights,
396            output_biases,
397            n_features,
398            task_outputs: self.task_outputs.clone(),
399            shared_layer_sizes: self.shared_layer_sizes.clone(),
400            task_specific_layer_sizes: self.task_specific_layer_sizes.clone(),
401            shared_activation: self.shared_activation,
402            task_activation: self.task_activation,
403            output_activations: self.output_activations.clone(),
404            task_loss_curves,
405            combined_loss_curve,
406            n_iter: self.max_iter,
407        };
408
409        Ok(MultiTaskNeuralNetwork {
410            state: trained_state,
411            shared_layer_sizes: self.shared_layer_sizes,
412            task_specific_layer_sizes: self.task_specific_layer_sizes,
413            task_outputs: self.task_outputs,
414            task_loss_functions: self.task_loss_functions,
415            task_weights: self.task_weights,
416            shared_activation: self.shared_activation,
417            task_activation: self.task_activation,
418            output_activations: self.output_activations,
419            learning_rate: self.learning_rate,
420            max_iter: self.max_iter,
421            tolerance: self.tolerance,
422            random_state: self.random_state,
423            alpha: self.alpha,
424            batch_size: self.batch_size,
425            early_stopping: self.early_stopping,
426            validation_fraction: self.validation_fraction,
427            task_balancing: self.task_balancing,
428        })
429    }
430}
431
432impl Predict<ArrayView2<'_, Float>, HashMap<String, Array2<Float>>>
433    for MultiTaskNeuralNetwork<MultiTaskNeuralNetworkTrained>
434{
435    fn predict(&self, X: &ArrayView2<'_, Float>) -> SklResult<HashMap<String, Array2<Float>>> {
436        let (n_samples, n_features) = X.dim();
437
438        if n_features != self.state.n_features {
439            return Err(SklearsError::InvalidInput(
440                "X has different number of features than training data".to_string(),
441            ));
442        }
443
444        let mut predictions = HashMap::new();
445
446        // Simplified prediction logic
447        for (task_name, &output_size) in &self.state.task_outputs {
448            let task_pred = Array2::<Float>::zeros((n_samples, output_size));
449            predictions.insert(task_name.clone(), task_pred);
450        }
451
452        Ok(predictions)
453    }
454}
455
456impl MultiTaskNeuralNetwork<MultiTaskNeuralNetworkTrained> {
457    /// Get the loss curves for all tasks
458    pub fn task_loss_curves(&self) -> &HashMap<String, Vec<Float>> {
459        &self.state.task_loss_curves
460    }
461
462    /// Get the combined loss curve
463    pub fn combined_loss_curve(&self) -> &[Float] {
464        &self.state.combined_loss_curve
465    }
466
467    /// Get training iterations
468    pub fn n_iter(&self) -> usize {
469        self.state.n_iter
470    }
471
472    /// Get task configurations
473    pub fn task_outputs(&self) -> &HashMap<String, usize> {
474        &self.state.task_outputs
475    }
476}
sklears_multioutput/multitask.rs

sklears_multioutput/
multitask.rs
