diff --git a/src/linreg.odin b/src/linreg.odin
index 0c66bee..7475b94 100644
--- a/src/linreg.odin
+++ b/src/linreg.odin
@@ -62,7 +62,6 @@ solve_linear_system :: proc(A: [][]f64, b: []f64) -> []f64 {
 train_linear_regression_cached :: proc(selected_features: []int, y: []f64) -> []f64 {
     m := len(selected_features)
 
-    // Extract submatrix from cache
     XtX := make([][]f64, m)
     defer {
         for row in XtX {delete(row)}
@@ -74,6 +73,8 @@ train_linear_regression_cached :: proc(selected_features: []int, y: []f64) -> []
         for j in 0 ..< m {
             XtX[i][j] = feature_state.XtX_cache[selected_features[i]][selected_features[j]]
         }
+        // Add ridge regularization to diagonal
+        XtX[i][i] += 0.01  // Small regularization term
     }
 
     Xty := make([]f64, m)
@@ -85,6 +86,7 @@ train_linear_regression_cached :: proc(selected_features: []int, y: []f64) -> []
     return solve_linear_system(XtX, Xty)
 }
 
+
 predict :: proc(X: [][]f64, beta: []f64) -> []f64 {
 	predictions := make([]f64, len(X))
 	for i in 0 ..< len(X) {