diff --git a/.gitignore b/.gitignore
index 501eeae..17d9625 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,12 +3,14 @@ generated/**
 wandb/**
 */.DS_Store
 .idea
+*.png
 
 /data
 /checkpoint
 /logs
 /plots
 /archive/imgs/
+/tmp
 
 
 # Created by https://www.gitignore.io/api/macos,python,R,linux,vim,emacs,visualstudiocode,intellij
diff --git a/src/asap/__init__.py b/src/asap/__init__.py
index b412d77..f2bdad5 100644
--- a/src/asap/__init__.py
+++ b/src/asap/__init__.py
@@ -1,2 +1,5 @@
 from .dataset import training_datasets, peak_dataset, wg_dataset, robustness_peak_dataset, robustness_wg_dataset 
-from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac
+
+
+from .task import train_model, eval_model, eval_robustness, export_predictions, predict_snv_atac, train_new_head_lp, train_new_head_ft, train_multiheaded_model, eval_multihead_model, train_multiheaded_model_progressively, extend_multiheaded_model_progressively
+
diff --git a/src/asap/dataloader/base.py b/src/asap/dataloader/base.py
index 473eba7..e6545c6 100644
--- a/src/asap/dataloader/base.py
+++ b/src/asap/dataloader/base.py
@@ -31,10 +31,12 @@ def __init__(
     ):
         super().__init__()
         self.genome = genome
-        if signal_files is not None:
-            self.signal_files = [signal_files]
-        else:
+        if signal_files is None:
             self.signal_files = None
+        elif type(signal_files) is list:
+            self.signal_files = signal_files
+        else:
+            self.signal_files = [signal_files]
 
         # if we randomly shift to augment the data
         self.random_shift = random_shift
diff --git a/src/asap/models/convnext_dcnn.py b/src/asap/models/convnext_dcnn.py
index 37dbbad..f3f01b4 100644
--- a/src/asap/models/convnext_dcnn.py
+++ b/src/asap/models/convnext_dcnn.py
@@ -20,7 +20,8 @@
     'kernel2': 1,
     'dropout': 0.3,
     'final_dropout': 0.05,
-    'use_map': False
+    'use_map': False,
+    'num_heads': 1,
 }
 
 
@@ -50,7 +51,8 @@ def __init__(self, **kwargs) -> None:
                       kernel1=config['kernel1'],
                       kernel2=config['kernel2'],
                       dropout=config['dropout'],
-                      final_dropout=config['final_dropout'])
+                      final_dropout=config['final_dropout'],
+                      num_heads=config['num_heads'])
     
     def forward(self, x:torch.Tensor, return_unmap=False) -> torch.Tensor:
         x = torch.transpose(x, dim0=-1, dim1=-2)
@@ -70,7 +72,7 @@ class BasenjiCoreBlock(nn.Module):
     def __init__(self, nr_tracks: int, window: int, filters_in,
                   nr_res_blocks: int = 11, rate_mult: float = 1.5, bin_size: int = 100, filters1: int = 128,
                   filters3: int = None, kernel1: int = 3, kernel2: int = 1, dropout: float = 0.3,
-                  final_dropout: float = 0.05):
+                  final_dropout: float = 0.05, num_heads=1):
         super().__init__()
         if not filters3:
             filters3 = window
@@ -93,10 +95,11 @@ def __init__(self, nr_tracks: int, window: int, filters_in,
         pool_size = bin_size // 2
         self.pool = nn.AvgPool1d(kernel_size=pool_size, stride=pool_size)
 
-        self.linear_out = nn.Linear(
-            in_features=filters3,
-            out_features=nr_tracks,
-        )
+        print("creating output layer with tracks:", num_heads)
+        self.heads = nn.ModuleList([
+            nn.Linear(filters3, nr_tracks)
+            for _ in range(num_heads)
+        ])
         self.activation = nn.Softplus()
     
     def forward(self, x:torch.Tensor) -> torch.Tensor:
@@ -115,9 +118,12 @@ def forward(self, x:torch.Tensor) -> torch.Tensor:
         x = self.final_dropout(x)
         x = self.pool(x)
         x = torch.transpose(x, -2, -1)
-        x = self.linear_out(x)
-        x = self.activation(x)
-        return x
+        outputs = []
+        for head in self.heads:
+            out = head(x)
+            out = self.activation(out)
+            outputs.append(out)
+        return outputs
 
 
 class ConvBlock(nn.Module):
diff --git a/src/asap/snv/predict.py b/src/asap/snv/predict.py
index 619609d..8de81f4 100644
--- a/src/asap/snv/predict.py
+++ b/src/asap/snv/predict.py
@@ -20,7 +20,7 @@ def _idx_to_ohe(idx: np.ndarray) -> np.ndarray:
     one_hot_encoded = eye[idx]
     return one_hot_encoded
 
-def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: nn.Module, genome: str, margin_size: int, window_size: int, bin_size: int, device: torch.device) -> None:
+def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: nn.Module, genome: str, margin_size: int, window_size: int, bin_size: int, device: torch.device, target_head=0) -> None:
     # Ensure output columns exist
     output_columns = ['signal_true', 'signal_pred_ref', 'signal_pred_alt']
     for col in output_columns:
@@ -56,8 +56,8 @@ def add_predictions(snv: pd.DataFrame, chroms: List[int], bw_file: str, model: n
             x = _idx_to_ohe(x).astype(np.float32)
             x_var = _idx_to_ohe(x_var).astype(np.float32)
 
-            ref_pred = model(torch.from_numpy(x).to(device))
-            var_pred = model(torch.from_numpy(x_var).to(device))
+            ref_pred = model(torch.from_numpy(x).to(device))[target_head]
+            var_pred = model(torch.from_numpy(x_var).to(device))[target_head]
             for (i, _), true, ref, var in zip(chr_df_i.iterrows(), y, ref_pred.cpu().detach().numpy(), var_pred.cpu().detach().numpy()):
                 snv.loc[i, 'signal_true'] = json.dumps(true.flatten().tolist())
                 snv.loc[i, 'signal_pred_ref'] = json.dumps(ref.tolist())
diff --git a/src/asap/task.py b/src/asap/task.py
index 06d60e9..cf4213e 100644
--- a/src/asap/task.py
+++ b/src/asap/task.py
@@ -13,7 +13,7 @@
 from .models import VanillaCNN, CNN_LSTM, DilatedCNN, ConvNextTransformer, ConvNeXtCNN, ConvNeXtLSTM, ConvNeXtDCNN
 from .snv import make_pcawg_df, add_predictions
 
-def _get_model(model_name: str, use_map: bool = False):
+def _get_model(model_name: str, use_map: bool = False, num_heads=1):
     if model_name == 'cnn':
         return VanillaCNN(use_map=use_map)
     elif model_name == 'lstm':
@@ -27,7 +27,7 @@ def _get_model(model_name: str, use_map: bool = False):
     elif model_name == 'convnext_lstm':
         return ConvNeXtLSTM(use_map=use_map)
     elif model_name == 'convnext_dcnn':
-        return ConvNeXtDCNN(use_map=use_map)
+        return ConvNeXtDCNN(use_map=use_map, num_heads=num_heads)
     else:
         raise ValueError(f'Unknown model name: {model_name}')
 
@@ -75,7 +75,185 @@ def train_model(experiment_name : str, model: str, train_dataset: BaseDataset, v
     # Start training
     trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate)
 
-def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False):
+def train_new_head_ft(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1):
+    '''
+    Add a new head to a bse model and train it using finetuining
+    
+    Args:
+        base_experiment_name(str): The name of the base model 
+        new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved
+        model (str): The model to train.
+        train_dataset: The training dataset.
+        val_dataset: The validation dataset.
+        logs_dir (str): The directory to save logs.
+        n_gpus (int): The number of GPUs to use for training.
+        max_epochs (int): The maximum number of epochs to train.
+        learning_rate (float): The learning rate for the optimizer.
+        batch_size (int): The batch size for training.
+        use_map (bool): Whether to use mappability for training.
+        num_heads (int): The number of heads (=num of prediction signals) of the new model (=num base heads + 1)
+    '''
+    if n_gpus > 0 and not torch.cuda.is_available():
+        n_gpus = 0
+        print("No GPU available, using CPU instead.")
+    
+    # Count the number of GPUs available
+    if n_gpus > torch.cuda.device_count():
+        n_gpus = torch.cuda.device_count()
+        print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.")
+
+    # Initialize the model
+    model = _get_model(model, use_map=use_map, num_heads=num_heads)
+    
+    # freeze old heads
+    for param in model.core.heads[:-1].parameters():
+        param.requires_grad = False
+
+    # set modes
+    model.train()
+    for head in model.core.heads[:-1]:
+        head.eval()
+
+    # Initialize the trainer with the model and datasets
+    trainer = Trainer(
+        filename=new_experiment_name, 
+        model=model,
+        criterion=nn.PoissonNLLLoss(log_input=False),
+        unmap_criterion=use_map,
+        batch_size=batch_size,
+        logger=TextLogger(logs_dir=logs_dir), 
+        n_gpus=n_gpus,
+        fine_tune=True,
+        num_heads=num_heads,
+    )
+
+    print(f'Loading best model weights from {base_experiment_name}')
+    checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / base_experiment_name / 'checkpoint.pth'
+
+    # Load the file manually to use strict=False
+    checkpoint = torch.load(checkpoint_path, map_location='cpu') # Load to CPU first to avoid OOM
+    trainer.model.load_state_dict(checkpoint, strict=False)
+
+
+    # Start training
+    trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate)
+    print("trained a new model")
+
+
+def train_new_head_lp(base_experiment_name: str, new_experiment_name: str, model: str, train_dataset: BaseDataset, val_dataset: BaseDataset, logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_original_heads: int=1):
+    '''
+    Add a new head to a bse model and train it using linear probing
+    
+    Args:
+        base_experiment_name(str): The name of the base model 
+        new_experiment_name (str): The name of the new experiment. This is how the finetuned model will be saved
+        model (str): The model to train.
+        train_dataset: The training dataset.
+        val_dataset: The validation dataset.
+        logs_dir (str): The directory to save logs.
+        n_gpus (int): The number of GPUs to use for training.
+        max_epochs (int): The maximum number of epochs to train.
+        learning_rate (float): The learning rate for the optimizer.
+        batch_size (int): The batch size for training.
+        use_map (bool): Whether to use mappability for training.
+        num_original_heads (int): The number of heads (=num of prediction signals) of the old model
+    '''
+    if n_gpus > 0 and not torch.cuda.is_available():
+        n_gpus = 0
+        print("No GPU available, using CPU instead.")
+    
+    # Count the number of GPUs available
+    if n_gpus > torch.cuda.device_count():
+        n_gpus = torch.cuda.device_count()
+        print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.")
+
+    # Initialize the model
+    model = _get_model(model, use_map=use_map, num_heads=num_original_heads+1)
+
+    # freeze everything
+    for param in model.parameters():
+        param.requires_grad = False
+
+    # unfreeze head
+    for param in model.core.heads[-1].parameters():
+        param.requires_grad = True
+
+    # set modes such that there is no dropout for the core
+    model.eval()
+    model.core.heads[-1].train()
+
+    # Initialize the trainer with the model and datasets
+    trainer = Trainer(
+        filename=new_experiment_name, 
+        model=model,
+        criterion=nn.PoissonNLLLoss(log_input=False),
+        unmap_criterion=use_map,
+        batch_size=batch_size,
+        logger=TextLogger(logs_dir=logs_dir), 
+        n_gpus=n_gpus,
+        linear_probe=True,
+        num_heads=num_original_heads+1,
+    )
+
+    print(f'Loading best model weights from {base_experiment_name}')
+    checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / base_experiment_name / 'checkpoint.pth'
+
+    # Load the file manually to use strict=False
+    checkpoint = torch.load(checkpoint_path, map_location='cpu') # Load to CPU first to avoid OOM
+    trainer.model.load_state_dict(checkpoint, strict=False)
+
+
+    # Start training
+    trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate)
+    print("trained a new model")
+
+def train_multiheaded_model(experiment_name : str, model: str,  train_dataset: List[BaseDataset], val_dataset: List[BaseDataset], logs_dir: str, n_gpus: int=0, max_epochs: int=70, learning_rate: float=1e-3, batch_size: int=64, use_map: bool=False, num_heads: int=1):
+    """
+    Train the model with the given datasets and parameters.
+    
+    Args:
+        experiment_name (str): The name of the experiment.
+        model (str): The model to train.
+        train_dataset: The training dataset.
+        val_dataset: The validation dataset.
+        logs_dir (str): The directory to save logs.
+        n_gpus (int): The number of GPUs to use for training.
+        max_epochs (int): The maximum number of epochs to train.
+        learning_rate (float): The learning rate for the optimizer.
+        batch_size (int): The batch size for training.
+        use_map (bool): Whether to use mappability for training.
+        num_heads (int): The number of heads (=num of prediction signals)
+    """
+
+    # Check if gpu is available
+    if n_gpus > 0 and not torch.cuda.is_available():
+        n_gpus = 0
+        print("No GPU available, using CPU instead.")
+    
+    # Count the number of GPUs available
+    if n_gpus > torch.cuda.device_count():
+        n_gpus = torch.cuda.device_count()
+        print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.")
+
+    # Initialize the model
+    model = _get_model(model, use_map=use_map, num_heads=num_heads)
+
+    # Initialize the trainer with the model and datasets
+    trainer = Trainer(
+        filename=experiment_name, 
+        model=model,
+        criterion=nn.PoissonNLLLoss(log_input=False),
+        unmap_criterion=use_map,
+        batch_size=batch_size,
+        logger=TextLogger(logs_dir=logs_dir), 
+        n_gpus=n_gpus,
+        num_heads=num_heads
+    )
+
+    # Start training
+    trainer.fit(train_dset=train_dataset, val_dset=val_dataset, nr_epochs=max_epochs, learning_rate=learning_rate)
+
+def eval_multihead_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False,  num_heads: int=1, target_head:int = 0):
     '''
     Evaluate the model on the given dataset.
     Args:
@@ -89,7 +267,55 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs
     n_gpus = 1 if torch.cuda.is_available() else 0
 
     # Initialize the model
-    model = _get_model(model, use_map=use_map)
+    model = _get_model(model, use_map=use_map, num_heads=num_heads)
+
+    trainer = Trainer(
+        filename=experiment_name, 
+        model=model,
+        criterion=nn.PoissonNLLLoss(log_input=False),
+        unmap_criterion=use_map,
+        batch_size=batch_size,
+        logger=TextLogger(logs_dir=logs_dir), 
+        n_gpus=n_gpus,
+        num_heads=num_heads
+    )
+
+    # for evaluation use the checkpoint of the best model
+    print(f'Loading best model weights from {trainer.filename}')
+    checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth'
+    trainer.load_weights(checkpoint_path)
+    
+    test_chroms = eval_dataset.chroms
+    scores = {}
+    for chrom in test_chroms:
+        eval_dataset.set_chroms([chrom])
+        test_gen = make_dataloader(
+            ddp_enabled=False,
+            dataset=eval_dataset,
+            batch_size=batch_size, 
+            is_train=False
+        )
+
+        _, _, result_metrics = trainer.predict_and_evaluate_multihead(test_gen, target_head=target_head)
+        scores[chrom] = {key: result_metrics[key] for key in ['pearson_r', 'mse', 'poisson_nll', 'spearman_r', 'kendall_tau']}
+    return scores
+
+
+def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False,  num_heads: int=1):
+    '''
+    Evaluate the model on the given dataset.
+    Args:
+        experiment_name (str): The name of the experiment.
+        model (str): The model to evaluate.
+        eval_dataset: The evaluation dataset.
+        logs_dir (str): The directory to load model checkpoints from.
+        batch_size (int): The batch size for evaluation.
+        use_map (bool): If mappability information was used during training.
+    '''
+    n_gpus = 1 if torch.cuda.is_available() else 0
+
+    # Initialize the model
+    model = _get_model(model, use_map=use_map, num_heads=num_heads)
 
     trainer = Trainer(
         filename=experiment_name, 
@@ -121,6 +347,243 @@ def eval_model(experiment_name: str, model: str, eval_dataset: BaseDataset, logs
         scores[chrom] = {key: result_metrics[key] for key in ['pearson_r', 'mse', 'poisson_nll', 'spearman_r', 'kendall_tau']}
     return scores
 
+def train_multiheaded_model_progressively(
+        experiment_name: str,
+        model: str,
+        train_dataset: List[BaseDataset],
+        val_dataset: List[BaseDataset],
+        logs_dir: str,
+        n_gpus: int=0,
+        max_epochs: int=20,
+        learning_rate: float=1e-3,
+        batch_size: int=64,
+        use_map: bool=False,
+        num_heads: List[int] = [1, 2],
+        checkpoint_on_new_heads_only: bool=False,
+):
+    '''
+    Train the model with the given datasets and parameters progressively. In Progressive Joint Training (PJT),
+    the number of heads should be strictly increasing at every step. In NJT*, the number of heads should be constant,
+    equal to the final number of heads. 
+
+    Args:
+        experiment_name (str): The base name of the experiments. Actual names are {experiment_name}_{i}_{num_heads[i]}, with i the current step.
+        model (str): The model to train.
+        train_dataset (List[BaseDataset]): The training dataset.
+        val_dataset (List[BaseDataset]): The validation dataset.
+        logs_dir (str): The directory to save logs.
+        n_gpus (int): The number of GPUs to use for training.
+        max_epochs (int): The maximum number of epochs to train (per step).
+        learning_rate (float): The learning rate for the optimizer.
+        batch_size (int): The batch size for training.
+        use_map (bool): Whether to use mappability for training.
+        num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on.
+        checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads.
+    '''
+
+    # Validate num_heads
+    if (len(num_heads) < 2):
+        raise ValueError("At least two values must be provided for num_heads when training progressively.")
+    for i in range(len(num_heads)):
+        if num_heads[i] <= 0:
+            raise ValueError("Cannot train on less than one head at any point.")
+        if i > 0:
+            if num_heads [i-1] > num_heads[i]:
+                raise ValueError("Number of heads must be nondecreasing.")
+
+
+    # Check if gpu is available
+    if n_gpus > 0 and not torch.cuda.is_available():
+        n_gpus = 0
+        print("No GPU available, using CPU instead.")
+
+    # Count the number of GPUs available
+    if n_gpus > torch.cuda.device_count():
+        n_gpus = torch.cuda.device_count()
+        print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.")
+
+    # Initialize the first model
+    model_type = model
+    model = _get_model(model_type, use_map=use_map, num_heads=num_heads[0])
+    trainer = Trainer(
+        filename=experiment_name + f'_0_{num_heads[0]}',
+        model=model,
+        criterion=nn.PoissonNLLLoss(log_input=False),
+        unmap_criterion=use_map,
+        batch_size=batch_size,
+        logger=TextLogger(logs_dir=logs_dir),
+        n_gpus=n_gpus,
+        num_heads=num_heads[0],
+    )
+
+    print(f'Starting training for the initial step with {num_heads[0]} heads.')
+
+    # Train the first model
+    trainer.fit(
+        train_dset=train_dataset[0],
+        val_dset=val_dataset[0],
+        nr_epochs=max_epochs,
+        learning_rate=learning_rate,
+        # validation always occurs on all heads in the very first step 
+    )
+
+    print(f'Finished training for the initial step. Starting training for the next step with {num_heads[1]} heads.')
+
+    # Progressively learn next heads
+    for i in range(1, len(num_heads)):
+        # print(f'Starting progressive training for step {i} with {num_heads[i]} heads.')
+
+        # Initialize the next model with the previous model's weights
+        model_tmp = _get_model(model_type, use_map=use_map, num_heads=num_heads[i])
+
+        print(f'Loading best model weights from {trainer.filename}')
+        checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth'
+        trainer.load_weights(checkpoint_path)
+
+        print('Loading model weights from previous step.')
+        model_tmp.load_state_dict(trainer.model.state_dict(), strict=False) # strict=False should allow loading when number of heads changes
+        model = model_tmp
+
+        print('Successfully loaded model weights.')
+
+        # Initialize new trainer
+        print(f'Initializing trainer for step {i}.')
+        trainer = Trainer(
+            filename=experiment_name + f'_{i}_{num_heads[i]}',
+            model=model,
+            criterion=nn.PoissonNLLLoss(log_input=False),
+            unmap_criterion=use_map,
+            batch_size=batch_size,
+            logger=TextLogger(logs_dir=logs_dir),
+            n_gpus=n_gpus,
+            num_heads=num_heads[i],
+        )
+
+        # Train the new model
+        print(f'Starting training for the step {i} with {num_heads[i]} heads.')
+        trainer.fit(
+            train_dset=train_dataset[0],
+            val_dset=val_dataset[0],
+            nr_epochs=max_epochs,
+            learning_rate=learning_rate,
+            val_on_heads=list(range(num_heads[i-1], num_heads[i])) if checkpoint_on_new_heads_only else None,
+        )
+        print(f'Finished training for the step {i} with {num_heads[i]} heads.')
+
+    print('Finished progressively training the model.')
+    print(f'Final model is saved as {experiment_name}_{len(num_heads)-1}_{num_heads[-1]} in {logs_dir}.')
+
+def extend_multiheaded_model_progressively(
+        base_experiment_name: str,
+        new_experiment_name: str,
+        model: str,
+        train_dataset: List[BaseDataset],
+        val_dataset: List[BaseDataset],
+        logs_dir: str,
+        n_gpus: int=0,
+        max_epochs: int=20,
+        learning_rate: float=1e-3,
+        batch_size: int=64,
+        use_map: bool=False,
+        num_heads: List[int] = [3, 4], # num_heads[0] is the number of heads in the already trained model, with num_heads[1] the first step of JT extension
+        checkpoint_on_new_heads_only: bool=False,
+):
+    '''
+    Train the model with the given datasets and parameters progressively. Start with an already trained model 
+    with some number of heads (num_heads[0]). The following stages can have an increasing number of heads (PJT),
+    or the same as the pretrained model (NJT*).
+
+    Args:
+        base_experiment_name (str): The base name of the previous experiment. The model will be loaded from {base_experiment_name}_0_{num_heads[0]}.
+        new_experiment_name (str): The base name of the new experiment. Actual names are {new_experiment_name}_{i}_{num_heads[i]}, with i the current step (ignoring step 0).
+        model (str): The model to train.
+        train_dataset (List[BaseDataset]): The training dataset.
+        val_dataset (List[BaseDataset]): The validation dataset.
+        logs_dir (str): The directory to save logs.
+        n_gpus (int): The number of GPUs to use for training.
+        max_epochs (int): The maximum number of epochs to train (per step).
+        learning_rate (float): The learning rate for the optimizer.
+        batch_size (int): The batch size for training.
+        use_map (bool): Whether to use mappability for training.
+        num_heads (List[int]): Number of heads at each step of the progressive training. Model is trained with num_heads[0] heads first, then num_heads[1], and so on.
+        checkpoint_on_new_heads_only (bool): Whether to do validation only on the new heads.
+    '''
+    # Validate num_heads
+    if (len(num_heads) < 2):
+        raise ValueError("At least two values must be provided for num_heads when training progressively.")
+    for i in range(len(num_heads)):
+        if num_heads[i] <= 0:
+            raise ValueError("Cannot train on less than one head at any point.")
+        if i > 0:
+            if num_heads [i-1] > num_heads[i]:
+                raise ValueError("Number of heads must be nondecreasing.")
+
+    # Check if gpu is available
+    if n_gpus > 0 and not torch.cuda.is_available():
+        n_gpus = 0
+        print("No GPU available, using CPU instead.")
+
+    # Count the number of GPUs available
+    if n_gpus > torch.cuda.device_count():
+        n_gpus = torch.cuda.device_count()
+        print(f"Requested {n_gpus} GPUs, but only {torch.cuda.device_count()} are available. Using {n_gpus} GPUs instead.")
+
+    # Create model with more heads
+    model_type = model
+    model = _get_model(model_type, use_map=use_map, num_heads=num_heads[0])
+
+    # Load the previous model
+    trainer = Trainer(
+        filename=base_experiment_name + f'_0_{num_heads[0]}',
+        model=model,
+        criterion=nn.PoissonNLLLoss(log_input=False),
+        unmap_criterion=use_map,
+        batch_size=batch_size,
+        logger=TextLogger(logs_dir=logs_dir),
+        n_gpus=n_gpus,
+        num_heads=num_heads[0],
+    )
+
+    for i in range(1, len(num_heads)):
+        model_tmp = _get_model(model_type, use_map=use_map, num_heads=num_heads[i])
+
+        print(f'Loading best model weights from {trainer.filename}')
+        checkpoint_path = pathlib.Path(trainer.logger.logs_dir) / trainer.filename / 'checkpoint.pth'
+        trainer.load_weights(checkpoint_path)
+
+        print('Loading model weights from previous step.')
+        model_tmp.load_state_dict(trainer.model.state_dict(), strict=False) # strict=False should allow loading when number of heads changes
+        model = model_tmp
+
+        print('Successfully loaded model weights.')
+
+        # Initialize new trainer
+        print(f'Initializing trainer for step {i}.')
+        trainer = Trainer(
+            filename=new_experiment_name + f'_{i}_{num_heads[i]}',
+            model=model,
+            criterion=nn.PoissonNLLLoss(log_input=False),
+            unmap_criterion=use_map,
+            batch_size=batch_size,
+            logger=TextLogger(logs_dir=logs_dir),
+            n_gpus=n_gpus,
+            num_heads=num_heads[i],
+        )
+
+        # Train the new model
+        print(f'Starting training for the step {i} with {num_heads[i]} heads.')
+        trainer.fit(
+            train_dset=train_dataset[0],
+            val_dset=val_dataset[0],
+            nr_epochs=max_epochs,
+            learning_rate=learning_rate,
+            val_on_heads=list(range(num_heads[i-1], num_heads[i])) if checkpoint_on_new_heads_only else None,
+        )
+        print(f'Finished training for the step {i} with {num_heads[i]} heads.')
+    print('Finished progressively extending the model.')
+    print(f'Final model is saved as {new_experiment_name}_{len(num_heads)-1}_{num_heads[-1]} in {logs_dir}.')
+
+
 def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset, logs_dir: str, batch_size: int=64, use_map: bool=False, nr_samples_for_var: int=17):
     '''
     Evaluate the robustness of the model on the given dataset.
@@ -173,7 +636,8 @@ def eval_robustness(experiment_name: str, model: str, eval_dataset: BaseDataset,
         scores[chrom] = {'cov': float(np.nanmean(cov)), 'cov_per_bin': {f'bin_{i}': float(cov_per_bin[i]) for i in range(len(cov_per_bin))}}
     return scores
 
-def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_file: str, logs_dir: str, out_dir: str, genome: str, chroms: List[int]=[*range(1,23)], use_map: bool=False, export_bigwig: str=None, scale: dict | float=1.0):
+def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_file: str, logs_dir: str, out_dir: str, genome: str, chroms: List[int]=[*range(1,23)], use_map: bool=False, export_bigwig: str=None, scale: dict | float=1.0, num_heads=1,
+        target_head=0):
     """
     Predict ATAC-seq for SNVs using the trained model.
     Args:
@@ -200,7 +664,7 @@ def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_fil
     snv_file_name = pathlib.Path(snv_file).stem.split('.')[0]
 
     # Initialize the model
-    model = _get_model(model, use_map=use_map)
+    model = _get_model(model, use_map=use_map, num_heads=num_heads)
     checkpoint_path = pathlib.Path(logs_dir) / experiment_name / 'checkpoint.pth'
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     state_dict = torch.load(checkpoint_path, map_location=device)
@@ -224,6 +688,7 @@ def predict_snv_atac(experiment_name: str, model: str, snv_file: str, signal_fil
         window_size=window_size,
         bin_size=bin_size,
         device=device,
+        target_head=target_head
     )
 
     # Save the results to a CSV file
diff --git a/src/asap/trainer/trainer.py b/src/asap/trainer/trainer.py
index ef02d66..4f4308f 100644
--- a/src/asap/trainer/trainer.py
+++ b/src/asap/trainer/trainer.py
@@ -28,6 +28,10 @@ def __init__(self,
                  batch_size: int = None,
                  logger: Logger = None,
                  n_gpus: int = None,
+                 nr_tracks: int = 1,
+                 num_heads: int = 1,
+                 linear_probe=False,
+                 fine_tune=False,
                  ):
         self.filename = filename
         self.model = model
@@ -40,6 +44,9 @@ def __init__(self,
         self.nr_tracks = 1
         self.nr_devices = n_gpus
         self.batch_size = batch_size
+        self.num_heads = num_heads
+        self.linear_probe = linear_probe
+        self.fine_tune = fine_tune
         if self.nr_devices > 1: 
             self.ddp_enabled = True
             self.device = 'cuda'
@@ -53,7 +60,7 @@ def __init__(self,
             self.device = 'cpu'
             self.model.to(self.device)
 
-    def fit(self, train_dset, val_dset, nr_epochs, learning_rate):
+    def fit(self, train_dset, val_dset, nr_epochs, learning_rate, val_on_heads=None):
         print(f'Training {self.filename}...')
         if self.nr_devices > 1:
             port = 10000 + randint(0,2355)
@@ -71,7 +78,11 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate):
                     self.logger,
                     self.filename,
                     self.nr_devices,
-                    port
+                    self.linear_probe,
+                    self.fine_tune,
+                    port,
+                    self.num_heads,
+                    val_on_heads
                 ),
                 nprocs=self.nr_devices
             )
@@ -100,7 +111,11 @@ def fit(self, train_dset, val_dset, nr_epochs, learning_rate):
                 self.unmap_criterion,
                 self.logger,
                 self.filename,
-                ddp_enabled=False
+                ddp_enabled=False,
+                num_heads=self.num_heads,
+                linear_probe=self.linear_probe,
+                fine_tune=self.fine_tune,
+                val_on_heads=val_on_heads
             )
 
     def predict(self, gen):
@@ -115,6 +130,42 @@ def evaluate(self, test_gen) -> dict:
         print(result_metrics)
         return result_metrics
 
+    def predict_and_evaluate_multihead(self, gen, metrics_for_track=None, no_eval=False, target_head=0) -> Tuple[np.ndarray, np.ndarray, dict]:
+        self.model.eval()
+
+        try:
+            gen.dataset.margin_size
+        except AttributeError:
+            print("Generator has no margin size -- assuming full prediction.")
+
+        predictions, true = self.predict(gen)
+        predictions = predictions[..., target_head: target_head+1]
+        predictions = predictions.reshape((-1, 1)).detach().cpu().numpy()
+        true = true.reshape((-1, 1)).detach().cpu().numpy()
+        
+        if no_eval:
+            return true, predictions, None
+
+        metric_results = {}
+        if metrics_for_track is None:
+            # compute metrics for all tracks
+            metrics_for_track = range(self.nr_tracks)
+            metrics_ = compute_metrics(
+                    predictions.flatten(),
+                    true.flatten(),
+                    logspace_input=self.logspace
+                )
+            metric_results.update(metrics_)
+        else:
+            for track in metrics_for_track:
+                metrics_ = compute_metrics(
+                    predictions[:, track],
+                    true[:, track],
+                    logspace_input=self.logspace
+                )
+                metric_results.update(prepend_to_keys(metrics_, f"_{track}_"))
+        return true, predictions, metric_results
+
     def predict_and_evaluate(self, gen, metrics_for_track=None, no_eval=False) -> Tuple[np.ndarray, np.ndarray, dict]:
         self.model.eval()
 
@@ -276,7 +327,11 @@ def _ddp_and_fit(
         logger,
         filename,
         world_size,
-        port=12355
+        linear_probe,
+        fine_tune,
+        port=12355,
+        num_heads=1,
+        val_on_heads=None
     ):
     #model = nn.SyncBatchNorm.convert_sync_batchnorm(model)
     model = setup_ddp(rank, world_size, model, port)
@@ -303,7 +358,11 @@ def _ddp_and_fit(
         unmap_criterion=unmap_criterion,
         logger=logger,
         filename=filename,
-        ddp_enabled=True
+        ddp_enabled=True,
+        num_heads=num_heads,
+        linear_probe=linear_probe,
+        fine_tune=fine_tune,
+        val_on_heads=val_on_heads
     )
     dist.destroy_process_group()
 
@@ -319,9 +378,19 @@ def _fit(
         unmap_criterion,
         logger: Logger,
         filename: str,
-        ddp_enabled: bool
+        ddp_enabled: bool,
+        num_heads: int,
+        linear_probe=False,
+        fine_tune=False,
+        val_on_heads: Union[None, list]=None
     ):
-    optimizer = configure_adamw(model, lr=learning_rate)
+
+
+    if linear_probe:
+        optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learning_rate)
+    else:
+        optimizer = configure_adamw(model, lr=learning_rate)
+
     scheduler: torch.optim.lr_scheduler.SequentialLR = make_warmupCAWR(
         optimizer=optimizer,
         warmup_steps=int(len(train_gen) * 0.25),
@@ -333,6 +402,8 @@ def _fit(
     best_val_score = -1
 
     for epoch in range(nr_epochs):
+        before = datetime.now()
+        print(f'\nEpoch {epoch}: start training at {before}')
         if ddp_enabled:
             train_gen.sampler.set_epoch(epoch)
         train_log_payload = _train_epoch(
@@ -342,7 +413,13 @@ def _fit(
             optimizer,
             scheduler,
             criterion,
-            unmap_criterion)
+            unmap_criterion,
+            linear_probe,
+            fine_tune,
+            num_heads)
+        after = datetime.now()
+        print(f'Epoch {epoch}: stop training at {after}')
+        print(f'Epoch {epoch}: train duration {(after - before).total_seconds()}')
 
         if train_log_payload is not None and (not ddp_enabled or rank == 0):
             logger.log(train_log_payload, step=epoch)
@@ -352,27 +429,40 @@ def _fit(
         # For synchronous loop breaking
         stop_early = torch.zeros(1).to(rank)
 
+        epoch_val_sum = 0.0
+
         if not ddp_enabled or rank == 0:
             logger.log({'lr': scheduler.get_last_lr()[0]})
             predictions, true = val_res
+            del val_res
             predictions, true = torch.cat(predictions).cpu(), torch.cat(true).cpu()
-            predictions, true = predictions[..., 0].flatten().numpy(), true[..., 0].flatten().numpy()
-            val_log_payload = compute_metrics(
-                predictions,
-                true,
-                logspace_input=val_gen.dataset.logspace
-            )
-            val_log_payload = prepend_to_keys(val_log_payload, 'val/')
-            logger.log(val_log_payload, step=epoch)
-            print(f'Epoch {epoch} - {datetime.now()}')
-            if train_log_payload is not None:
-                print(f'\tTrain loss: {train_log_payload["train/loss"]}')
-            print(f'\tVal pearson r: {val_log_payload["val/pearson_r"]}')
-            print('-----------------------------------------')
-            if val_log_payload['val/pearson_r'] > best_val_score:
+            start_head = num_heads-1 if (linear_probe or fine_tune) else 0
+            if start_head == 0 and val_on_heads is not None:
+                start_head = val_on_heads[0]
+                # TODO verify 
+            for head in range(start_head, num_heads):
+                predictions_head = predictions[..., head].flatten().numpy()
+                if linear_probe or fine_tune:
+                    true_head =  true[..., 0].flatten().numpy()
+                else:
+                    true_head =  true[..., head].flatten().numpy()
+                val_log_payload = compute_metrics(
+                    predictions_head,
+                    true_head,
+                    logspace_input=val_gen.dataset.logspace
+                )
+                val_log_payload = prepend_to_keys(val_log_payload, 'val/')
+                logger.log(val_log_payload, step=epoch)
+                print(f'Epoch {epoch} - {datetime.now()}')
+                if train_log_payload is not None:
+                    print(f'\tTrain loss: {train_log_payload["train/loss"]}')
+                print(f'\tVal pearson r: {val_log_payload["val/pearson_r"]}')
+                print('-----------------------------------------')
+                epoch_val_sum += val_log_payload['val/pearson_r']
+            if(epoch_val_sum / (num_heads - start_head)) > best_val_score:
                 best_val_score = val_log_payload['val/pearson_r']
                 logger.save_model(model, filename)
-                 # handle early stopping
+                # handle early stopping
                 no_improvement_for = 0
             else:
                 no_improvement_for += 1
@@ -393,14 +483,24 @@ def _fit(
     print('Completed training!')
 
 
-def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion):
-    model.train()
+def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_criterion, linear_probe, fine_tune, num_heads=1):
+    if linear_probe:
+        model.eval()
+        model.core.heads[-1].train()
+    elif fine_tune:
+        model.train()
+        for head in model.core.heads[:-1]:
+            head.eval() 
+    else:
+        model.train()
 
     train_unmap = unmap_criterion is not None
 
     if rank == 0:
         # pbar if on rank 0
         train_gen = tqdm(train_gen)
+    
+    start_head = num_heads-1 if linear_probe or fine_tune else 0
 
     for X_i, m_i, y_i in train_gen:
         X_i = X_i.to(rank)
@@ -413,12 +513,23 @@ def _train_epoch(rank, model, train_gen, optimizer, scheduler, criterion, unmap_
             m_len = output_m_i.shape[1]
             unmap_loss = unmap_criterion(output_m_i, m_i[:, :m_len])
 
-            base_loss = criterion(output, y_i)
+            # only use the last head
+            base_loss = 0
+            for head in range(start_head, num_heads):
+                if linear_probe or fine_tune:
+                    base_loss += criterion(output[head], y_i)
+                else:
+                    base_loss += criterion(output[head], y_i[..., head:head+1])
             loss = base_loss + unmap_loss
         else:
             output = model(X_i)
-            loss = criterion(output, y_i)
-
+            loss = 0
+            for head in range(start_head, num_heads):
+                if linear_probe or fine_tune:
+                    print("Training head ", head, "out of",  len(output))
+                    loss += criterion(output[head], y_i)
+                else:
+                    loss += criterion(output[head], y_i[..., head:head+1])
         loss.backward()
         optimizer.step()
         scheduler.step()
@@ -455,31 +566,32 @@ def _predict(model, gen, rank, ddp_enabled):
     predictions = []
     true = []
 
-    for i, (X_i, _, y_i) in enumerate(gen):
-        X_i = X_i.to(rank)
-        y_i = y_i.to(rank)
-        
+    with torch.no_grad():
+        for i, (X_i, _, y_i) in enumerate(gen):
+            X_i = X_i.to(rank)
+            y_i = y_i.to(rank)
 
-        with torch.no_grad():
-            p_i = model(X_i)
+            with torch.no_grad():
+                p_i = model(X_i)
+                p_i = torch.cat(p_i, dim=-1)
 
-        if margin_size is not None:
-            y_i.flatten()
-            p_i = p_i[..., trim:-trim, :]
-        
-        y_i, p_i = y_i.contiguous(), p_i.contiguous()
-        if ddp_enabled:
-            all_predictions = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())]
-            all_true = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())]
-            dist.all_gather(all_predictions, p_i)
-            dist.all_gather(all_true, y_i)
-
-            if rank == 0:
-                predictions.extend(all_predictions)
-                true.extend(all_true)
-        else:
-            predictions.append(p_i)
-            true.append(y_i)
+            if margin_size is not None:
+                y_i.flatten()
+                p_i = p_i[..., trim:-trim, :]
+            
+            y_i, p_i = y_i.contiguous(), p_i.contiguous()
+            if ddp_enabled:
+                all_predictions = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())]
+                all_true = [torch.zeros_like(y_i) for _ in range(dist.get_world_size())]
+                dist.all_gather(all_predictions, p_i)
+                dist.all_gather(all_true, y_i)
+
+                if rank == 0:
+                    predictions.extend([t.detach().cpu() for t in all_predictions])
+                    true.extend([t.detach().cpu() for t in all_true])
+            else:
+                predictions.append(p_i.detach().cpu())
+                true.append(y_i.detach().cpu())
 
     print(i)
     return predictions, true