matrix scaling
3D matrix
data preprocessing
standardization
data normalization

How to standard scale a 3D matrix?

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Standard scaling means subtracting the mean and dividing by the standard deviation. With a 3D matrix, the hard part is not the formula but deciding what each axis means. In many machine-learning pipelines the array shape is something like (samples, timesteps, features), and the usual goal is to standardize each feature consistently across all observed samples and timesteps.

Start by Interpreting the Axes

A 3D array can represent very different things:

  • '(samples, timesteps, features) for sequence data'
  • '(batch, height, width) for image-like tensors'
  • '(items, channels, measurements) for sensor data'

You cannot scale correctly until you know which axis should share the same mean and standard deviation.

If the last axis is the feature axis, a common approach is to compute one mean and one standard deviation per feature using all observations from the first two axes.

Reshape to 2D, Scale, Then Reshape Back

StandardScaler expects shape (observations, features). For a tensor shaped (samples, timesteps, features), flatten the first two axes into one observation axis.

python
1import numpy as np
2from sklearn.preprocessing import StandardScaler
3
4rng = np.random.default_rng(42)
5X = rng.normal(loc=[10.0, 100.0, -5.0], scale=[2.0, 20.0, 0.5], size=(4, 3, 3))
6
7n_samples, n_steps, n_features = X.shape
8X_2d = X.reshape(-1, n_features)
9
10scaler = StandardScaler()
11X_scaled_2d = scaler.fit_transform(X_2d)
12X_scaled = X_scaled_2d.reshape(n_samples, n_steps, n_features)
13
14print(X_scaled.shape)
15print(X_scaled.reshape(-1, n_features).mean(axis=0).round(6))
16print(X_scaled.reshape(-1, n_features).std(axis=0).round(6))

This keeps the feature columns intact while preserving the original 3D structure in the final output.

Fit Only on Training Data

Just like 2D preprocessing, fit the scaler on training data only and reuse it for validation, test, and inference.

python
1import numpy as np
2from sklearn.preprocessing import StandardScaler
3
4X_train = np.random.randn(50, 10, 4)
5X_test = np.random.randn(10, 10, 4)
6
7n_features = X_train.shape[-1]
8train_2d = X_train.reshape(-1, n_features)
9test_2d = X_test.reshape(-1, n_features)
10
11scaler = StandardScaler().fit(train_2d)
12X_train_scaled = scaler.transform(train_2d).reshape(X_train.shape)
13X_test_scaled = scaler.transform(test_2d).reshape(X_test.shape)

If you fit on the full dataset, the preprocessing step leaks information from held-out data.

Manual NumPy Scaling Works Too

If you want explicit control or do not want to depend on scikit-learn, use broadcasting directly.

python
1import numpy as np
2
3X = np.random.randn(8, 5, 2)
4mean = X.mean(axis=(0, 1), keepdims=True)
5std = X.std(axis=(0, 1), keepdims=True)
6X_scaled = (X - mean) / std
7
8print(mean.shape)
9print(std.shape)
10print(X_scaled.shape)

keepdims=True keeps the feature axis aligned so the subtraction and division broadcast correctly.

The "Right" Scaling Rule Depends on the Problem

Feature-wise global scaling is common, but it is not universal. Sequence models sometimes use per-sequence normalization when relative change matters more than absolute level. Image models often use fixed dataset-wide channel statistics rather than generic feature scaling.

So the shape alone does not tell you the correct preprocessing rule. The meaning of the data does.

In real pipelines, save the fitted scaler alongside the model so inference uses the same statistics as training. Scaling rules are part of the model contract, not a disposable preprocessing detail.

Common Pitfalls

  • Reshaping the wrong axis so columns no longer represent features.
  • Fitting the scaler on validation or test data.
  • Forgetting to reshape the transformed data back to 3D.
  • Using feature-wise scaling when the problem actually needs per-sequence normalization.

Summary

  • Standard-scaling a 3D matrix starts with understanding what each axis represents.
  • For (samples, timesteps, features), flatten to 2D, scale by feature, then reshape back.
  • Fit the scaler on training data only.
  • NumPy broadcasting works well when you want explicit control.
  • Good scaling depends on data semantics, not just array rank.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.