# Linear Regression — MLX 0.0.7 documentation

Skip to main content
`Ctrl`+`K`

# Linear Regression

Let’s implement a basic linear regression model as a starting point to learn MLX. First import the core package and setup some problem metadata:

    import mlx.core as mx

    num_features = 100
    num_examples = 1_000
    num_iters = 10_000  # iterations of SGD
    lr = 0.01  # learning rate for SGD

We’ll generate a synthetic dataset by:

1. Sampling the design matrix `X`.

2. Sampling a ground truth parameter vector `w_star`.

3. Compute the dependent values `y` by adding Gaussian noise to `X @ w_star`.

    # True parameters
    w_star = mx.random.normal((num_features,))

    # Input examples (design matrix)
    X = mx.random.normal((num_examples, num_features))

    # Noisy labels
    eps = 1e-2 * mx.random.normal((num_examples,))
    y = X @ w_star + eps

We will use SGD to find the optimal weights. To start, define the squared loss and get the gradient function of the loss with respect to the parameters.

    def loss_fn(w):
        return 0.5 * mx.mean(mx.square(X @ w - y))

    grad_fn = mx.grad(loss_fn)

Start the optimization by initializing the parameters `w` randomly. Then repeatedly update the parameters for `num_iters` iterations.

    w = 1e-2 * mx.random.normal((num_features,))

    for _ in range(num_iters):
        grad = grad_fn(w)
        w = w - lr * grad
        mx.eval(w)

Finally, compute the loss of the learned parameters and verify that they are close to the ground truth parameters.

    loss = loss_fn(w)
    error_norm = mx.sum(mx.square(w - w_star)).item() ** 0.5

    print(
        f"Loss {loss.item():.5f}, |w-w*| = {error_norm:.5f}, "
    )
    # Should print something close to: Loss 0.00005, |w-w*| = 0.00364

Complete linear regression and logistic regression examples are available in the MLX GitHub repo.