numeric stability for affine_transform and operator* for matrix vector multiplication
- Dominant language
- C++
- Stars
- 3.4k
- Forks
- 701
- PR merge metrics
- No merged PRs in 30d
Description
Do affine_transform({b, W_1, x_1, W_2, x_2, ...}) and b + W_1 * x_1 + W_2 * x_2 have any implementation difference or numeric stability difference?
I recently spent a lot of time debugging an "nan" problem. The code below could run several times before the forward pass would yield 'nan'. I have check every detail of my code, and finaly change the matrix-vector multiplication into an affine_transform({0, W, x}). And only after that, the code runs well without yielding 'nan'.
before
```c++
// lstm_outputs vector, shape of each Expression is {2 * lstm_dim}
// s_bias shape {attn_dim}
// s_q2a shape {attn_dim, decoder_dim}
// prev_h shape {decoder_dim}
// s_k2a shape {attn_dim, 2 * lstm_dim}
// s_v shape {attn_dim}
Expression s_cxt;
{
// stack conditioned input
vector vs;
vs.push_back(lstm_outputs[0]);
for (size_t i = 1; i < stacki.size(); ++i) {
vs.push_back(lstm_outputs.at(stacki[i] + 1));
}
vector ts;
for (auto& x: vs) {
ts.push_back(tanh(affine_transform({s_bias, s_q2a, prev_h, s_k2a, x})));
}
Expression ts_cols = transpose(concatenate_cols(ts));
Expression vs_cols = concatenate_cols(vs);
cerr << ts_cols.dim() << endl << s_v.dim() << endl;
Expression energy = ts_cols * s_v; // would cause 'nan'
//Expression energy = affine_transform({zeros(hg, {ts.size()}), ts_cols, s_v});
Expression attn = softmax(energy);
s_cxt = vs_cols * attn; //would cause 'nan'
//s_cxt = affine_transform({zeros(hg, lstm_outputs[0].dim()), vs_cols, attn});
}
```
after
```c++
// lstm_outputs vector, shape of each Expression is {2 * lstm_dim}
// s_bias shape {attn_dim}
// s_q2a shape {attn_dim, decoder_dim}
// prev_h shape {decoder_dim}
// s_k2a shape {attn_dim, 2 * lstm_dim}
// s_v shape {attn_dim}
Expression s_cxt;
{
// stack conditioned input
vector vs;
vs.push_back(lstm_outputs[0]);
for (size_t i = 1; i < stacki.size(); ++i) {
vs.push_back(lstm_outputs.at(stacki[i] + 1));
}
vector ts;
for (auto& x: vs) {
ts.push_back(tanh(affine_transform({s_bias, s_q2a, prev_h, s_k2a, x})));
}
Expression ts_cols = transpose(concatenate_cols(ts));
Expression vs_cols = concatenate_cols(vs);
cerr << ts_cols.dim() << endl << s_v.dim() << endl;
//Expression energy = ts_cols * s_v; would cause 'nan'
Expression energy = affine_transform({zeros(hg, {ts.size()}), ts_cols, s_v});
Expression attn = softmax(energy);
// s_cxt = vs_cols * attn; would cause 'nan'
s_cxt = affine_transform({zeros(hg, lstm_outputs[0].dim()), vs_cols, attn});
}
```
I am also looking into the code for them. Thanks!
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by comparing the implementations and evaluation paths for affine_transform and matrix-vector operator* using the provided C++ example. Reproduce the NaN behavior, inspect intermediate values and operation ordering, and document whether the two forms differ numerically and what test coverage demonstrates the result.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100