SIM (State-driven Implicit Models)¶

State-driven Implicit Modeling (SIM) is an advanced training methodology for implicit models, introduced in “State-driven Implicit Models”. SIM distills implicit models from pre-trained explicit networks by matching internal state representations.

Theoretical Foundation¶

A standard Implicit Model is defined by:

\[\begin{split}\begin{aligned} X &= \phi(A X + B U) \quad &\text{(Equilibrium equation)}, \\ \hat{Y} &= C X + D U \quad &\text{(Prediction equation)}, \end{aligned}\end{split}\]

SIM Training Objective:

In SIM, we assume the hidden states \(X\) are fixed and extracted from a classical neural network. We then train the parameters \(A,B,C,D\) to match the behavior of this explicit network.

Given:

Input matrix \(U \in \mathbb{R}^{p \times m}\)
Pre-activation states \(Z \in \mathbb{R}^{n \times m}\) from explicit network
Post-activation states \(X \in \mathbb{R}^{n \times m}\) from explicit network
Outputs \(\hat{Y} \in \mathbb{R}^{q \times m}\) from explicit network

SIM solves the following convex optimization problem:

\[\begin{split}\begin{aligned} & \min_{A,B,C,D} \quad f(A,B,C,D)\\ & \text{subject to:} \\ & \quad Z = AX + BU, \\ & \quad \hat{Y} = CX + DU, \\ & \quad \|A\|_\infty \leq \kappa, \end{aligned}\end{split}\]

where \(f\) is an objective function that typically includes regularization terms to promote sparsity or other desirable model properties.

Implementation Components¶

The SIM training process consists of two main phases:

State Extraction: The idl.sim.sim.SIM class contains the method to extract internal state vectors from a given neural network and formulates the optimization problem.
Convex Optimization: Various solvers are already provided in the next sections to solve the resulting optimization problem efficiently. Moreover, custom solvers can be applied by inheriting from the idl.sim.solvers.solver.BaseSolver class.

API Reference¶

class torchidl.sim.sim.SIM(activation_fn=<function relu>, kappa=0.99, atol=1e-06, skip_layers=None, standardize=False, device='cpu', dtype=torch.float32)[source]¶

SIM base class.

Parameters:

activation_fn (Callable) – Activation function used in the implicit model.
kappa (float) – Parameter to ensure convergence of the fixed-point iteration. Default is 0.99.
atol (float) – Absolute tolerance for the fixed-point iteration. Default is 1e-6.
skip_layers (int, optional) – If not None, only the last activation of each layer block will be used. The block size is controlled by the parameter. Defaults to None.
standardize (bool, optional) – Whether to standardize the input data using scipy StandardScaler. Defaults to False.
device (str or torch.device, optional) – Device to use for SIM. Defaults to “cpu”.
dtype (str or torch.dtype, optional) – Data type for SIM. Defaults to torch.float32.

__call__(input)[source]¶

Forward pass of a standard implicit model.

Parameters:: input (torch.Tensor) – Input data with shape (batch_size, input_dim).
Returns:: Output data with shape (batch_size, output_dim).
Return type:: output (torch.Tensor)

evaluate(dataloader)[source]¶

Evaluate the SIM model on the given test data.

Parameters:: dataloader (torch.utils.data.DataLoader) – Test data loader.
Returns:: Accuracy of the SIM model on the given test data.
Return type:: test_accuracy (float)

get_states(model, dataloader)[source]¶

Extract the states data (pre-activations, post-activations, inputs, outputs) from the explicit model. The dataloader should only load a small amount of data to avoid memory issues.

Parameters:

model (torch.nn.Module) – Explicit model (teacher) to extract state data.
dataloader (torch.utils.data.DataLoader) – Data loader.

Returns:

Dictionary containing the states data:

U (np.ndarray): Input data with shape (batch_size, input_dim).
Z (np.ndarray): Pre-activations with shape (batch_size, hidden_dim).
X (np.ndarray): Post-activations with shape (batch_size, hidden_dim).
Y (np.ndarray): Output data with shape (batch_size, output_dim).

Return type:

states_data (dict)

train(solver, model, dataloader=None, states_data_path=None, save_states_path=None)[source]¶

Train the SIM model.

Parameters:

solver (Callable) – Solver to use for training.
model (torch.nn.Module) – Explicit model (teacher) to extract state data.
dataloader (torch.utils.data.DataLoader, optional) – Training data loader.
states_data_path (str, optional) – Path to the states data file.
save_states_path (str, optional) – Path to save the states data file.

Example usage:

import torch
from torchidl import SIM
from torchidl import CVXSolver

explicit_model = ...
dataloader = ...

# Define the SIM model
sim = SIM(activation_fn=torch.nn.functional.relu, device="cuda", dtype=torch.float32)

# Define the solver
solver = CVXSolver()

# Train SIM
sim.train(solver=solver, model=explict_model, dataloader=dataloader)