Tests the correctness of one step of GGT.
(param,
g_t,
lr,
grad_buffer,
m,
window,
t,
beta1=0.9,
eps=1e-4,
svd_eps=1e-6,
sigma_eps=1e-2)
| 31 | |
| 32 | |
| 33 | def ggt_update_numpy(param, |
| 34 | g_t, |
| 35 | lr, |
| 36 | grad_buffer, |
| 37 | m, |
| 38 | window, |
| 39 | t, |
| 40 | beta1=0.9, |
| 41 | eps=1e-4, |
| 42 | svd_eps=1e-6, |
| 43 | sigma_eps=1e-2): |
| 44 | """Tests the correctness of one step of GGT.""" |
| 45 | m_t = m * beta1 + (1 - beta1) * g_t |
| 46 | grad_buffer[((t - 1) % window), :] = m_t |
| 47 | m_matrix = np.transpose(grad_buffer / np.sqrt(np.minimum(t, window))) |
| 48 | mm = np.dot(np.transpose(m_matrix), m_matrix) |
| 49 | damping = np.eye(window) * svd_eps |
| 50 | u, sigma, _ = np.linalg.svd(mm + damping) |
| 51 | |
| 52 | sigma_sqrt_inv = np.power(np.sqrt(sigma) + sigma_eps, -3) |
| 53 | new_step = np.linalg.multi_dot([ |
| 54 | m_matrix, u, |
| 55 | np.diag(sigma_sqrt_inv), |
| 56 | np.transpose(u), |
| 57 | np.transpose(m_matrix), m_t |
| 58 | ]) |
| 59 | |
| 60 | sigma_sqrt_min = np.sqrt(sigma).min() |
| 61 | |
| 62 | if sigma_sqrt_min > eps: |
| 63 | new_step += (m_t - np.linalg.multi_dot([ |
| 64 | m_matrix, u, |
| 65 | np.diag(1.0 / sigma), |
| 66 | np.transpose(u), |
| 67 | np.transpose(m_matrix), m_t |
| 68 | ])) * (1.0 / sigma_sqrt_min) |
| 69 | |
| 70 | param_t = param - lr * new_step |
| 71 | return param_t, m_t, grad_buffer |
| 72 | |
| 73 | |
| 74 | class GGTOptimizerTest(test.TestCase): |