Extraction of linearized models from pre-trained networks via knowledge distillation

Fumito Kimura; Jun Ohkubo

cs.LGApr 8, 2026

Extraction of linearized models from pre-trained networks via knowledge distillation

Authors: Fumito Kimura, Jun Ohkubo

Organizations: Graduate School of Science and Engineering, Saitama University, Sakura, Saitama, 338–8570 Japan

Abstract

Recent developments in hardware, such as photonic integrated circuits and optical devices, are driving demand for research on constructing machine learning architectures tailored for linear operations. Hence, it is valuable to explore methods for constructing learning machines with only linear operations after simple nonlinear preprocessing. In this study, we propose a framework to extract a linearized model from a pre-trained neural network for classification tasks by integrating Koopman operator theory with knowledge distillation. Numerical demonstrations on the MNIST and the Fashion-MNIST datasets reveal that the proposed model consistently outperforms the conventional least-squares-based Koopman approximation in both classification accuracy and numerical stability.

Figures & tables

Figure 1: The naive method for extracting a linearized model from a pre-trained network. The snapshot pair for the EDMD algorithm consists of the first and last nodes of the hidden layer.

Figure 2: A difference between outputs on classification tasks and those in the Koopman operator method. In conventional classification tasks, we often use one-hot vectors as outputs, as depicted in (a). On the other hand, since the Koopman operator method, as denoted in Sec. 2.2, employs the least-squares method, the final output vector is not necessarily one-hot, as depicted in (b).

Figure 3: The architecture of the proposed model. In the proposed model, we first compress the input vectors using PCA before constructing the Koopman matrix. Furthermore, we incorporate not only the output of the teacher model but also the one-hot vector representing the original output into the learning process.

Number of epochs	10
Batch size	32
Optimization algorithm	AdaDelta
Hyperparameter for AdaDelta $\rho$	0.9
Hyperparameter for AdaDelta weight_decay	0.0001
Initial learning rate	1.0
Decay of learning rate	$\times$ 0.75 per epoch

Table 1: Training settings of the neural networks.

Number of epochs	10
Batch size	32
Optimization algorithm	AdaDelta
Hyperparameter for AdaDelta $\rho$	0.9
Hyperparameter for AdaDelta weight_decay	0.0001
Initial learning rate	1.0
Decay of learning rate	$\times$ 0.75 per epoch

Table 2: Training settings of the Koopman matrix in the proposed model.

Model	Dictionary (# of dictionary functions)	Test accuracy (%)
Naive method	Monomials up to 2nd degree (231)	93.96 (0.42)
Naive method (PCA)	Monomials up to 2nd degree (231)	93.25 (0.30)
Proposed model	Monomials up to 2nd degree (231)	95.98 (0.07)
Naive method	Monomials up to 3rd degree (1771)	95.89 (0.19)
Naive method (PCA)	Monomials up to 3rd degree (1771)	95.86 (0.17)
Proposed model	Monomials up to 3rd degree (1771)	96.77 (0.16)

Table 3: Comparison test accuracy with the naive method on MNIST.

Figure 4: Comparison of test accuracy between the proposed model and the naive method with PCA using ResNet18 on the MNIST dataset. (a) and (b) corresponds to results with 20 and 40 principal components, respectively.

Figure 5: Comparison of test accuracy between the proposed model and the naive method with PCA using ResNet18 on the Fashion-MNIST dataset. (a) and (b) corresponds to results with 20 and 40 principal components, respectively.

Koopman Operator Learning

Koopman operator learning builds finite-dimensional linear models of nonlinear dynamics by finding useful observable representations from data.

See all 79 papers

Spectral Distillation: From Nonlinear Dynamics to Linear State-Space Models
Aug 5, 2026Liane Galanti, Devan Shah, Shlomo Fortgang +1Dynamical Systems Nonlinear System Identification
Princeton University
Closed-Form Linear-Probe Dataset Distillation for Pre-trained Vision Models
May 8, 2026Bincheng Peng, Guang Li, Ping Liu +2Linear Probing Dataset Distillation
Hokkaido University · University of Nevada, Reno
Data-driven Koopman mode approximation: A neural power iteration algorithm
Aug 27, 2026Guillaume O. Berger, Raphaël M. JungersKoopman Operator Learning Spectral Methods
ICTEAM, UCLouvain, Louvain-la-Neuve, Belgium

Aug 5, 2026J/K move · Enter open · S save

Spectral Distillation: From Nonlinear Dynamics to Linear State-Space Models

Liane Galanti, Devan Shah, Shlomo Fortgang, Elad Hazan

Princeton University

Dynamical Systems Nonlinear System Identification Spectral Methods Knowledge Distillation

Open paper AskPDF arXiv

Can nonlinear dynamical systems be learned through a compact linear state-space representation, without directly solving a non-convex system-identification problem? We give a provable pipeline for doing so. Starting from observations of an unknown nonlinear dynamical system, we first learn an implicit spectral predictor using Observation Spectral Filtering (OSF), a convex method that competes with the best linear observer for the system. We then apply spectral-to-LDS distillation to convert this predictor into an explicit recurrent linear dynamical system. Our main theorem shows that the average prediction error of the distilled LDS decomposes into an exponentially-small distillation term and the OSF learning term governed by the Luenberger complexity of the best observer. The guarantee is dimension-free: it depends on observer complexity rather than on the latent dimension needed to represent the nonlinear system. To our knowledge, this yields the first end-to-end provable method for extracting a best-in-hindsight LDS representation of nonlinear dynamics through convex learning followed by provable distillation. Experiments on linear LDS benchmarks and MuJoCo behavior cloning show that the train-then-distill pipeline produces compact LDS predictors that match or outperform directly trained baselines.

Abstract

Koopman Operator Learning

Explore similar work

Spectral Distillation: From Nonlinear Dynamics to Linear State-Space Models

Closed-Form Linear-Probe Dataset Distillation for Pre-trained Vision Models

Data-driven Koopman mode approximation: A neural power iteration algorithm