tensorflow optimization with added noise near first-order critical points

noisy_opt_demo.py

import numpy as np
import tensorflow as tf

# initialize variable x=10
x = tf.Variable(10.0, dtype=tf.float32)

# objective is x ** 3 which has a saddle point at x=0
f = x ** 3

# create optimizer and compute gradients
opt = tf.train.GradientDescentOptimizer(1e-3)
grads_and_vars = opt.compute_gradients(f)
dfdx = grads_and_vars[0][0]

# create training op (does a step of gradient descent)
train_op = opt.apply_gradients(grads_and_vars)

# store stuff in a list
datastore = []

# threshold for determining whether or not we should add noise to the gradient
thr = 0.3

# start a session
sess = tf.Session()
sess.run(tf.global_variables_initializer())

# loop through optimization
for k in range(3000):
    
    # run gradient descent step
    x_np, dfdx_np, f_np, _ = sess.run([x, dfdx, f, train_op])
    
    # store stuff
    ds.append((x_np, dfdx_np, f_np))
    
    # add noise if the gradient is small
    # this means that if we get close to a saddle point, the noise will kick us out of it
    if np.linalg.norm(dfdx_np) < thr:
        sess.run(x.assign_add(np.random.randn()))

running the above code should generate something like the following:

we optimize for a bit, getting drawn to the saddle point at x=0. but when the gradient becomes small, we start adding some random noise, which sometimes kicks up back up the hill (the upward jumps) but then around iteration 2000 we happened to jump past the point x=0 and then start to accelerate to negative infinity.