Skip to main content

Linear Regression with TensorFlow

A Keras Dense(1) layer with no activation computes the linear predictor Xw+bXw+b. This is useful when a linear head belongs to a larger TensorFlow model. For ordinary tabular least squares, a dedicated linear-model library is usually simpler and provides a direct solution.

1. Split, then standardize​

All data are generated locally. The scaler learns its mean and variance from the training partition only.

import numpy as np
import tensorflow as tf
from sklearn.datasets import make_regression
from sklearn.metrics import root_mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

tf.keras.utils.set_random_seed(42)
tf.config.threading.set_inter_op_parallelism_threads(1)
tf.config.threading.set_intra_op_parallelism_threads(1)

X, y = make_regression(
    n_samples=240, n_features=4, noise=4.0, random_state=42
)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=42
)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train).astype("float32")
X_test = scaler.transform(X_test).astype("float32")

model = tf.keras.Sequential([
    tf.keras.Input(shape=(X_train.shape[1],)),
    tf.keras.layers.Dense(1),
])
model.compile(
    optimizer=tf.keras.optimizers.Adam(learning_rate=0.05),
    loss=tf.keras.losses.MeanSquaredError(),
)
model.fit(X_train, y_train, epochs=120, batch_size=32, verbose=0, shuffle=False)
y_pred = model.predict(X_test, verbose=0).squeeze(1)

rmse = root_mean_squared_error(y_test, y_pred)
assert np.isfinite(rmse) and y_pred.shape == y_test.shape
print(f"Test RMSE: {rmse:.2f}")

2. What the model is—and is not​

The layer is linear because no activation follows it. MeanSquaredError supplies the OLS objective, while Adam finds an approximate numerical optimum. Training for more epochs is not a substitute for checking validation behavior, residual structure, or whether a linear conditional mean is appropriate.

Keras does not automatically provide classical coefficient standard errors or confidence intervals. Use a statistical model when inference is the goal. For prediction, retain the scaler, validate input schema, and treat out-of-range inputs as extrapolation.