github links:
Cellbox
SFENO

Introduction to SFENO

SFENO (Structured Framework for Exploring Network ODEs) is a flexible framework designed to model and simulate biological networks using ordinary differential equations (ODEs). It was developed as an improvement over the existing Cellbox model, addressing key limitations in flexibility, hardware compatibility, and experimental design.

Key Limitations of Cellbox

The original Cellbox model had several limitations:

  1. Tightly Coupled Architecture: The ODE equations were deeply embedded within the Cellbox model itself, making it difficult to modify or experiment with different mathematical approaches without rewriting the entire class.

  2. Limited Learning Methods: The learning code was designed specifically for the approach described in the paper, lacking flexibility for applying various learning techniques.

  3. Hardware Constraints: Being implemented in TensorFlow, the code was not easily adaptable to different hardware environments without significant modifications.

SFENO’s Architectural Improvements

Structural differences between Cellbox and SFENO

SFENO addresses these limitations through several key architectural changes:

1. Modular Architecture

SFENO

The architecture clearly separates different components:

  • LitSfeno: The top-level wrapper that integrates PyTorch Lightning functionality
  • GeneralModel: A middle-layer that handles ODE solving and model coordination
  • User-defined ODE modules: Customizable components that define how the state changes

This separation allows researchers to modify just the ODE equations without changing the rest of the system, significantly enhancing experimental flexibility.

Here’s an example of the LitSfeno class showing how it wraps the model and integrates PyTorch Lightning:

class LitSfeno(pl.LightningModule):
    def __init__(self, config, args, ud_class_nm):
        super().__init__()
        
        # Create model using the user-defined ODE class
        self.model = GeneralModel.build_model(config=config,
                                              args=args,
                                              user_defined_dxdt=ud_class_nm,
                                              master=self)
        
        self.lr = 0.001
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        
    def forward(self, t, y0, arg_dict):
        return self.model(self, t, y0, arg_dict)
        
    def get_device(self):
        return self.device
        
    def training_step(self, batch, batch_idx):
        inputs, targets = batch
        y0, additional_args = inputs
        
        outputs = self.model(t=0, y0=y0, arg_dict=additional_args)
        loss = self.loss_fn(outputs, targets)
        
        self.log("train_loss", loss, on_step=False, on_epoch=True)
        return loss

2. PyTorch Lightning Integration

SFENO replaces TensorFlow with PyTorch and introduces PyTorch Lightning, providing:

  • Hardware agnostic operation across CPUs, GPUs, and TPUs without code changes
  • Support for distributed training (DDP - Distributed Data Parallel)
  • Built-in callbacks at appropriate points during model training
  • Automatic tensor management across devices

The following code demonstrates how SFENO leverages PyTorch Lightning for hardware-agnostic training:

# Example of setting up training with different hardware configurations
if pargs.ddp:  # Distributed Data Parallel (multi-GPU)
    trainer = Trainer(devices=pargs.num_gpu,
                      accelerator="gpu",
                      max_epochs=config.config['max_epoch'],
                      strategy="ddp_find_unused_parameters_false")
else:  # Single GPU
    trainer = Trainer(devices=pargs.num_gpu,
                      accelerator="gpu",
                      max_epochs=config.config['max_epoch'])

# Train the model with the same code regardless of hardware configuration
trainer.fit(lit_sfeno, datamodule=lit_datamodule)

3. Flexible Masking System

SFENO implements a generalized masking system for network constraints:

  • Four possible interaction types between nodes: no influence, bidirectional influence, positive-only influence, or negative-only influence
  • User-defined network structure information can be directly incorporated
  • The StructInfoLayer class manages these constraints through an elegant filtering mechanism

Here’s how the masking system is implemented:

class StructInfoLayer(nn.Module):
    def __init__(self, config, master):
        super().__init__()
        self.master = [master]
        
        # Initialize adjacency matrix
        if config.adj_mat==1:
            self.adj_mat = np.ones((config.config["n_x"],config.config["n_x"]))
        else:
            self.adj_mat = config.adj_mat
            
        self.node_names = config.node_nm_ls
        self.node_num = len(self.adj_mat)
        
        # Network weight parameters
        self.register_parameter(
            'W', nn.parameter.Parameter(torch.zeros((self.node_num, self.node_num), 
                                                    requires_grad=True))
        )
        
        # Define different parameter transformation functions
        def zero_param(x):
            return x * 0
            
        def R_param(x):
            return x
            
        def R_pos_param(x):
            return F.softplus(x)  # Positive-only influence
            
        def R_ngt_param(x):
            return(-1) * F.softplus(x)  # Negative-only influence
            
        self.adj_val2get_param_func = {
            0: zero_param,      # No influence
            1: R_param,         # Bidirectional influence
            2: R_pos_param,     # Positive-only influence
            3: R_ngt_param      # Negative-only influence
        }

    def load_mask(self):
        # Create masks for each interaction type
        self.mask = torch.zeros(
            (self.node_num, self.node_num),
            requires_grad=False).type_as(next(self.parameters()))
            
        for adj_val, func in self.adj_val2get_param_func.items():
            mask = func(self.W * self.adj_val2mask[adj_val].type_as(next(self.parameters())))
            self.mask = self.mask + mask

4. Easy Experimental Configuration

The framework supports:

  • Hyperparameter scheduling across training stages
  • Custom loss function definitions within user-defined ODE modules
  • Automatic network weight saving at specified epochs
  • Performance tracking and evaluation

Example of implementing a custom ODE function with hyperparameter scheduling:

class Example_Cellbox(GeneralODE):
    def __init__(self, config, master):
        super().__init__(config, master)
        
        # Initialize network model with structural information
        self.net_model = StructInfoLayer(config, master=self)
        self.net_model.init_params()
        self.net_model.load_mask()
        
        # Register learnable parameters
        self.register_parameter(
            'epsilon', nn.parameter.Parameter(
                F.softplus(torch.ones((config.config["n_x"]),
                           dtype=torch.float32, requires_grad=True)))
        )
        
        self.register_parameter(
            'alpha', nn.parameter.Parameter(
                F.softplus(torch.ones((config.config["n_x"]),
                           dtype=torch.float32, requires_grad=True)))
        )
        
        # Setup hyperparameters
        self.hyper_params = {
            'l1': 1e-4,
            'l2': 1e-4
        }
        
        self.mse = torch.nn.MSELoss()
        self.stages = config.config['stages']
        
    def forward(self, t0, x, arg_dict):
        '''Differential equation implementation'''
        y = self.net_model(x) + arg_dict['mu']
        y = torch.tanh(y)
        y = self.epsilon * y
        y = y - self.alpha * x
        return y
        
    # Custom loss function
    def loss_fn(self, output, target):
        l1 = torch.sum(torch.abs(self.net_model.W))
        l2 = torch.sum(torch.pow(self.net_model.W, 2))
        
        mse_loss = self.mse(torch.squeeze(output), torch.squeeze(target))
        
        total_loss = self.hyper_params['l1'] * l1 + self.hyper_params['l2'] * l2 + mse_loss
        return total_loss
        
    # Hyperparameter scheduling based on training stage
    def training_epoch_end(self, outputs):
        cur_epoch = self.master[0].master[0].current_epoch
        if not self.stages:
            return
        else:
            idx_stage, change_hype = self.epoch2stage_idx(cur_epoch)
            stage = self.stages[idx_stage]
            if change_hype:
                # Update hyperparameters based on stage configuration
                lr = stage['lr_val'] if 'lr_val' in stage else self.config.config['default']['lr_val']
                l1 = stage['l1'] if 'l1' in stage else self.config.config['default']['l1']
                l2 = stage['l2'] if 'l2' in stage else self.config.config['default']['l2']
                nt = stage['nT'] if 'nT' in stage else self.config.config['default']['nT']
                
                print('Changing Hyperparameters')
                print(f'lr {lr}, l1 {l1}, l2 {l2}, nt {nt}')
                
                # Apply new hyperparameters
                self.master[0].master[0].lr = lr
                self.hyper_params['l1'] = l1
                self.hyper_params['l2'] = l2
                self.master[0].nt = nt
                
                self.master[0].master[0].configure_optimizers()
                self.save_net_weight(cur_epoch)

Implementation Details

Core Components

  1. LitSfeno Class: Wraps the model with PyTorch Lightning functionality, handling training, validation, and testing steps while maintaining device compatibility.

  2. GeneralModel Class: Manages the ODE solving process and coordinates between the Lightning module and the user-defined ODE equations.

class GeneralModel_v2(nn.Module):
    def __init__(self, config, dt, ud_dxdt, master):
        super().__init__()
        # Initialize user defined module
        self.master = [master]
        self.dxdt = ud_dxdt(config, master=master)
        
        self.ode_solver = None
        self.dt = dt
        self.nt = 200
        self.config = config
        
    def forward(self, t, y0, arg_dict):
        if self.nt == 0:
            raise Exception('Number of iteration nt on model is zero')
            
        batch_size = y0.shape[0]
        
        # Solve the ODE system
        ys = self.ode_solver.general_solve(
            ode=self.dxdt,
            dt=self.dt,
            nt=self.nt,
            t0=t,
            x=y0,
            arg_dict=arg_dict
        )
        
        # Return the final state
        yhat = ys[-1].reshape(batch_size, -1, 1)
        return yhat

  1. Example_Cellbox ODE Class: Demonstrates how to implement custom ODE equations, including parameters such as epsilon and alpha, as well as custom loss functions.

  2. StructInfoLayer: Handles network structural constraints through sophisticated masking techniques.

Hardware Compatibility

SFENO achieves hardware agnosticism by following key design principles:

  • Avoiding explicit device specifications (.cuda() or .to())
  • Using PyTorch Lightning’s device management system
  • Properly registering parameters and buffers for automatic device placement

For instance, when creating new tensors, SFENO uses this approach to maintain device compatibility:

# Create a device-agnostic tensor
self.mask = torch.zeros(
    (self.node_num, self.node_num),
    requires_grad=False).type_as(next(self.parameters()))

Example of ODE solver implementation that works across different devices:

class HeunSolver:
    """Heun's ODE solver"""
    def __init__(self):
        pass

    def general_solve(self, ode, dt, nt, t0, x, arg_dict):
        batch_s = x.shape[0]
        xs = []
        
        for i in range(nt):
            t = i*dt
            dxdt_current = ode(t, x, arg_dict)
            dxdt_next = ode(t + dt, x + dt * dxdt_current, arg_dict)
            x = x + dt * 0.5 * (dxdt_current + dxdt_next)
            xs.append(x)

        xs = torch.stack(xs, dim=0)
        return xs

Performance Benchmarks

The system was tested on multiple biological datasets:

  • Borisov 2009 (Cyanobacteria photosynthesis)
  • Molinelli 2013 (Immunology disease treatment)
  • Nelander 2008 (Brain tumor research)
  • Pezze 2012 (Metabolic network analysis)
  • Schliemann 2011 (Protein interaction networks)
  • Korkut 2015 (Cancer biology data)

Example of running experiments across multiple datasets:

experiment_path_ls = [
    'sfeno/datasets/borisov_2009/sfeno_data',
    'sfeno/datasets/molinelli_2013/sfeno_data',
    'sfeno/datasets/nelander_2008/sfeno_data',
    'sfeno/datasets/pezze_2012/sfeno_data',
    'sfeno/datasets/schliemann_2011/sfeno_data',
    'sfeno/datasets/korkut_2015a/sfeno_data',
]

for net_info in [False, True]:
    for data_path in experiment_path_ls:
        config.config['data_path'] = data_path
        config.load_paths_n_info()
        
        # Configure the model based on parameters
        config.config['batchsize'] = config.config['batchsize'] if pargs.num_batch==0 else pargs.num_batch
        config.config['max_epoch'] = config.config['max_epoch'] if pargs.num_epoch==0 else pargs.num_epoch
        
        # Setup data loaders
        trainloader, validloader, testloader = get_dataloaders(config)
        lit_datamodule = DataModule(trainloader, validloader, testloader)
        
        # Create model with appropriate ODE implementation
        if pargs.nn_ode:
            lit_sfeno = LitSfeno(config=config, args=args, ud_class_nm=general_model.NN_dxdt)
        else:
            lit_sfeno = LitSfeno(config=config, args=args, ud_class_nm=general_model.Example_Cellbox)
            
        # Train and evaluate
        trainer.fit(lit_sfeno, datamodule=lit_datamodule)
        pred_ls = trainer.predict(lit_sfeno, datamodule=lit_datamodule)

Initial performance scaling tests showed promising results:

  • Single GPU: 21 hours for 3000 epochs with 1000 nodes
  • 4-GPU DDP: 12 hours for the same workload (approximately 1.75x speedup)

While these results are preliminary and will require more extensive testing across different hardware configurations and model sizes, they suggest that SFENO can effectively utilize distributed computing resources.

Conclusion

SFENO offers a more flexible alternative to the original Cellbox model by providing a modular, hardware-agnostic framework for modeling biological networks. The separation of ODE equations from the core model architecture, along with PyTorch Lightning integration and the enhanced masking system, simplifies experimentation with different mathematical approaches.

The framework aims to reduce the technical barriers to modeling complex biological systems, allowing researchers to focus more on the biological questions rather than implementation details. While SFENO is still in development and has been tested on a limited number of datasets, early results suggest it could be a useful tool for systems biology research.

Researchers interested in customizing ODE models can create a class that inherits from GeneralODE, implement the forward method to define their differential equations, and specify a custom loss_fn method for training. This design provides flexibility without requiring modifications to the core architecture, potentially making it easier to explore different modeling approaches in systems biology.

github links:
Cellbox
SFENO

<
Previous Post
Building Resilient Client-Server Game Communication
>
Next Post
Beyond Hierarchy: Rethinking Visual Processing in the Brain