<jupyter_start><jupyter_text>Gradient DescentIn this notebook, we will explore what gradient-based optimization is.Are you wondering why we're starting here?Here's why: almost every neural network model that is seen out in the wildis trained by some form of gradient descent.If you have a strong intuition for gradient descent,you will have a tremendous chunk of neural network land demystified.Are you ready? Let's get going! 😃  A brief introduction to _optimization_Firstly, we need to understand what we mean by the term _optimization_.Optimization is something we might do in our daily lives,and that's where we can find great analogies.For example, finding driving directionsthat minimize the amount of time you take to get from A to B,and we would tweak the route we take to achieve that goal.Another example is finding stock picks to maximize financial returns,and we would tweak our stock picks to achieve that goal.These two analogies illuminate the core of optimization:What we mean by optimization really is_finding the minimum or maximum value_ of some functionby tweaking some inputs to that function.What other analogies can you think of?Try writing them down for yourself and evaluate whether they follow the same intuitionof tweaking inputs/levers to minimize/maximize an outcome.  An anchoring example for the math of optimizationIn the deep learning world,optimization almost always takes place in the context of some math function.I am always a fan of anchoring examples,so let's take a math function drawn from a familythat many of us would be familiar with -the 2nd-degree polynomial.Let's use a specific one, $f(x) = x^2 + 3x - 5$.Plotted in code, it looks something like the following:<jupyter_code>import matplotlib.pyplot as plt
import jax.numpy as np


def f(x):
    return x ** 2 + 3 * x - 5


fig, ax = plt.subplots(figsize=(4, 4))
x = np.linspace(-6, 3, 1000)
plt.plot(x, f(x));<jupyter_output><empty_output><jupyter_text>Analytically minimizing $f(x)$$f(x)$ is a smiley-faced function,thus we know that $f(x)$ has a minimum value.For those of you who know how to take derivatives,you can probably work out that the minimum value occurs at $x=-1.5$.As a refresher for those who need it,we start by figuring out what $f'(x)$,the derivative of $f(x)$ w.r.t. $x$, looks like.From introductory calculus, to find the minima or maxima of a function,we know that the value of the derivative at that point is equal to zero,so to find the value of $x$ at which $f'(x)$ is zero, we solve the equation:$$f'(x) = 2x + 3 = 0$$$$x = -1.5$$Just to be doubly-sure that we are finding a minima,we might also want to check the second derivative to make sure it is positive.That's an exercise you may attempt on your own if you desire.The derivative function $f'(x)$ is also known as the _gradient function_.The derivative function tells us how $y$ will change as we change $x$;it always points us in the way that we need to change $x$to increase the value of $f(x)$.You can verify this by plugging indifferent values of $x$ into the derivative function.At values to the left of the minima, you will get negative numberswhile at values to the right of the minima, you will get positive numbers.  Minimizing the polynomial _computationally_While for simple polynomials we can take gradients easily,for complicated math functions,finding an analytical solution to identify the minima or maximamight not be easy!But fret not, that's where computers come into the picture.There is a general purpose way of finding minima or maxima analytically,and that is by gradient descent.The core of gradient descent is as follows:We pick a starting point at will -there are heuristics to pick good starting points,but for now, let's be satisfied with "at random".We take little steps in the _negative_ direction of the gradient.At each step, we evaluate the gradientand then take a small step in $x$ that lets us followthe negative direction of the gradient.We repeat this process over and overfor a fixed number of steps(which needs to be configured beforehand).This routine is what we call _gradient descent_.Those of you who are attuned to the algorithmmight see the following key knobs for gradient descent:- Where the initial point is.- How large of a step to take.- How many steps to take.We won't immediately address how to set these,but you should know that those knobs exist.At this point in the course,I'd ask that you accept the default values that we'll provide for youknowing that pragmatically speaking,they'll help you solve the exercises you encounterin a reasonable amount of time.  The workhorse powered by _automatic differentiation_Gradient descent is the workhorse of all of modern deep learning.Whenever we fit a deep learning model,we are also minimizing a function -- the _loss function_ -- that tells us how bad our model is.We'll visit the idea of a loss function in the second notebook,so keep that name in your mind for the time being.For now, though, in order to help you get familiar with gradient descentand how to write a program that uses gradient descent in it,we're going to introduce to you an _automatic differentiation system_:a program that can return the derivative functionof any function that performs numerical computation.We'll be using JAX,which provides automatic differentiation on top of the NumPy API --an API that many of us in the scientific Python world should be familiar with.  JAX's `grad` functionRemember, when we are training a model to do some prediction,we need a way of calculating gradient functions so that we can do gradient-based optimization.Calculating the gradient function by hand is sometimes easy,especially if you're competent with algebra and calculus.However, in many cases, calculating the gradient function by hand is tedious and repetitive.As it turns out, though, tedious and repetitive tasks are _exactly_ what computers are great at!This is where JAX comes in.JAX is a Python package developed by friends at Google.You can think of JAX as the NumPy API turned on steroids. JAX ships with an automatic differentiation system,exposed to end-users in a simple function: `grad`.Here's what `grad` does.It takes in a function that returns a scalar output,and returns another _transformed_ function.That newly returned function is the gradient function of the original.The gradient function takes in the same arguments as the original,but will evaluate the gradient of the original function w.r.t. the arguments passed in.  Example use of `grad`Let's see a concrete example of this in action.We have here an $x^2$ polynomial function:$$ f(x) = 3x^2 + 4x - 2 $$Implemented in NumPy, it'll look like this:<jupyter_code>def func(x):
    return 3 * x ** 2 + 4 * x - 2<jupyter_output><empty_output><jupyter_text>Analytically, `func`'s gradient function is going to be:$$ f'(x) = 6x + 4 $$If we wanted, we could write that gradient function in NumPy by hand.However, we're instead going to obtain the gradient function using JAX's `grad` function.<jupyter_code>from jax import grad

func_p = grad(func)<jupyter_output><empty_output><jupyter_text>Pictorially, we'd represent the transformation as follows: In the diagram above, the functions that we want to take derivatives of are the nodes,and `grad` is an edge that transforms a function.In other words, `grad` acts as a program transformation devicethat takes `f` and turns it into `f'`, a.k.a. the first-order derivative of `f`.(You can keep passing the result through `grad` if you desire!) Now, we can evaluate the gradient function w.r.t. any value of $x$.Let's sanity-check that we have it right.If we plug in x = 3, then $f'(x) = 6(3) + 4 = 22$.Calling the same on `func_p`:<jupyter_code>func_p(3.0)<jupyter_output><empty_output><jupyter_text>In the Jupyter notebook, I would encourage you to hand-calculate the gradient evaluated at various values of $x$ to convince yourself that `grad` does the right thing.  Gradient DescentSeeing how `grad` works isn't the main point though! We want to use the gradient function `func_p` in an optimization routine to _minimize_ the function `func`. Here's how we do it.  Step 1: Initialize `x`Firstly, we start by initializing $x$ at some value.<jupyter_code>x = 4.2<jupyter_output><empty_output><jupyter_text>I chose a fixed value to start, but you could have just as well drawn a number from a random number generator instead.  Step 2: Take small steps in negative direction of gradientNext, we take steps in the negative direction of the gradient a fixed number of times.<jupyter_code># hyperparameters
num_steps = 300
step_size = 5e-3

# loop
for i in range(num_steps):
    x -= func_p(x) * step_size<jupyter_output><empty_output><jupyter_text>There are two _hyperparameters_ that we can configure here - the number of steps and the step size. Step sizes are canonically smaller than 1.0, and most of the time set by default to 0.005. (There's an <NAME> tweet about this.) The smaller the step size, the more steps may be needed to reach convergence. The number of steps is the other hyperparameter; the longer we run, the more likely we are to converge on the minima.  Step 3: Verify proximity to stationary pointLet's make sure that we have hit someplace near a _stationary point_ in our optimization routine. Stationary points refer to places where the derivative of the function, or the evaluated gradient, is equal to zero.<jupyter_code>func_p(x)<jupyter_output><empty_output><jupyter_text>Indeed, we have!And that's a great sign that we're doing the right thing _computationally_.  Step 4: Check second derivativeLet's also evaluate the second derivative to sanity-check that we have hit a minima.<jupyter_code>func_pp = grad(func_p)
func_pp(x)<jupyter_output><empty_output><jupyter_text>And just like that, we see that the second derivative at our final value of `x` is a positive number,so we know we're in a regime which is a minima!  ExercisesComing up is a pair of exercises for you to practice what you should have learned here today:1. How to obtain the derivative function of a NumPy/Python function, and2. How to use the evaluated derivative to find the minima of that function.  Exercise: Obtaining the derivative function of a Python functionFor this exercise, we would like you to use `grad` to obtain the first and second derivative functions of a power-4 polynomialand evaluate the derivatives.First off, here is the function in code and in plot:<jupyter_code>def cubic_polynomial(x):
    return 2 * x ** 4 - 3 * x ** 3 - 9 * x ** 2 + 4 * x + 3
x = np.linspace(-3, 4, 1000)
y = cubic_polynomial(x)

plt.plot(x, y)<jupyter_output><empty_output><jupyter_text>Now, we would like you to evaluate the gradient at the following values of `x`: `[3.0, 0.5, -2.0]`. (Be sure to pass them in as floats, i.e. `3.0` rather than `3`.)<jupyter_code>dcubic_polynomial = grad(cubic_polynomial)

values = [3.0, 0.5, -2.0]

for value in values:
    print(dcubic_polynomial(value))<jupyter_output><empty_output><jupyter_text>Now, we would like you to evaluate the second derivative at the same values.<jupyter_code>ddcubic_polynomial = grad(dcubic_polynomial)

for value in values:
    print(ddcubic_polynomial(value))<jupyter_output><empty_output><jupyter_text>Do the values of the second derivative make sense?For the values `3.0` and `-2.0`, they should,because they are within the smiley-faced regime of the function.The value for 0.5 should make sense too,because it's in the frowny-faced regime of the function.  Exercise: Finding a minima of a Python function using derivativesWe are now going to attempt to find one of the minima of that functionusing what you know about writing gradient descent loops.Using gradient descent, computationally find one of the minima of the power-4 polynomial.<jupyter_code>x = -5.0

for i in range(1000):
    x -= dcubic_polynomial(x) * 0.001

print(x)<jupyter_output><empty_output><jupyter_text>Multidimensional gradientsThus far, you've seen how to use the `grad` functionin the context of functions where we want to obtain the derivative with respect to one parameter.Thinking forward to our machine learning models,we'll need to obtain the derivative of functions with respect to more than one parameter.This is important because even the simplest linear models have two model parameters,while the vast majority of neural networks have hundreds to millions of parameters.So, then, how do we take the derivative of a function that accepts more than one parameter?  Derivatives of a 2D functionTo illustrate, we're going to play a game where we try to find the minima of a 2D function.I will introduce to you the setting here, which will form the basis of your exercise below.You're programming a robot to find the low points in a landscape that has peaks and valleys.This is important because the low points are the likely place where _gold_ could be found.I'm going to provide the math equation that describes your gold field:$$ f(x, y) = sin(x) + sin(y) $$It's possible to write this equation in Python and NumPy in two forms:one where it accepts `(x, y)` coordinates as a tuple,and one where it accepts `x` and `y` independently:<jupyter_code>def goldfield1(coords):
    x, y = coords
    return np.arctan(x)**2 + np.power(np.sin(y), 2)


def goldfield2(x, y):
    return np.arctan(x)**2 + np.power(np.sin(y), 2)<jupyter_output><empty_output><jupyter_text>This is what the math function looks like:<jupyter_code>import matplotlib.pyplot as plt
from matplotlib import cm
from matplotlib.ticker import LinearLocator

fig, ax = plt.subplots(subplot_kw={"projection": "3d"})

# Make data.
X = np.arange(-5, 5, 0.25)
Y = np.arange(-5, 5, 0.25)
X, Y = np.meshgrid(X, Y)
Z = goldfield2(X, Y)

# Plot the surface.
surf = ax.plot_surface(X, Y, Z, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False)

# Customize the z axis.
# ax.set_zlim(-1.01, 1.01)
ax.zaxis.set_major_locator(LinearLocator(10))
# A StrMethodFormatter is used automatically
ax.zaxis.set_major_formatter('{x:.02f}')

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()<jupyter_output><empty_output><jupyter_text>It should be obvious to you that both `goldfield1` and `goldfield2` are mathematically equivalent,but their function signatures are different, which will make the difference here.By default, `grad` takes the derivative of a function w.r.t. the first argument.Thus, it's easy to take the derivative of `goldfield1`:<jupyter_code>coords = np.array([1.0, 1.0])
coords
dgoldfield1 = grad(goldfield1)
dgoldfield1(coords)<jupyter_output><empty_output><jupyter_text>However, what do we do for the second function, `goldfield2`? As it turns out, `grad` has an `argnums` keyword argument that tells us with respect to which argument to take the derivative.We can take the derivative w.r.t. both arguments:<jupyter_code>dgoldfield2 = grad(goldfield2, argnums=(0, 1))
x, y = coords
dgoldfield2(x, y)<jupyter_output><empty_output><jupyter_text>And we can also take partial derivatives, starting with the derivative w.r.t. `x`:<jupyter_code>dgoldfield2_dx = grad(goldfield2, argnums=0)
dgoldfield2_dx(x, y)<jupyter_output><empty_output><jupyter_text>and then the derivative w.r.t. `y`:<jupyter_code>dgoldfield2_dy = grad(goldfield2, argnums=1)
dgoldfield2_dy(x, y)<jupyter_output><empty_output><jupyter_text>Whether you choose to go with partial or full derivativeswill depend on how much complexity you wish to exposeto whoever is using the function.In my opinion, `dgoldfield1` is less complex, but is also more restrictive.By contrast, structuring your function in the style of `dgoldfield2`allows you to take a _partial_ derivative with respect to either `x` or `y`,If your particular project requires partial derivatives,you can easily specify, using `argnums` which partial derivative you require. Alrighty, to reinforce these concepts, we've got a bunch of exercises present for you. As always, these examples are present in the notebook for you to review as you attempt the exercises!  ExercisesWe are going to leverage the goldfield function for the following exercises.  Exercise 1: Find a minima w.r.t. both `x` and `y`As you can see in the plot, there are multiple local minima in the function. We'd like you to find one of them by using gradient descent.<jupyter_code>coords = np.array([0., -2])
for i in range(1000):
    g = dgoldfield1(coords)
    coords -= g * 0.01
coords<jupyter_output><empty_output><jupyter_text>Exercise 2: Find a minima along one axis, either `x` or `y`Using partial derivatives, find a minima of the goldfieldwhen fixing either `x` or `y` to a particular value.(You get to choose which you wish to fix!)<jupyter_code># This answer assumes that we are finding a minima along the y-axis,
# i.e. fixing the value of `x`.

x, y = 1.0, 1.0

for i in range(1000):
    g = dgoldfield2_dy(x, y)
    y -= g * 0.01
x, y<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.06797493793592624, 'fraction_numerical': 0.011289750561532096, 'mean_word_length': 4.963694043003172, 'pattern_counts': {'":': 1, '<': 82, '<?xml version=': 0, '>': 82, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '20', 'dir': 'jupyter-structured-clean-dedup', 'id': '90849', 'license': 'MIT', 'n_tokens_mistral': 4966, 'n_tokens_neox': 4680, 'n_words': 2379, 'path': 'notebooks/01-gradient-descent.ipynb', 'repo_name': 'NFAcademy/2021_course_dev-ericmjl', 'size': '16927'}