Java: linear regression with gradient descent

Fit a straight line to data by repeatedly stepping the weights against the gradient of the squared error.

Java

import java.util.Locale;

class Main {
  public static void main(String[] args) {
    double[] x = {1, 2, 3, 4};
    double[] y = {3, 5, 7, 9};
    double w = 0, b = 0, rate = 0.05;
    for (int epoch = 0; epoch < 2000; epoch++) {
      double gw = 0, gb = 0;
      for (int i = 0; i < x.length; i++) {
        double error = w * x[i] + b - y[i];
        gw += 2 * error * x[i] / x.length;
        gb += 2 * error / x.length;
      }
      w -= rate * gw;
      b -= rate * gb;
    }
    System.out.printf(Locale.ROOT, "w=%.2f b=%.2f%n", w, b);
    System.out.printf(Locale.ROOT, "prediction for x=5: %.2f%n", w * 5 + b);
  }
}

Output

w=2.00 b=1.00
prediction for x=5: 11.00

Explanation

Linear regression fits y = w·x + b to data. The mean squared error measures how far the predictions are from the targets. Gradient descent reduces it by computing the gradient, the direction in which the error grows fastest, and moving the parameters a small step the other way.

For each point the error is (w·x + b − y). The gradient of the mean squared error is 2·error·x averaged over the points for w, and 2·error averaged for b. The learning rate controls the step size. The data follow y = 2x + 1, so the fitted values should approach w = 2 and b = 1.

Explore the dry run

Guided walkthroughReady

1import java.util.Locale;

3class Main {

4 public static void main(String[] args) {

5 double[] x = {1, 2, 3, 4};

6 double[] y = {3, 5, 7, 9};

7 double w = 0, b = 0, rate = 0.05;

8 for (int epoch = 0; epoch < 2000; epoch++) {

9 double gw = 0, gb = 0;

10 for (int i = 0; i < x.length; i++) {

11 double error = w * x[i] + b - y[i];

12 gw += 2 * error * x[i] / x.length;

13 gb += 2 * error / x.length;

14 }

15 w -= rate * gw;

16 b -= rate * gb;

17 }

18 System.out.printf(Locale.ROOT, "w=%.2f b=%.2f%n", w, b);

19 System.out.printf(Locale.ROOT, "prediction for x=5: %.2f%n", w * 5 + b);

20 }

21}

Follow the example

Start to inspect the source, values, and output. These prepared steps explain the published sample; grouped loop steps describe multiple iterations.

View the complete dry-run table

Prepared trace of the published program. Some rows group repeated iterations.
What happens	State after step	Output so far
Start from a flat line and a modest learning rate.	`w = 0, b = 0`	—
Prediction minus target for one point.	`w = 0, b = 0`	—
Step the parameters against the gradient.	`w and b move toward 2 and 1`	—
After 2000 passes the line matches the data.	`w ≈ 2.00, b ≈ 1.00`	w=2.00 b=1.00
The fitted line predicts 2·5 + 1 for an unseen x.	`w ≈ 2.00, b ≈ 1.00`	w=2.00 b=1.00 prediction for x=5: 11.00

Cost and common mistakes

O(epochs · n) time and O(1) extra space for n points.

A learning rate that is too large makes the error grow instead of shrink; too small makes learning very slow.
Use Locale.ROOT with printf so the decimal separator is a point on every machine.
Unscaled features with very different ranges slow gradient descent down.

Try it yourself

What would the weights become if the data were y = 3x?

Reveal answer and reasoning

w would approach 3 and b would approach 0.

Java: linear regression with gradient descent

Explanation

Follow the example

Cost and common mistakes

Try it yourself

Further reading