library

[fredshone]

https://arxiv.org/abs/1907.00503

https://pdfs.semanticscholar.org/2522/9cb956c80d00874281f4481676d57fae394b.pdf

[fredshone]

TVAE looks good. Key additions:

• decomposed scalars
• nll loss of gaussian used for scalars, see following code:

def _loss_function(recon_x, x, sigmas, mu, logvar, output_info, factor):
    st = 0
    loss = []
    for column_info in output_info:
        for span_info in column_info:
            if span_info.activation_fn != 'softmax':
                ed = st + span_info.dim
                std = sigmas[st]
                eq = x[:, st] - torch.tanh(recon_x[:, st])
                loss.append((eq**2 / 2 / (std**2)).sum())
                loss.append(torch.log(std) * x.size()[0])
                st = ed

            else:
                ed = st + span_info.dim
                loss.append(
                    cross_entropy(
                        recon_x[:, st:ed], torch.argmax(x[:, st:ed], dim=-1), reduction='sum'
                    )
                )
                st = ed

    assert st == recon_x.size()[1]
    KLD = -0.5 * torch.sum(1 + logvar - mu**2 - logvar.exp())
    return sum(loss) * factor / x.size()[0], KLD / x.size()[0]

What the two terms mean

If we assume
[
eq \sim \mathcal{N}(0, \sigma^2),
]

then the per-sample negative log-likelihood is
[
\text{NLL}(eq;\sigma)=\frac{eq^2}{2\sigma^2} + \log\sigma + \frac{1}{2}\log(2\pi).
]

Summing over a batch of size (N):
[
\text{NLL}\text{batch} = \sum{i=1}^{N}\left(\frac{eq_i^2}{2\sigma^2} + \log\sigma\right) + \frac{N}{2}\log(2\pi).
]

• (eq2 / 2 / (std2)).sum()` is the quadratic term (\sum \frac{eq^2}{2\sigma^2}).
• torch.log(std) * x.size()[0]` is the log-σ term (N\log\sigma).
• The constant (\tfrac{N}{2}\log(2\pi)) is omitted (it doesn’t affect optimization).

So the total loss you’re optimizing is (up to an additive constant):
[
\text{loss} = \sum_{i=1}^{N}\frac{eq_i^2}{2\sigma^2} + N\log\sigma,
]
which is precisely the Gaussian NLL for zero-mean residuals with shared std.

Why you need both terms

If you only used loss1, the optimizer could drive std → ∞ to make the loss arbitrarily small (since the denominator grows). The + N·log(std) term prevents this by penalizing large std. Conversely, making std → 0 makes the quadratic term blow up, so there’s a finite optimum.

What value of std is learned?

Set the derivative w.r.t. (\sigma) to zero:

[
\frac{\partial}{\partial\sigma}\left(\sum \frac{eq_i^2}{2\sigma^2} + N\log\sigma\right)
= \sum \left(-\frac{eq_i^2}{\sigma^3}\right) + \frac{N}{\sigma} = 0.
]

Multiply by (\sigma^3):

[

• \sum eq_i^2 + N\sigma^2 = 0 \quad \Rightarrow \quad \sigma^2 = \frac{1}{N}\sum_{i=1}^{N} eq_i^2.
  ]

So the MLE estimate is
[
\boxed{\sigma^2 = \text{mean of squared residuals},}
]
i.e., std converges to the root mean square (RMS) of the residuals (the batch-wise empirical standard deviation, up to the mean-zero assumption). The initialization at 0.1 just seeds the optimization; the log(std) term ensures the parameter doesn’t collapse to 0 or blow up.

Intuition

• loss1 says: “If residuals are large relative to the current σ, penalize strongly.”
• loss2 says: “But don’t let σ become huge; larger σ adds a cost per sample.”
• The balance point makes σ reflect the noise level in your residuals. If your model improves (residuals shrink), the learned σ will also shrink; if residuals are inherently noisy, σ rises to match that scale.

Practical notes

• Positivity constraint: std must stay positive. A common pattern is to parameterize via log_std and set std = softplus(log_std) or std = exp(log_std). If you’re directly optimizing std, ensure it can’t go non-positive (e.g., clamp or reparameterize).
• Global vs. per-sample/per-feature: Your current term N·log(std) suggests a single shared σ for the whole batch. You can generalize to per-output or per-sample σ for heteroscedastic modeling (then the log term sums over those dimensions).
• Constant term: You can add (\tfrac{N}{2}\log(2\pi)) for the exact NLL; it’s constant w.r.t. parameters, so it won’t change gradients for σ or the mean.

If you’d like, I can help you refactor this into a numerically stable PyTorch implementation with a positive σ parameterization and optional heteroscedasticity. Is your std intended to be global, per-output, or per-sample?