use super::config::{ActivationFunction, MetaLearningConfig};
use super::feedforward::FeedForwardNetwork;
use super::layers::LayerNormalization;
use crate::error::Result;
use scirs2_core::ndarray::{Array1, Array2, Array3, Axis};
use scirs2_core::numeric::Float;
use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;
use std::time::Instant;
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum MetaLearningStrategy {
MAML,
FOMAML,
Reptile,
GradientBased,
MemoryAugmented,
}
pub struct TransformerMetaLearning<T: Float + Debug + Send + Sync + 'static> {
strategy: MetaLearningStrategy,
config: MetaLearningConfig<T>,
meta_optimizer: MetaOptimizer<T>,
adaptation_network: AdaptationNetwork<T>,
memory_bank: MemoryBank<T>,
performance_tracker: PerformanceTracker<T>,
meta_state: MetaState<T>,
}
impl<T: Float + Debug + Send + Sync + 'static + scirs2_core::ndarray::ScalarOperand>
TransformerMetaLearning<T>
{
pub fn new(config: &super::config::TransformerBasedOptimizerConfig<T>) -> Result<Self> {
let meta_config = config.meta_learning_config.clone();
let strategy = MetaLearningStrategy::MAML;
let meta_optimizer = MetaOptimizer::new(&meta_config)?;
let adaptation_network =
AdaptationNetwork::new(config.model_dimension, config.feedforward_dimension)?;
let memory_bank = MemoryBank::new(1000, config.model_dimension)?;
let performance_tracker = PerformanceTracker::new();
let meta_state = MetaState::new(config.model_dimension)?;
Ok(Self {
strategy,
config: meta_config,
meta_optimizer,
adaptation_network,
memory_bank,
performance_tracker,
meta_state,
})
}
pub fn meta_step(
&mut self,
tasks: &[TaskBatch<T>],
support_data: &[Array2<T>],
query_data: &[Array2<T>],
) -> Result<MetaLearningResult<T>> {
match self.strategy {
MetaLearningStrategy::MAML => self.maml_step(tasks, support_data, query_data),
MetaLearningStrategy::FOMAML => self.fomaml_step(tasks, support_data, query_data),
MetaLearningStrategy::Reptile => self.reptile_step(tasks, support_data, query_data),
MetaLearningStrategy::GradientBased => {
self.gradient_based_step(tasks, support_data, query_data)
}
MetaLearningStrategy::MemoryAugmented => {
self.memory_augmented_step(tasks, support_data, query_data)
}
}
}
fn maml_step(
&mut self,
tasks: &[TaskBatch<T>],
support_data: &[Array2<T>],
query_data: &[Array2<T>],
) -> Result<MetaLearningResult<T>> {
let start_time = Instant::now();
let mut total_loss = T::zero();
let mut task_adaptations = Vec::new();
for (i, task) in tasks.iter().enumerate() {
let mut adapted_params = self.meta_state.get_parameters().clone();
for inner_step in 0..self.config.inner_steps {
let support_loss =
self.compute_task_loss(&adapted_params, &support_data[i], task)?;
let gradients = self.compute_gradients(&adapted_params, support_loss)?;
for (param, grad) in adapted_params.iter_mut().zip(gradients.iter()) {
*param = *param - self.config.inner_learning_rate * (*grad);
}
}
let query_loss = self.compute_task_loss(&adapted_params, &query_data[i], task)?;
total_loss = total_loss + query_loss;
task_adaptations.push(TaskAdaptation {
task_id: task.id.clone(),
adapted_parameters: adapted_params,
support_loss: self.compute_task_loss(
self.meta_state.get_parameters(),
&support_data[i],
task,
)?,
query_loss,
adaptation_steps: self.config.inner_steps,
});
}
let meta_loss = total_loss / T::from(tasks.len()).expect("unwrap failed");
let meta_gradients = self.compute_meta_gradients(&task_adaptations)?;
self.meta_optimizer
.update(&mut self.meta_state, &meta_gradients)?;
for (i, adaptation) in task_adaptations.iter().enumerate() {
self.memory_bank.store_experience(
&tasks[i],
&adaptation.adapted_parameters,
adaptation.query_loss,
)?;
}
let result = MetaLearningResult {
meta_loss: meta_loss.to_f64().unwrap_or(0.0),
task_adaptations,
computation_time: start_time.elapsed(),
convergence_rate: self.estimate_convergence_rate()?,
};
self.performance_tracker.record_meta_step(result.clone());
Ok(result)
}
fn fomaml_step(
&mut self,
tasks: &[TaskBatch<T>],
support_data: &[Array2<T>],
query_data: &[Array2<T>],
) -> Result<MetaLearningResult<T>> {
self.maml_step(tasks, support_data, query_data)
}
fn reptile_step(
&mut self,
tasks: &[TaskBatch<T>],
support_data: &[Array2<T>],
query_data: &[Array2<T>],
) -> Result<MetaLearningResult<T>> {
let start_time = Instant::now();
let mut parameter_updates = Vec::new();
let mut total_loss = T::zero();
for (i, task) in tasks.iter().enumerate() {
let mut adapted_params = self.meta_state.get_parameters().clone();
for _ in 0..self.config.inner_steps {
let loss = self.compute_task_loss(&adapted_params, &support_data[i], task)?;
let gradients = self.compute_gradients(&adapted_params, loss)?;
for (param, grad) in adapted_params.iter_mut().zip(gradients.iter()) {
*param = *param - self.config.inner_learning_rate * (*grad);
}
}
let original_params = self.meta_state.get_parameters();
let param_diff: Vec<T> = adapted_params
.iter()
.zip(original_params.iter())
.map(|(adapted, original)| *adapted - *original)
.collect();
parameter_updates.push(param_diff);
let query_loss = self.compute_task_loss(&adapted_params, &query_data[i], task)?;
total_loss = total_loss + query_loss;
}
let mut meta_update = vec![T::zero(); self.meta_state.get_parameters().len()];
for param_update in ¶meter_updates {
for (i, &update) in param_update.iter().enumerate() {
meta_update[i] = meta_update[i] + update;
}
}
let num_tasks = T::from(tasks.len()).expect("unwrap failed");
for update in meta_update.iter_mut() {
*update = *update / num_tasks;
}
self.meta_state
.update_parameters(&meta_update, self.config.meta_learning_rate)?;
let result = MetaLearningResult {
meta_loss: (total_loss / num_tasks).to_f64().unwrap_or(0.0),
task_adaptations: Vec::new(), computation_time: start_time.elapsed(),
convergence_rate: self.estimate_convergence_rate()?,
};
self.performance_tracker.record_meta_step(result.clone());
Ok(result)
}
fn gradient_based_step(
&mut self,
tasks: &[TaskBatch<T>],
support_data: &[Array2<T>],
query_data: &[Array2<T>],
) -> Result<MetaLearningResult<T>> {
let start_time = Instant::now();
let mut total_loss = T::zero();
for (i, task) in tasks.iter().enumerate() {
let context_embedding = self.adaptation_network.encode_task_context(task)?;
let predicted_params = self
.adaptation_network
.predict_parameters(&context_embedding)?;
let mut adapted_params = predicted_params;
for _ in 0..self.config.inner_steps {
let loss = self.compute_task_loss(&adapted_params, &support_data[i], task)?;
let gradients = self.compute_gradients(&adapted_params, loss)?;
for (param, grad) in adapted_params.iter_mut().zip(gradients.iter()) {
*param = *param - self.config.inner_learning_rate * (*grad);
}
}
let query_loss = self.compute_task_loss(&adapted_params, &query_data[i], task)?;
total_loss = total_loss + query_loss;
}
let result = MetaLearningResult {
meta_loss: (total_loss / T::from(tasks.len()).expect("unwrap failed"))
.to_f64()
.unwrap_or(0.0),
task_adaptations: Vec::new(),
computation_time: start_time.elapsed(),
convergence_rate: self.estimate_convergence_rate()?,
};
Ok(result)
}
fn memory_augmented_step(
&mut self,
tasks: &[TaskBatch<T>],
support_data: &[Array2<T>],
query_data: &[Array2<T>],
) -> Result<MetaLearningResult<T>> {
let start_time = Instant::now();
let mut total_loss = T::zero();
for (i, task) in tasks.iter().enumerate() {
let relevant_experiences = self.memory_bank.retrieve_similar_experiences(task, 5)?;
let memory_guided_params = self.initialize_from_memory(&relevant_experiences)?;
let mut adapted_params = memory_guided_params;
for _ in 0..self.config.inner_steps {
let loss = self.compute_task_loss(&adapted_params, &support_data[i], task)?;
let gradients = self.compute_gradients(&adapted_params, loss)?;
for (param, grad) in adapted_params.iter_mut().zip(gradients.iter()) {
*param = *param - self.config.inner_learning_rate * (*grad);
}
}
let query_loss = self.compute_task_loss(&adapted_params, &query_data[i], task)?;
total_loss = total_loss + query_loss;
self.memory_bank
.store_experience(task, &adapted_params, query_loss)?;
}
let result = MetaLearningResult {
meta_loss: (total_loss / T::from(tasks.len()).expect("unwrap failed"))
.to_f64()
.unwrap_or(0.0),
task_adaptations: Vec::new(),
computation_time: start_time.elapsed(),
convergence_rate: self.estimate_convergence_rate()?,
};
Ok(result)
}
pub fn generate_update(
&mut self,
transformer_output: &Array2<T>,
current_parameters: &Array1<T>,
) -> Result<Array1<T>> {
let update = self
.adaptation_network
.generate_parameter_update(transformer_output, current_parameters)?;
let scaled_update = self.apply_meta_scaling(&update)?;
Ok(scaled_update)
}
pub fn update_from_loss(&mut self, loss: T) -> Result<()> {
self.meta_state.update_loss_history(loss);
self.performance_tracker
.record_loss(loss.to_f64().unwrap_or(0.0));
Ok(())
}
pub fn set_strategy(&mut self, strategy: MetaLearningStrategy) {
self.strategy = strategy;
}
pub fn get_strategy(&self) -> MetaLearningStrategy {
self.strategy
}
fn compute_task_loss(&self, params: &[T], data: &Array2<T>, task: &TaskBatch<T>) -> Result<T> {
let prediction_error = self.compute_prediction_error(params, data, task)?;
Ok(prediction_error)
}
fn compute_prediction_error(
&self,
_params: &[T],
data: &Array2<T>,
_task: &TaskBatch<T>,
) -> Result<T> {
let mean_squared_error = data
.iter()
.map(|&x| x * x)
.fold(T::zero(), |acc, x| acc + x);
Ok(mean_squared_error / T::from(data.len()).expect("unwrap failed"))
}
fn compute_gradients(&self, params: &[T], loss: T) -> Result<Vec<T>> {
let gradients = params
.iter()
.map(|_| loss / T::from(params.len()).expect("unwrap failed"))
.collect();
Ok(gradients)
}
fn compute_meta_gradients(&self, adaptations: &[TaskAdaptation<T>]) -> Result<Vec<T>> {
let param_count = adaptations[0].adapted_parameters.len();
let mut meta_gradients = vec![T::zero(); param_count];
for adaptation in adaptations {
for (i, ¶m) in adaptation.adapted_parameters.iter().enumerate() {
meta_gradients[i] = meta_gradients[i] + param * adaptation.query_loss;
}
}
let num_tasks = T::from(adaptations.len()).expect("unwrap failed");
for grad in meta_gradients.iter_mut() {
*grad = *grad / num_tasks;
}
Ok(meta_gradients)
}
fn estimate_convergence_rate(&self) -> Result<f64> {
let loss_history = self.performance_tracker.get_loss_history();
if loss_history.len() < 2 {
return Ok(0.0);
}
let recent_losses: Vec<_> = loss_history.iter().rev().take(5).cloned().collect();
let improvement = recent_losses.last().expect("unwrap failed")
- recent_losses.first().expect("unwrap failed");
Ok(improvement.clamp(0.0, 1.0))
}
fn apply_meta_scaling(&self, update: &Array1<T>) -> Result<Array1<T>> {
let scale_factor = self.meta_state.get_scale_factor();
Ok(update * scale_factor)
}
fn initialize_from_memory(&self, experiences: &[MemoryExperience<T>]) -> Result<Vec<T>> {
if experiences.is_empty() {
return Ok(self.meta_state.get_parameters().clone());
}
let param_count = experiences[0].parameters.len();
let mut averaged_params = vec![T::zero(); param_count];
for experience in experiences {
for (i, ¶m) in experience.parameters.iter().enumerate() {
averaged_params[i] = averaged_params[i] + param;
}
}
let num_experiences = T::from(experiences.len()).expect("unwrap failed");
for param in averaged_params.iter_mut() {
*param = *param / num_experiences;
}
Ok(averaged_params)
}
}
pub struct MetaOptimizer<T: Float + Debug + Send + Sync + 'static> {
learning_rate: T,
momentum: Option<T>,
velocity: Option<Vec<T>>,
}
impl<T: Float + Debug + Send + Sync + 'static> MetaOptimizer<T> {
pub fn new(config: &MetaLearningConfig<T>) -> Result<Self> {
Ok(Self {
learning_rate: config.meta_learning_rate,
momentum: None,
velocity: None,
})
}
pub fn update(&mut self, state: &mut MetaState<T>, gradients: &[T]) -> Result<()> {
let params = state.get_parameters_mut();
if let Some(momentum) = self.momentum {
if self.velocity.is_none() {
self.velocity = Some(vec![T::zero(); params.len()]);
}
if let Some(ref mut velocity) = self.velocity {
for i in 0..params.len() {
velocity[i] = momentum * velocity[i] + self.learning_rate * gradients[i];
params[i] = params[i] - velocity[i];
}
}
} else {
for (param, &grad) in params.iter_mut().zip(gradients.iter()) {
*param = *param - self.learning_rate * grad;
}
}
Ok(())
}
}
pub struct AdaptationNetwork<T: Float + Debug + Send + Sync + 'static> {
context_encoder: FeedForwardNetwork<T>,
parameter_predictor: FeedForwardNetwork<T>,
update_generator: FeedForwardNetwork<T>,
model_dimension: usize,
}
impl<T: Float + Debug + Send + Sync + 'static> AdaptationNetwork<T> {
pub fn new(model_dimension: usize, hidden_dimension: usize) -> Result<Self> {
let context_encoder =
FeedForwardNetwork::new(model_dimension, hidden_dimension, ActivationFunction::ReLU)?;
let parameter_predictor =
FeedForwardNetwork::new(hidden_dimension, model_dimension, ActivationFunction::Tanh)?;
let update_generator = FeedForwardNetwork::new(
model_dimension * 2, model_dimension,
ActivationFunction::ReLU,
)?;
Ok(Self {
context_encoder,
parameter_predictor,
update_generator,
model_dimension,
})
}
pub fn encode_task_context(&mut self, _task: &TaskBatch<T>) -> Result<Array2<T>> {
let task_features = Array2::ones((1, self.model_dimension));
self.context_encoder.forward(&task_features)
}
pub fn predict_parameters(&mut self, context: &Array2<T>) -> Result<Vec<T>> {
let predicted = self.parameter_predictor.forward(context)?;
Ok(predicted.row(0).to_vec())
}
pub fn generate_parameter_update(
&mut self,
transformer_output: &Array2<T>,
current_parameters: &Array1<T>,
) -> Result<Array1<T>> {
let batch_size = transformer_output.shape()[0];
let mut input = Array2::zeros((batch_size, self.model_dimension * 2));
for i in 0..batch_size {
for j in 0..self.model_dimension {
input[[i, j]] = transformer_output[[i, j]];
if j < current_parameters.len() {
input[[i, j + self.model_dimension]] = current_parameters[j];
}
}
}
let update = self.update_generator.forward(&input)?;
Ok(update.row(0).to_owned())
}
}
pub struct MemoryBank<T: Float + Debug + Send + Sync + 'static> {
experiences: VecDeque<MemoryExperience<T>>,
max_size: usize,
parameter_dimension: usize,
}
impl<T: Float + Debug + Send + Sync + 'static> MemoryBank<T> {
pub fn new(max_size: usize, parameter_dimension: usize) -> Result<Self> {
Ok(Self {
experiences: VecDeque::new(),
max_size,
parameter_dimension,
})
}
pub fn store_experience(
&mut self,
task: &TaskBatch<T>,
parameters: &[T],
performance: T,
) -> Result<()> {
let experience = MemoryExperience {
task_signature: self.compute_task_signature(task),
parameters: parameters.to_vec(),
performance: performance.to_f64().unwrap_or(0.0),
timestamp: Instant::now(),
};
self.experiences.push_back(experience);
if self.experiences.len() > self.max_size {
self.experiences.pop_front();
}
Ok(())
}
pub fn retrieve_similar_experiences(
&self,
task: &TaskBatch<T>,
k: usize,
) -> Result<Vec<MemoryExperience<T>>> {
let target_signature = self.compute_task_signature(task);
let mut scored_experiences: Vec<_> = self
.experiences
.iter()
.map(|exp| {
let similarity = self.compute_similarity(&target_signature, &exp.task_signature);
(similarity, exp.clone())
})
.collect();
scored_experiences.sort_by(|a, b| b.0.partial_cmp(&a.0).expect("unwrap failed"));
Ok(scored_experiences
.into_iter()
.take(k)
.map(|(_, exp)| exp)
.collect())
}
fn compute_task_signature(&self, task: &TaskBatch<T>) -> Vec<f64> {
vec![task.difficulty, task.complexity, task.data_characteristics]
}
fn compute_similarity(&self, sig1: &[f64], sig2: &[f64]) -> f64 {
if sig1.len() != sig2.len() {
return 0.0;
}
let dot_product: f64 = sig1.iter().zip(sig2.iter()).map(|(a, b)| a * b).sum();
let norm1: f64 = sig1.iter().map(|x| x * x).sum::<f64>().sqrt();
let norm2: f64 = sig2.iter().map(|x| x * x).sum::<f64>().sqrt();
if norm1 == 0.0 || norm2 == 0.0 {
0.0
} else {
dot_product / (norm1 * norm2)
}
}
}
#[derive(Debug, Clone)]
pub struct TaskBatch<T: Float + Debug + Send + Sync + 'static> {
pub id: String,
pub difficulty: f64,
pub complexity: f64,
pub data_characteristics: f64,
pub _phantom: std::marker::PhantomData<T>,
}
#[derive(Debug, Clone)]
pub struct TaskAdaptation<T: Float + Debug + Send + Sync + 'static> {
pub task_id: String,
pub adapted_parameters: Vec<T>,
pub support_loss: T,
pub query_loss: T,
pub adaptation_steps: usize,
}
#[derive(Debug, Clone)]
pub struct MetaLearningResult<T: Float + Debug + Send + Sync + 'static> {
pub meta_loss: f64,
pub task_adaptations: Vec<TaskAdaptation<T>>,
pub computation_time: std::time::Duration,
pub convergence_rate: f64,
}
#[derive(Debug, Clone)]
pub struct MemoryExperience<T: Float + Debug + Send + Sync + 'static> {
pub task_signature: Vec<f64>,
pub parameters: Vec<T>,
pub performance: f64,
pub timestamp: Instant,
}
pub struct PerformanceTracker<T: Float + Debug + Send + Sync + 'static> {
loss_history: VecDeque<f64>,
meta_results: VecDeque<MetaLearningResult<T>>,
}
impl<T: Float + Debug + Send + Sync + 'static> Default for PerformanceTracker<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Float + Debug + Send + Sync + 'static> PerformanceTracker<T> {
pub fn new() -> Self {
Self {
loss_history: VecDeque::new(),
meta_results: VecDeque::new(),
}
}
pub fn record_loss(&mut self, loss: f64) {
self.loss_history.push_back(loss);
if self.loss_history.len() > 1000 {
self.loss_history.pop_front();
}
}
pub fn record_meta_step(&mut self, result: MetaLearningResult<T>) {
self.meta_results.push_back(result);
if self.meta_results.len() > 100 {
self.meta_results.pop_front();
}
}
pub fn get_loss_history(&self) -> &VecDeque<f64> {
&self.loss_history
}
}
#[derive(Debug, Clone)]
pub struct MetaState<T: Float + Debug + Send + Sync + 'static> {
parameters: Vec<T>,
loss_history: VecDeque<T>,
scale_factor: T,
}
impl<T: Float + Debug + Send + Sync + 'static> MetaState<T> {
pub fn new(parameter_count: usize) -> Result<Self> {
Ok(Self {
parameters: vec![T::zero(); parameter_count],
loss_history: VecDeque::new(),
scale_factor: T::one(),
})
}
pub fn get_parameters(&self) -> &Vec<T> {
&self.parameters
}
pub fn get_parameters_mut(&mut self) -> &mut Vec<T> {
&mut self.parameters
}
pub fn update_parameters(&mut self, updates: &[T], learning_rate: T) -> Result<()> {
for (param, &update) in self.parameters.iter_mut().zip(updates.iter()) {
*param = *param + learning_rate * update;
}
Ok(())
}
pub fn update_loss_history(&mut self, loss: T) {
self.loss_history.push_back(loss);
if self.loss_history.len() > 100 {
self.loss_history.pop_front();
}
}
pub fn get_scale_factor(&self) -> T {
self.scale_factor
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
#[ignore]
fn test_meta_learning_creation() {
let config = super::super::config::TransformerBasedOptimizerConfig::<f32>::default();
let meta_learning = TransformerMetaLearning::new(&config);
assert!(meta_learning.is_ok());
}
#[test]
fn test_memory_bank() {
let memory = MemoryBank::<f32>::new(100, 64);
assert!(memory.is_ok());
let mut bank = memory.expect("unwrap failed");
let task = TaskBatch {
id: "test".to_string(),
difficulty: 0.5,
complexity: 0.7,
data_characteristics: 0.3,
_phantom: std::marker::PhantomData,
};
let params = vec![0.1f32; 64];
assert!(bank.store_experience(&task, ¶ms, 0.8).is_ok());
}
#[test]
fn test_adaptation_network() {
let network = AdaptationNetwork::<f32>::new(128, 256);
assert!(network.is_ok());
let mut net = network.expect("unwrap failed");
let task = TaskBatch {
id: "test".to_string(),
difficulty: 0.5,
complexity: 0.7,
data_characteristics: 0.3,
_phantom: std::marker::PhantomData,
};
let context = net.encode_task_context(&task);
assert!(context.is_ok());
}
}