python - なんらかの変換によってパラメータを定義する OR グラフ全体ではなくサブグラフを保持する

Question

今まで見たことのない問題に出くわしています。私はベイジアン機械学習に取り組んでおり、PyTorch のディストリビューションをよく利用しています。よくあることの 1 つは、分布のパラメーターのいくつかをパラメーターの対数で定義することです。これにより、最適化で負になることはありません (正規分布の標準偏差など)。

ただし、分布に依存しないようにするために、このパラメーターの変換を手動で再計算する必要はありません。例を介して実証するには：

次のコードは実行されません。最初の後方パスの後、パラメーターの指数関数を計算するグラフの部分は自動的に削除され、再度追加されることはありません。

import torch
import torch.nn as nn
import torch.distributions as dd

log_std = nn.Parameter(torch.Tensor([1])) # Define the log of the parameter as an nn.Parameter, this is what we want to optimise
std = torch.exp(log_std) # Define the transformation we want to apply to the parameter to using it in the distribution
mean = nn.Parameter(torch.Tensor([1])) # A normal parameter
dist = dd.Normal(loc=mean, scale=std) # Define the distribution. From here I want to ONLY refer to this, not the other variables

optim = torch.optim.SGD([log_std, mean], lr=0.01) # Standard optimiser
target = dd.Normal(5,5) # Target distribution to match

for i in range(50):
    optim.zero_grad()

    samples = dist.rsample((1000,)) # Sample our model, note no reference to log_std

    cost = -(target.log_prob(samples) - dist.log_prob(samples)).sum() # KLdivergence cost metric
    cost.backward()
    optim.step()
    print(i)
    print(log_std, mean, cost)
    print()

次のコード セットは実行されlog_stdますが、ループ内でパラメーターを明示的に参照し、ディストリビューションを再作成する必要があります。分散タイプを変更したい場合、特定のケースを考慮しないと不可能です。

import torch
import torch.nn as nn
import torch.distributions as dd

log_std = nn.Parameter(torch.Tensor([1])) # Define the log of the parameter as an nn.Parameter, this is what we want to optimise
mean = nn.Parameter(torch.Tensor([1])) # A normal parameter

optim = torch.optim.SGD([log_std, mean], lr=0.001) # Standard optimiser
target = dd.Normal(5,5) # Target distribution to match

for i in range(50):
    optim.zero_grad()

    std = torch.exp(log_std)  # Define the transformation we want to apply to the parameter to using it in the distribution
    dist = dd.Normal(loc=mean, scale=std)  # Define the distribution.

    samples = dist.rsample((1000,)) # Sample our model, note no reference to log_std

    cost = -(target.log_prob(samples) - dist.log_prob(samples)).sum() # KL divergence cost metric
    cost.backward()
    optim.step()
    print(i)
    print(mean, std, cost)
    print()

ただし、グラフが静的であるため、最初の例は Tensorflow で機能します。これを修正する方法について誰かアイデアがありますか? 関係を定義するグラフの部分のみを保持することが可能であれば、std = torch.exp(log_std)これはうまくいく可能性があります。逆方向勾配フックも試してみましたが、残念ながら、新しい勾配を適切に計算するには、パラメーター値と学習率にアクセスする必要があります。

前もって感謝します！マイケル

編集

ディストリビューションを変更する方法の例を求められました。現在機能しないコードを取得し、分布をガンマ分布に変更します。

import torch
import torch.nn as nn
import torch.distributions as dd

log_rate = nn.Parameter(torch.Tensor([1])) # Define the log of the parameter as an nn.Parameter, this is what we want to optimise
rate = torch.exp(log_std) # Define the transformation we want to apply to the parameter to usi it in the distribution
concentration = nn.Parameter(torch.Tensor([1])) # A normal parameter
dist = dd.Gamma(concentration=concentration, rate=std) # Define the distribution. From here I want to ONLY refer to this, not the other variables

optim = torch.optim.SGD([log_rate, concentration], lr=0.01) # Standard optimiser
target = dd.Gamma(5,5) # Target distribution to match

for i in range(50):
    optim.zero_grad()

    samples = dist.rsample((1000,)) # Sample our model, note no reference to log_std

    cost = -(target.log_prob(samples) - dist.log_prob(samples)).sum() # KL divergence cost metric
    cost.backward()
    optim.step()
    print(i)
    print(log_std, mean, cost)
    print()

ただし、現在動作しているコードを見ると:

import torch
import torch.nn as nn
import torch.distributions as dd

log_rate = nn.Parameter(torch.Tensor([1])) # Define the log of the parameter as an nn.Parameter, this is what we want to optimise
mean = nn.Parameter(torch.Tensor([1])) # A normal parameter

optim = torch.optim.SGD([log_rate, concentration], lr=0.001) # Standard optimiser
target = dd.Gamma(5,5) # Target distribution to match

for i in range(50):
    optim.zero_grad()

    rate = torch.exp(log_rate)  # Define the transformation we want to apply to the parameter to usi it in the distribution
    dist = dd.Gamma(concentration=concentration, rate=rate)  # Define the distribution.

    samples = dist.rsample((1000,)) # Sample our model, note no reference to log_std

    cost = -(target.log_prob(samples) - dist.log_prob(samples)).sum() # KL divergence cost metric
    cost.backward()
    optim.step()
    print(i)
    print(mean, std, cost)
    print()

そして、アルゴリズムが機能するようにするには、ループ内のコードを変更する必要があることがわかります。この小さな例では大きな問題ではありませんが、これは、心配する必要がないことが信じられないほど有益な、はるかに大きなアルゴリズムのデモンストレーションにすぎません。