When a straight line isn't enough

The shape problem, illustrated

A small synthetic example. Rent depends on distance from city centre, but not linearly — the cheapest neighbourhoods are 5-10 km out (suburban), with prices rising again at 15+ km (rural luxury) and far higher in the centre.

import numpy as np
import matplotlib.pyplot as plt

rng = np.random.default_rng(42)
distance = rng.uniform(0, 20, 200)

# A U-shaped relationship: cheap at moderate distances, expensive near and far
rent = 60000 - 8000 * distance + 400 * distance**2 + rng.normal(0, 5000, 200)

plt.scatter(distance, rent, alpha=0.5)
plt.xlabel("Distance from centre (km)")
plt.ylabel("Rent (NPR)")
plt.show()

You see a U-shape. Fit linear regression to this and the model draws a flat-ish line through it, with terrible predictions everywhere — high residuals near 0 km (under-predict) and at 20 km (under-predict) and over-predict in the middle.

Non-linear option 1: k-Nearest Neighbours regression

KNeighborsRegressor works just like KNeighborsClassifier, but instead of voting on a category, it averages the target values of the k nearest neighbours:

from sklearn.neighbors import KNeighborsRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline

model = Pipeline([
    ("scale", StandardScaler()),
    ("regressor", KNeighborsRegressor(n_neighbors=10)),
])
model.fit(X_train, y_train)

For a new flat at distance 12 km, k-NN finds the 10 training flats nearest to 12 km in feature-space, averages their rents, and reports that average as the prediction. No equation, no fitted curve — just local averaging.

On the U-shaped data above, k-NN tracks the curve well. It does not assume linearity; it does not assume any shape. The shape comes from the data.

Strengths of k-NN regression:

• Handles arbitrary non-linear shapes without effort.
• Easy to explain (“we predicted Rs 47,000 because that was the average rent of the 10 most similar past listings”).
• No assumptions about the data’s distribution.

Weaknesses:

• Same as classification: distance-based, so scaling is mandatory.
• Slow on large datasets (every prediction = distance to every training point).
• Poor near the edges of the feature space (the 10 nearest neighbours to a new flat 30 km away are probably 20+ km away, which doesn’t make for a good local average).

Non-linear option 2: decision tree regression

DecisionTreeRegressor builds a tree just like the classifier, but each leaf holds an average of the target values for the training points that fall in that leaf, rather than a class label.

from sklearn.tree import DecisionTreeRegressor

model = DecisionTreeRegressor(max_depth=5, random_state=42)
model.fit(X_train, y_train)

The tree partitions the feature space into rectangular regions, and predicts the mean rent of the training points in each region. The predictions look like a staircase: piecewise-constant within regions, jumping at the splits.

Strengths of decision tree regression:

• Handles arbitrary non-linear shapes and interactions automatically.
• No need to scale features.
• Interpretable (export_text works the same way).
• Robust to outliers (a single weird point doesn’t shift the mean of a region much).

Weaknesses:

• Piecewise-constant predictions can look ugly on smooth real-world relationships.
• Single trees can be brittle — a small change in data shifts the splits.
• Tends to overfit without max_depth (same as the classifier).

A single tree is rarely the final regression model. Random forests and gradient boosting — both ensembles of many trees, both regressor classes in scikit-learn (RandomForestRegressor, GradientBoostingRegressor) — almost always beat a single tree. We cover them in Chapter 5.

A side-by-side comparison

Three regressors on the same U-shaped data:

from sklearn.linear_model import LinearRegression
from sklearn.tree import DecisionTreeRegressor
from sklearn.neighbors import KNeighborsRegressor
from sklearn.metrics import mean_absolute_error
from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    distance.reshape(-1, 1), rent, test_size=0.2, random_state=42,
)

models = {
    "Linear":  LinearRegression(),
    "k-NN":    KNeighborsRegressor(n_neighbors=10),
    "Tree":    DecisionTreeRegressor(max_depth=5, random_state=42),
}

for name, model in models.items():
    model.fit(X_train, y_train)
    pred = model.predict(X_test)
    mae = mean_absolute_error(y_test, pred)
    print(f"{name:8s}  MAE: Rs {mae:.0f}")

Output (illustrative):

Linear    MAE: Rs 9,820
k-NN      MAE: Rs 4,310
Tree      MAE: Rs 5,140

Linear regression is doing its best with a straight line through a U-shape. k-NN and the tree both shape themselves to the curve and predict roughly half as wrong.

The feature-engineering escape hatch

You don’t always need a non-linear model. Sometimes adding the right derived features lets a linear model fit a non-linear pattern:

X_engineered = pd.DataFrame({
    "distance": distance,
    "distance_squared": distance**2,
})

model = LinearRegression()
model.fit(X_engineered, rent)

The model can now combine a positive weight on distance_squared with a negative weight on distance to produce a U-shape. This works because the model is still linear in its weights, but the features include non-linear transformations.

This is the basis expansion trick, and a lot of classical statistics is built on it. Polynomial features, interaction terms (size * location_score), and log-transformations all fit here. scikit-learn has PolynomialFeatures to generate these automatically:

from sklearn.preprocessing import PolynomialFeatures

poly = PolynomialFeatures(degree=2, include_bias=False)
X_poly = poly.fit_transform(X)

This generates all interactions and squares up to degree 2. Combined with linear regression, you get a surprisingly flexible model that is still interpretable.

Predictions outside the training range

A subtle but important difference between models:

• Linear regression extrapolates — predicts confidently for inputs outside the training range, by extending the line.
• k-NN and tree regressors clamp — predictions for inputs outside the training range just match the nearest training neighbours, since there’s nothing else to average from.

Whether extrapolation is helpful depends on the problem. For a model predicting tomorrow’s rent from this year’s listings, modest extrapolation is fine. For a model predicting rent at a building twice as large as any in training data, linear regression will produce a confident number that may be wildly wrong — and you should not trust it.

When you ship a regression model, always document the training-data range, and warn callers when their input is outside it.

Check your understanding

Quick check

—

A linear regression for predicting NEPSE close from recent trading days has an R² of 0.31 and clearly curved residuals. A teammate suggests just throwing in a quadratic feature (`day_index**2`). Will that work?

1. ANo — linear regression can't fit curves under any circumstances.
2. BYes, it can help: adding squared (or other transformed) features lets a linear regression fit non-linear shapes while keeping the model linear in its parameters. This is the basis-expansion trick. It often produces strong, interpretable models without needing to switch to a non-linear model class. Try it; if the residuals straighten out, you're done. If they don't, then try a non-linear model.
3. CYes, but only if you also remove all the linear features first.
4. DQuadratic features are deprecated in scikit-learn.

Check answer→

Quick check

—

A teammate trains a decision tree regressor with no max_depth and reports a near-perfect training MAE. They are surprised that the test MAE is twice as large. What's happening and what's the fix?

1. ADecision trees are not suitable for regression; switch to linear regression.
2. BAn unrestricted decision tree memorises the training data — each training point can end up in its own leaf, where the prediction is exactly that point's target value. Training MAE is therefore near zero, but the tree is overfit and generalises poorly. Fix: set `max_depth` (typically 3 to 8 for tabular data), or use cross-validation to pick a good depth. Chapter 5 covers this systematically.
3. CThe test set is corrupted.
4. DUse a smaller training set.

Check answer→

What comes next

Chapter 3 is done. We’ve built two regressors, a worked rent-prediction project, and learned what to reach for when a line is not enough. Chapter 4 is the most important conceptual chapter of Course 03: how to honestly evaluate a model. Accuracy lies. Precision and recall tell more honest stories. Confusion matrices reveal the asymmetries. Overfitting hides in plain sight if you don’t know where to look. Without Chapter 4, none of the metrics you’ve seen so far are trustworthy.