Logistic Regression (Gradient Descent)

Implement binary logistic regression trained with batch gradient descent, from scratch with NumPy. Logistic regression is the canonical linear classifier, and interviewers use it to check that you can assemble the full loop: add a bias, push features through a stable sigmoid, form the gradient of the cross-entropy loss, and update the weights. The math is short but the two classic traps (forgetting the bias, and an overflowing sigmoid) are exactly what they are watching for.

Model

Stack a column of ones onto X so the intercept is folded into the weight vector. With that augmented design matrix, the model predicts

p = sigmoid(X w)          where sigmoid(z) = 1 / (1 + exp(-z))

Training minimizes the mean binary cross-entropy loss. Its gradient with respect to w collapses to a single clean expression:

grad = (1 / n) * X^T (sigmoid(X w) - y)

and each gradient-descent step is w <- w - lr * grad. Note the gradient depends only on the residual p - y, so you never have to evaluate the loss itself to train.

Task

Implement two functions with the signatures shown in the starter.

• train_logistic_regression(X, y, lr, n_iters) prepends a bias column of ones to X, initializes the weights to zeros, and runs n_iters full-batch gradient-descent steps using the gradient above. It returns the length n_features + 1 weight vector, where w[0] is the bias.
• predict_proba(X, w) prepends the same bias column and returns sigmoid(X w) as probabilities in (0, 1).

Use a numerically stable sigmoid so large-magnitude scores do not overflow. Do not call sklearn, scipy, or any built-in logistic/expit helper.

Example

import numpy as np

# Two clearly separable clusters in 1D.
X = np.array([[0.0], [1.0], [2.0], [8.0], [9.0], [10.0]])
y = np.array([0, 0, 0, 1, 1, 1])

w = train_logistic_regression(X, y, lr=0.5, n_iters=5000)
p = predict_proba(X, w)

np.round(p, 2)                # -> array([0., 0., 0., 1., 1., 1.])
(p >= 0.5).astype(int)        # -> array([0, 0, 0, 1, 1, 1])  (all training points correct)

The learned weights split the two clusters: rows near 0 get probabilities close to 0, rows near 10 get probabilities close to 1.

Constraints

• 1 <= n_samples <= 10^5, 1 <= n_features <= 10^3; use vectorized NumPy, not Python loops over samples.
• y contains only 0 and 1.
• Scores X w may reach large magnitudes; the sigmoid must not raise an overflow warning or return nan/inf.
• Values fit in float64.

Approach

Fold the bias into the weights by prepending a column of ones to X, so a single matrix-vector product X w produces the raw scores. Push those scores through a sign-branched sigmoid so exp never overflows, then take gradient-descent steps using the closed-form gradient of the mean cross-entropy, (1/n) X^T (p - y). Weights start at zero; after n_iters steps we return them, and predict_proba reuses the same augmentation and sigmoid to score new data.

Solution

import numpy as np

def _sigmoid(z: np.ndarray) -> np.ndarray:
    # Sign-branched: exp() only ever sees a non-positive argument, so no overflow.
    z = np.asarray(z, dtype=np.float64)
    pos = 1.0 / (1.0 + np.exp(-np.abs(z)))
    neg = np.exp(-np.abs(z)) / (1.0 + np.exp(-np.abs(z)))
    return np.where(z >= 0, pos, neg)

def _add_bias(X: np.ndarray) -> np.ndarray:
    X = np.asarray(X, dtype=np.float64)
    ones = np.ones((X.shape[0], 1), dtype=np.float64)
    return np.hstack([ones, X])           # column 0 is the bias feature

def train_logistic_regression(
    X: np.ndarray,
    y: np.ndarray,
    lr: float = 0.1,
    n_iters: int = 1000,
) -> np.ndarray:
    Xb = _add_bias(X)                     # (n, d+1)
    y = np.asarray(y, dtype=np.float64)   # (n,)
    n = Xb.shape[0]
    w = np.zeros(Xb.shape[1], dtype=np.float64)
    for _ in range(n_iters):
        p = _sigmoid(Xb @ w)              # (n,) predicted probabilities
        grad = (Xb.T @ (p - y)) / n       # (d+1,) gradient of mean cross-entropy
        w -= lr * grad
    return w

def predict_proba(X: np.ndarray, w: np.ndarray) -> np.ndarray:
    Xb = _add_bias(X)
    return _sigmoid(Xb @ w)

Walkthrough

On the 1D example with two clusters:

• _add_bias(X) turns each row [x] into [1, x], so w = [w0, w1] where w0 is the bias and w1 the slope. Weights start at [0, 0], so the first prediction is sigmoid(0) = 0.5 for every row.
• Step one: p - y = [0.5, 0.5, 0.5, -0.5, -0.5, -0.5]. The gradient (1/n) Xb^T (p - y) is negative in the x component (large-x rows have negative residuals), so w1 increases and the boundary starts tilting to separate the clusters. The bias adjusts to place the 0.5 crossover between the groups.
• Repeating for n_iters steps drives p toward y: rows with x near 0 head to probability ~0, rows near 10 head to ~1.
• predict_proba(X, w) re-augments and returns sigmoid(Xb @ w), giving [0.0, 0.0, 0.0, 1.0, 1.0, 1.0] (rounded to two places); thresholding at 0.5 classifies all six training points correctly.

Complexity & notes

• Time O(n_iters * n * d), dominated by the two matrix-vector products per step (Xb @ w and Xb.T @ (p - y)), each O(n * d). Space O(n * d) for the augmented design matrix plus O(n) and O(d) temporaries.
• The bias must be a learned parameter. Dropping it forces the decision boundary through the origin and can make well-separated data unlearnable; prepending a ones column is the cleanest way to include it without special-casing the update.
• The stable sigmoid matters here specifically because scores grow during training: once the clusters separate, X w reaches large magnitudes, and a naive 1/(1+exp(-z)) would overflow to inf, produce nan gradients, and destroy the weights. Folding through np.abs keeps every exp argument in (-inf, 0].
• The gradient uses no log terms because they cancel when you differentiate cross-entropy through the sigmoid; you can train without ever computing the loss, though you would compute it to monitor convergence.
• Real-world extensions: add L2 regularization (grad += (lam/n) * w, usually leaving the bias term unpenalized), standardize features so a single learning rate works across dimensions, and switch to mini-batches for large n.