ニューラル ODE を使用した動的システムのモデル化

この例では、ニューラル常微分方程式 (ODE) を使用して、物理的なシステムのダイナミクスをニューラルネットワークに学習させる方法を説明します。

ニューラル ODE [1] は、ODE の解によって定義される深層学習演算です。具体的には、ニューラル ODE は、どのようなアーキテクチャでも使用できる演算で、与えられた入力に対する出力を ODE の数値解として定義します。

$y^{'} = f (t, y, θ)$

ただし、時間ホライズンを $(t_{0}, t_{1})$ 、初期状態を $y (t_{0}) = y_{0}$ とします。ODE の右辺にある $f (t, y, θ)$ は、学習プロセスにおいてモデルが学習する一連の学習可能なパラメーター $θ$ に依存します。この例の $f (t, y, θ)$ は、全結合演算と非線形活性化を含むモデル関数を使用してモデル化されています。初期状態 $y_{0}$ には、この例のようにアーキテクチャ全体の入力が使用されるか、前の演算の出力が使用されます。

この例では、次の ODE で表される物理的なシステムのダイナミクス $x$ を、ニューラル ODE を使用してニューラルネットワークに学習させる方法を説明します。

$x^{'} = A x$ ,

ここで、 $A$ は 2 行 2 列の行列です。

この例のニューラルネットワークは、初期状態を入力として受け取り、学習済みのニューラル ODE モデルを使用して ODE の解を計算します。

ニューラル ODE 演算は、初期状態が与えられると、ODE モデルの解を出力します。この例では、全結合層、tanh 層、および別の全結合層を含むブロックを ODE モデルとして指定します。

この例では、ドルマン・プリンス法による陽的ルンゲ・クッタ (4, 5) のペア [2] を使用して、モデルを定義する ODE の解を数値的に求めます。バックワードパスでは、自動微分を使用し、ODE ソルバーの各演算を逆伝播させて、学習可能なパラメーター $θ$ を学習させます。

学習済みの関数 $f (t, y, θ)$ を右辺として使用し、同じモデルの解を計算して追加の初期状態を求めます。

ターゲットダイナミクスのデータの合成

x0 を初期状態とする線形 ODE モデル $x^{'} = A x$ としてターゲットダイナミクスを定義し、ode45 を使用して時間間隔 [0 15] における数値解 xTrain を計算します。グラウンドトゥルースデータを正確に計算するため、ode45 数値ソルバーの相対許容誤差を $10^{- 7}$ に設定します。後ほど、ニューラル ODE モデルを使用して近似的なダイナミクスを学習させるために、xTrain の値をグラウンドトゥルースデータとして使用します。

x0 = [2; 0];
A = [-0.1 -1; 1 -0.1];
trueModel = @(t,y) A*y;

numTimeSteps = 2000;
T = 15;
odeOptions = odeset(RelTol=1.e-7);
t = linspace(0, T, numTimeSteps);
[~, xTrain] = ode45(trueModel, t, x0, odeOptions);
xTrain = xTrain';

学習データをプロットに可視化します。

figure
plot(xTrain(1,:),xTrain(2,:))
title("Ground Truth Dynamics") 
xlabel("x(1)") 
ylabel("x(2)")
grid on

モデルパラメーターの定義と初期化

このモデル関数は、1 つの dlode45 の呼び出しで構成され、近似的なダイナミクス $f (t, y, θ)$ によって定義された ODE の解をタイムステップ 40 で求めます。

neuralOdeTimesteps = 40;
dt = t(2);
timesteps = (0:neuralOdeTimesteps)*dt;

dlode45 の呼び出しで使用する学習可能なパラメーターを定義し、それらを変数 neuralOdeParameters にまとめます。関数 initializeGlorot は、学習可能なパラメーターのサイズ sz、および全結合演算の出力と入力の数を入力として受け取り、基となる型が single で、Glorot の初期化を使用して値が設定された dlarray オブジェクトを返します。関数 initializeZeros は、学習可能なパラメーターのサイズを入力として受け取り、基となる型が single である dlarray オブジェクトとしてパラメーターを返します。この初期化サンプル関数は、この例にサポートファイルとして添付されています。これらの関数にアクセスするには、この例をライブスクリプトとして開きます。モデル関数用の学習可能なパラメーターの初期化に関する詳細については、モデル関数の学習可能パラメーターの初期化を参照してください。

パラメーター構造体を初期化します。

neuralOdeParameters = struct;

ODE モデルの全結合演算のパラメーターを初期化します。最初の全結合演算は、サイズ stateSize のベクトルを入力として受け取り、その長さを hiddenSize まで増やします。逆に、2 番目の全結合演算は、長さ hiddenSize のベクトルを入力として受け取り、その長さを stateSize まで減らします。

stateSize = size(xTrain,1);
hiddenSize = 20;

neuralOdeParameters.fc1 = struct;
sz = [hiddenSize stateSize];
neuralOdeParameters.fc1.Weights = initializeGlorot(sz, hiddenSize, stateSize);
neuralOdeParameters.fc1.Bias = initializeZeros([hiddenSize 1]);

neuralOdeParameters.fc2 = struct;
sz = [stateSize hiddenSize];
neuralOdeParameters.fc2.Weights = initializeGlorot(sz, stateSize, hiddenSize);
neuralOdeParameters.fc2.Bias = initializeZeros([stateSize 1]);

モデルの学習可能なパラメーターを表示します。

neuralOdeParameters.fc1

ans = struct with fields:
    Weights: [20×2 dlarray]
       Bias: [20×1 dlarray]

neuralOdeParameters.fc2

ans = struct with fields:
    Weights: [2×20 dlarray]
       Bias: [2×1 dlarray]

ニューラル ODE モデルの定義

この例の ODE モデルのセクションにリストされている関数 odeModel を作成します。この関数は、時間入力 (これは使用されません)、対応する解、および ODE 関数のパラメーターを入力として受け取ります。この関数は、パラメーターによって与えられた重みとバイアスを使用して、全結合演算、tanh 演算、および別の全結合演算を入力データに適用します。

モデルの関数の定義

この例のモデル関数のセクションにリストされている関数 model を作成します。この関数は深層学習モデルの出力を計算します。関数 model は、モデルパラメーターと入力データを入力として受け取ります。この関数は、ニューラル ODE の解を出力します。

モデル損失関数の定義

この例のモデル損失関数のセクションにリストされている関数 modelLoss を作成します。この関数は、モデルパラメーター、入力データのミニバッチと対応するターゲットを入力として受け取り、学習可能なパラメーターについての損失とその損失の勾配を返します。

学習オプションの指定

Adam 最適化のオプションを指定します。

gradDecay = 0.9;
sqGradDecay = 0.999;
learnRate = 0.002;

ミニバッチサイズを 200 として、1200 回反復して学習させます。

numIter = 1200;
miniBatchSize = 200;

50 回の反復ごとに、学習済みのダイナミクスを求め、グラウンドトゥルースとともにその結果を位相ダイアグラムに表示して学習パスを示します。

plotFrequency = 50;

カスタム学習ループを使用したモデルの学習

Adam ソルバーのパラメーター averageGrad および averageSqGrad を初期化します。

averageGrad = [];
averageSqGrad = [];

TrainingProgressMonitor オブジェクトを初期化します。monitor オブジェクトを作成するとタイマーが開始されるため、学習ループに近いところでオブジェクトを作成するようにしてください。

monitor = trainingProgressMonitor(Metrics="Loss",Info=["Iteration","LearnRate"],XLabel="Iteration");

カスタム学習ループを使用してネットワークに学習させます。

それぞれの反復で次を行います。

この例のミニバッチ関数の作成のセクションにリストされている関数 createMiniBatch を使用して、合成データからデータのミニバッチを構築します。
この例のモデル損失関数のセクションにリストされている関数 dlfeval および関数 modelLoss を使用して、モデルの損失と勾配を評価します。
関数 adamupdate を使用してモデルパラメーターを更新します。
学習の進行状況プロットを更新します。

numTrainingTimesteps = numTimeSteps;
trainingTimesteps = 1:numTrainingTimesteps;
plottingTimesteps = 2:numTimeSteps;

iteration = 0;

while iteration < numIter && ~monitor.Stop
    iteration = iteration + 1;

    % Create batch
    [X, targets] = createMiniBatch(numTrainingTimesteps, neuralOdeTimesteps, miniBatchSize, xTrain);

    % Evaluate network and compute loss and gradients
    [loss,gradients] = dlfeval(@modelLoss,timesteps,X,neuralOdeParameters,targets);

    % Update network
    [neuralOdeParameters,averageGrad,averageSqGrad] = adamupdate(neuralOdeParameters,gradients,averageGrad,averageSqGrad,iteration,...
        learnRate,gradDecay,sqGradDecay);

    % Plot loss
    recordMetrics(monitor,iteration,Loss=loss);

    % Plot predicted vs. real dynamics
    if mod(iteration,plotFrequency) == 0  || iteration == 1

        % Use ode45 to compute the solution 
        y = dlode45(@odeModel,t,dlarray(x0),neuralOdeParameters,DataFormat="CB");

        plot(xTrain(1,plottingTimesteps),xTrain(2,plottingTimesteps),"r--")

        hold on
        plot(y(1,:),y(2,:),"b-")
        hold off

        xlabel("x(1)")
        ylabel("x(2)")
        title("Predicted vs. Real Dynamics")
        legend("Training Ground Truth", "Predicted")

        drawnow
    end
    updateInfo(monitor,Iteration=iteration,LearnRate=learnRate);
    monitor.Progress = 100*iteration/numIter;
end

モデルの評価

モデルを使用して、さまざまな初期状態について近似解を計算します。

モデルの学習に使用したのとは異なる 4 つの新しい初期状態を定義します。

tPred = t;

x0Pred1 = sqrt([2;2]);
x0Pred2 = [-1;-1.5];
x0Pred3 = [0;2];
x0Pred4 = [-2;0];

4 つの新しい初期状態について、ode45 を使用して ODE の真のダイナミクスを数値的に求めます。

[~, xTrue1] = ode45(trueModel, tPred, x0Pred1, odeOptions);
[~, xTrue2] = ode45(trueModel, tPred, x0Pred2, odeOptions);
[~, xTrue3] = ode45(trueModel, tPred, x0Pred3, odeOptions);
[~, xTrue4] = ode45(trueModel, tPred, x0Pred4, odeOptions);

学習済みのニューラル ODE ダイナミクスを使用して、ODE の解を数値的に求めます。

xPred1 = dlode45(@odeModel,tPred,dlarray(x0Pred1),neuralOdeParameters,DataFormat="CB");
xPred2 = dlode45(@odeModel,tPred,dlarray(x0Pred2),neuralOdeParameters,DataFormat="CB");
xPred3 = dlode45(@odeModel,tPred,dlarray(x0Pred3),neuralOdeParameters,DataFormat="CB");
xPred4 = dlode45(@odeModel,tPred,dlarray(x0Pred4),neuralOdeParameters,DataFormat="CB");

予測の可視化

この例の真の解と予測解をプロットする関数のセクションにリストされている関数 plotTrueAndPredictedSolutions を使用して、さまざまな初期状態における予測解をグラウンドトゥルース解とともに可視化します。

figure
subplot(2,2,1)
plotTrueAndPredictedSolutions(xTrue1, xPred1);
subplot(2,2,2)
plotTrueAndPredictedSolutions(xTrue2, xPred2);
subplot(2,2,3)
plotTrueAndPredictedSolutions(xTrue3, xPred3);
subplot(2,2,4)
plotTrueAndPredictedSolutions(xTrue4, xPred4);

補助関数

モデル関数

予測に使用されるニューラルネットワークを定義する関数 model は、1 つのニューラル ODE 呼び出しで構成されます。この関数は、各観測値について、長さ stateSize のベクトルを右辺として受け取り、また数値解が出力されるタイミングを定義する時間点 tspan のベクトルを受け取ります。この長さのベクトルは、求解対象となる ODE の学習可能な右辺 $f (t, y, θ)$ を表す関数 odeModel を使用して ODE の解を数値的に求める際の初期状態として使用されます。学習済みのシステムは自律しているため、初期状態にかかわらず、この関数は各観測値についてベクトル tspan を使用します。つまり、関数 odeModel は時間に明示的に依存しません。

function X = model(tspan,X0,neuralOdeParameters)

X = dlode45(@odeModel,tspan,X0,neuralOdeParameters,DataFormat="CB");

end

ODE モデル

関数 odeModel は学習可能な右辺で、dlode45 の呼び出し時に使用されます。これは、サイズ stateSize のベクトルを入力として受け取り、長さが hiddenSize となるようにそれを拡張し、非線形関数 tanh を適用します。その後、この関数は、長さが stateSize となるようにこのベクトルを再度圧縮します。

function y = odeModel(~,y,theta)

y = tanh(theta.fc1.Weights*y + theta.fc1.Bias);
y = theta.fc2.Weights*y + theta.fc2.Bias;

end

モデル損失関数

この関数は、ベクトル tspan、一連の初期状態 X0、学習可能なパラメーター neuralOdeParameters、およびターゲットシーケンス targets を入力として受け取ります。これは、関数 model を使用して予測を計算し、与えられたターゲットシーケンスと予測結果を比較します。最後に、ニューラル ODE の学習可能なパラメーターについて、損失とその損失の勾配を計算します。

function [loss,gradients] = modelLoss(tspan,X0,neuralOdeParameters,targets)

% Compute predictions.
X = model(tspan,X0,neuralOdeParameters);

% Compute L1 loss.
loss = l1loss(X,targets,NormalizationFactor="all-elements",DataFormat="CBT");

% Compute gradients.
gradients = dlgradient(loss,neuralOdeParameters);

end

ミニバッチ関数の作成

関数 createMiniBatch は、ターゲットダイナミクスの観測値のバッチを作成します。これは、グラウンドトゥルースデータのタイムステップ数の合計 numTimesteps、観測値ごとに返される連続タイムステップの数 numTimesPerObs、観測値の数 miniBatchSize、およびグラウンドトゥルースデータ X を入力として受け取ります。

function [x0, targets] = createMiniBatch(numTimesteps,numTimesPerObs,miniBatchSize,X)

% Create batches of trajectories.
s = randperm(numTimesteps - numTimesPerObs, miniBatchSize);

x0 = dlarray(X(:, s));
targets = zeros([size(X,1) miniBatchSize numTimesPerObs]);

for i = 1:miniBatchSize
    targets(:, i, 1:numTimesPerObs) = X(:, s(i) + 1:(s(i) + numTimesPerObs));
end

end

真の解と予測解をプロットする関数

関数 plotTrueAndPredictedSolutions は、真の解 xTrue、学習済みのニューラル ODE モデルによって計算された近似解 xPred、および対応する初期状態 x0Str を入力として受け取ります。これは、真の解と予測解との誤差を計算して位相ダイアグラムにプロットします。

function plotTrueAndPredictedSolutions(xTrue,xPred)

xPred = squeeze(xPred)';

err = mean(abs(xTrue(2:end,:) - xPred), "all");

plot(xTrue(:,1),xTrue(:,2),"r--",xPred(:,1),xPred(:,2),"b-",LineWidth=1)

title("Absolute Error = " + num2str(err,"%.4f"))
xlabel("x(1)")
ylabel("x(2)")

xlim([-2 3])
ylim([-2 3])

legend("Ground Truth","Predicted")

end

[1] Chen, Ricky T. Q., Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. “Neural Ordinary Differential Equations.” Preprint, submitted December 13, 2019. https://arxiv.org/abs/1806.07366.

[2] Shampine, Lawrence F., and Mark W. Reichelt. “The MATLAB ODE Suite.” SIAM Journal on Scientific Computing 18, no. 1 (January 1997): 1–22. https://doi.org/10.1137/S1064827594276424.