TensorFlow
optimization
constrained optimization
machine learning
deep learning

Tensorflow how to minimize under constraints

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

TensorFlow is very good at unconstrained gradient-based optimization, but most real systems have constraints: weights must stay nonnegative, probabilities must sum to one, or a parameter must remain inside a physical range. There is no single built-in "constrained minimize" switch for every case, so the usual solution is to encode the constraint into the parameterization or into the loss.

Choose the Constraint Strategy First

Before writing code, decide what kind of constraint you have:

  • A bound such as x >= 0
  • A box such as 0 <= x <= 1
  • A sum constraint such as probabilities adding to 1
  • A soft business rule that can be violated, but should be penalized

The strategy depends on the type. Hard constraints are often easiest to enforce by reparameterizing the variable. Soft constraints are often handled with penalty terms added to the objective.

Reparameterize for Hard Constraints

If a variable must stay nonnegative, optimizing an unconstrained raw variable and passing it through softplus is a common TensorFlow pattern.

python
1import tensorflow as tf
2
3raw_w = tf.Variable(0.0)
4optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)
5
6target = tf.constant(3.0)
7
8for step in range(100):
9    with tf.GradientTape() as tape:
10        w = tf.nn.softplus(raw_w)
11        loss = tf.square(w - target)
12
13    grads = tape.gradient(loss, [raw_w])
14    optimizer.apply_gradients(zip(grads, [raw_w]))
15
16print("constrained w:", tf.nn.softplus(raw_w).numpy())

raw_w can take any real value, but the actual model parameter w is always positive. That means the optimizer never has to be corrected after the update.

Use Projection for Simple Bounds

Sometimes a direct projection step is clearer. After each gradient update, clip the value back into the allowed interval.

python
1import tensorflow as tf
2
3x = tf.Variable(2.5)
4optimizer = tf.keras.optimizers.SGD(learning_rate=0.2)
5
6for step in range(50):
7    with tf.GradientTape() as tape:
8        loss = tf.square(x - 0.3)
9
10    grad = tape.gradient(loss, x)
11    optimizer.apply_gradients([(grad, x)])
12    x.assign(tf.clip_by_value(x, 0.0, 1.0))
13
14print("bounded x:", x.numpy())

This is projected gradient descent. It is easy to implement and works well for box constraints, although the projection can interact with the optimizer state in ways that make convergence less smooth than reparameterization.

Add Penalty Terms for Soft Constraints

If violating the constraint is allowed but undesirable, add a penalty to the loss. Suppose you want x + y to stay close to 1:

python
1import tensorflow as tf
2
3x = tf.Variable(0.8)
4y = tf.Variable(0.8)
5optimizer = tf.keras.optimizers.Adam(learning_rate=0.05)
6
7for step in range(200):
8    with tf.GradientTape() as tape:
9        objective = tf.square(x - 0.2) + tf.square(y - 0.7)
10        penalty = 10.0 * tf.square((x + y) - 1.0)
11        loss = objective + penalty
12
13    grads = tape.gradient(loss, [x, y])
14    optimizer.apply_gradients(zip(grads, [x, y]))
15
16print("x:", x.numpy(), "y:", y.numpy(), "sum:", (x + y).numpy())

The coefficient on the penalty controls how strongly the optimizer respects the rule. Too small, and the constraint is weak. Too large, and optimization becomes numerically awkward.

Constraints Inside Keras Layers

Keras also provides simple built-in constraints for layer weights. This is useful when the constraint belongs to the model definition itself.

python
1import tensorflow as tf
2
3model = tf.keras.Sequential([
4    tf.keras.layers.Dense(
5        4,
6        input_shape=(3,),
7        kernel_constraint=tf.keras.constraints.NonNeg(),
8    )
9])

This approach is convenient for standard cases, but it is narrower than full constrained optimization. It helps with weight properties, not arbitrary coupled constraints between multiple variables.

Common Pitfalls

  • Expecting TensorFlow optimizers to solve general constrained optimization automatically usually leads to confusion; you often need to encode the constraint yourself.
  • Using a penalty coefficient that is far too small means the model optimizes the original objective and effectively ignores the constraint.
  • Using a huge penalty coefficient can make gradients unstable and training difficult to tune.
  • Clipping variables after every step is simple, but it can hide the fact that the optimizer keeps trying to push the variable outside the feasible region.
  • Forgetting that some constraints are better expressed by parameterization than by penalties leads to unnecessarily fragile code.

Summary

  • In TensorFlow, constrained minimization is usually implemented with reparameterization, projection, or penalty terms.
  • Reparameterization is often the cleanest solution for hard constraints such as positivity.
  • Projection with clip_by_value works well for simple bounds.
  • Penalty methods are useful for soft or coupled constraints, but the penalty scale must be tuned carefully.

Related reading
Course
Intermediate
27 lessons
15 hours
DSA Fundamentals

Master algorithmic patterns and data structures through hands-on LeetCode-style problems - from arrays and hashing to dynamic programming and advanced graphs.

View the 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.