What Is Argmax in Machine Learning?

Argmax is a mathematical function that you may encounter in applied machine learning.

For example, you may see “argmax” or “arg max” used in a research paper used to describe an algorithm. You may also be instructed to use the argmax function in your algorithm implementation.

This may be the first time that you encounter the argmax function and you may wonder what it is and how it works.

In this tutorial, you will discover the argmax function and how it is used in machine learning.

After completing this tutorial, you will know:

• Argmax is an operation that finds the argument that gives the maximum value from a target function.
• Argmax is most commonly used in machine learning for finding the class with the largest predicted probability.
• Argmax can be implemented manually, although the argmax() NumPy function is preferred in practice.

Kick-start your project with my new book Linear Algebra for Machine Learning, including step-by-step tutorials and the Python source code files for all examples.

Let’s get started.

What Is argmax in Machine Learning?
Photo by Bernard Spragg. NZ, some rights reserved.

Tutorial Overview

This tutorial is divided into three parts; they are:

1. What Is Argmax?
2. How Is Argmax Used in Machine Learning?
3. How to Implement Argmax in Python

What Is Argmax?

Argmax is a mathematical function.

It is typically applied to another function that takes an argument. For example, given a function g() that takes the argument x, the argmax operation of that function would be described as follows:

• result = argmax(g(x))

The argmax function returns the argument or arguments (arg) for the target function that returns the maximum (max) value from the target function.

Consider the example where g(x) is calculated as the square of the x value and the domain or extent of input values (x) is limited to integers from 1 to 5:

• g(1) = 1^2 = 1
• g(2) = 2^2 = 4
• g(3) = 3^2 = 9
• g(4) = 4^2 = 16
• g(5) = 5^2 = 25

We can intuitively see that the argmax for the function g(x) is 5.

That is, the argument (x) to the target function g() that results in the largest value from the target function (25) is 5. Argmax provides a shorthand for specifying this argument in an abstract way without knowing what the value might be in a specific case.

• argmax(g(x)) = 5

Note that this is not the max() of the values returned from function. This would be 25.

It is also not the max() of the arguments, although in this case the argmax and max of the arguments is the same, e.g. 5. The argmax() is 5 because g returns the largest value (25) when 5 is provided, not because 5 is the largest argument.

Typically, “argmax” is written as two separate words, e.g. “arg max“. For example:

• result = arg max(g(x))

It is also common to use the arg max function as an operation without brackets surrounding the target function. This is often how you will see the operation written and used in a research paper or textbook. For example:

• result = arg max g(x)

You can also use a similar operation to find the arguments to the target function that result in the minimum value from the target function, called argmin or “arg min.”

How Is Argmax Used in Machine Learning?

The argmax function is used throughout the field of mathematics and machine learning.

Nevertheless, there are specific situations where you will see argmax used in applied machine learning and may need to implement it yourself.

The most common situation for using argmax that you will encounter in applied machine learning is in finding the index of an array that results in the largest value.

Recall that an array is a list or vector of numbers.

It is common for multi-class classification models to predict a vector of probabilities (or probability-like values), with one probability for each class label. The probabilities represent the likelihood that a sample belongs to each of the class labels.

The predicted probabilities are ordered such that the predicted probability at index 0 belongs to the first class, the predicted probability at index 1 belongs to the second class, and so on.

Often, a single class label prediction is required from a set of predicted probabilities for a multi-class classification problem.

This conversion from a vector of predicted probabilities to a class label is most often described using the argmax operation and most often implemented using the argmax function.

Let’s make this concrete with an example.

Consider a multi-class classification problem with three classes: “red“, “blue,” and “green.” The class labels are mapped to integer values for modeling, as follows:

• red = 0
• blue = 1
• green = 2

Each class label integer values maps to an index of a 3-element vector that may be predicted by a model specifying the likelihood that an example belongs to each class.

Consider a model has made one prediction for an input sample and predicted the following vector of probabilities:

• yhat = [0.4, 0.5, 0.1]

We can see that the example has a 40 percent probability of belonging to red, a 50 percent probability of belonging to blue, and a 10 percent probability of belonging to green.

We can apply the argmax function to the vector of probabilities. The vector is the function, the output of the function is the probabilities, and the input to the function is a vector element index or an array index.

• arg max yhat

We can intuitively see that in this case, the argmax of the vector of predicted probabilities (yhat) is 1, as the probability at array index 1 is the largest value.

Note that this is not the max() of the probabilities, which would be 0.5. Also note that this is not the max of the arguments, which would be 2. Instead it is the argument that results in the maximum value, e.g. 1 that results in 0.5.

• arg max yhat = 1

We can then map this integer value back to a class label, which would be “blue.”

• arg max yhat = “blue”

How to Implement Argmax in Python

The argmax function can be implemented in Python for a given vector of numbers.

Argmax from Scratch

First, we can define a function called argmax() that enumerates a provided vector and returns the index with the largest value.

The complete example is listed below.

Running the example prints the argmax of our test data used in the previous section, which in this case is an index of 1.

Argmax with NumPy

Thankfully, there is a built-in version of the argmax() function provided with the NumPy library.

This is the version that you should use in practice.

The example below demonstrates the argmax() NumPy function on the same vector of probabilities.

Running the example prints an index of 1, as is expected.

It is more likely that you will have a collection of predicted probabilities for multiple samples.

This would be stored as a matrix with rows of predicted probabilities and each column representing a class label. The desired result of an argmax on this matrix would be a vector with one index (or class label integer) for each row of predictions.

This can be achieved with the argmax() NumPy function by setting the “axis” argument. By default, the argmax would be calculated for the entire matrix, returning a single number. Instead, we can set the axis value to 1 and calculate the argmax across the columns for each row of data.

The example below demonstrates this with a matrix of four rows of predicted probabilities for the three class labels.

Running the example first prints the shape of the matrix of predicted probabilities, confirming we have four rows with three columns per row.

The argmax of the matrix is then calculated and printed as a vector, showing four values. This is what we expect, where each row results in a single argmax value or index with the largest probability.

Further Reading

This section provides more resources on the topic if you are looking to go deeper.

• numpy.argmax API.
• Arg max, Wikipedia.

Summary

In this tutorial, you discovered the argmax function and how it is used in machine learning.

Specifically, you learned:

• Argmax is an operation that finds the argument that gives the maximum value from a target function.
• Argmax is most commonly used in machine learning for finding the class with the largest predicted probability.
• Argmax can be implemented manually, although the argmax() NumPy function is preferred in practice.

Do you have any questions?
Ask your questions in the comments below and I will do my best to answer.

Get a Handle on Linear Algebra for Machine Learning!

Develop a working understand of linear algebra

...by writing lines of code in python

Discover how in my new Ebook:
Linear Algebra for Machine Learning

It provides self-study tutorials on topics like:
Vector Norms, Matrix Multiplication, Tensors, Eigendecomposition, SVD, PCA and much more...

Finally Understand the Mathematics of Data

Skip the Academics. Just Results.

See What's Inside