Java · Course navigation

Java: linear regression with gradient descent

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

On this page

Before you start

Java arrays and floating-point numbers.

Fit a straight line to data by repeatedly stepping the weights against the gradient of the squared error.
Predict the sample output and explain the state changes.
Complete the practice task and check the boundary cases.

The concept

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.

Code example

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

Example details

Follow the code

Open the walkthrough to inspect the example one step at a time.

Show guided walkthrough

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

Behind the scenes

How Java processes this code

The Java compiler produces bytecode. The JVM loads the class and invokes its main method.

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.