# Gradient based optimization of step function w.r.t number of steps

I am trying to optimize the parameter b in the following simple function using gradient descent in PyTorch:

$$y = \frac{\lfloor{xb} \rfloor + 0.5}{b}$$

x is in $$[0,1]$$ and b is continuous and in $$[5,30]$$. With fixed b, y(x) is a multi-step function with b steps. This function is applied pixel-wise to multiple images as a color-quantization effect to create an artistic style. From a given set of original and quantized images, I want to be able to estimate the number of quantization bins (b). However, as the floor function has gradient 0 almost everywhere, optimization with gradient descent is not straightforward.

I have tried several ideas to overcome this limitation:

1. Replace the gradient of the floor function with identity (straight-through): Let $$y \approx \frac{{xb} + 0.5}{b}$$ then: $$\frac{\partial y}{\partial b} \approx \frac{- 0.5}{b^2}$$
2. Continuous approximation 1: Replacing $$\lfloor{xb} \rfloor$$ with a differentiable function. I tried $$f(x) = x - \frac{\sin(2 \pi x)}{(2 \pi)}$$ and then $$\lfloor{xb} \rfloor \approx f(f(xb))$$ according to this thread.
3. Continuous approximation 2 according to this thread: $$\lfloor{xb} \rfloor \approx xb - 0.5 - \frac{\tan^{-1}\big(\frac{-\gamma \sin(2 \pi xb)}{1 - \gamma \cos(2 \pi xb)}\big)}{\pi}$$

Unfortunately, it seems like none of these methods really works. I also tried to add gradient noise, change the learning rate, or use a different optimizer (Adam, SGD, RMSprop). I ran multiple experiments using random values of b to generate the training data and performed 500 gradient steps. See below for 2 typical examples using a continuous approximation (number 2) on the left and the straight-through gradient on the right. The straight line is the target value for b and below I plot the parameter value b after each gradient step. Clearly, the optimization gets stuck long before the estimated parameter reaches the target.  It seems puzzling that the performance is so poor considering that even just by looking at quantized images it is easy to make an educated guess about the number of different colors (quantization bins) in the processed image. Maybe I am doing something wrong?

• It does take continuous values in the interval. I think if it were discrete values one could probably just try them all. However, this quantization filter is a part of a larger image stylization stack, which I am trying to optimize altogether. This is the reason why I am trying to solve this with gradient descend if anyhow possible. Apr 7 at 20:20

The minimizing value of $$b \in [5, 30]$$ for a fixed $$x \in [0, 1]$$ is given by $$b = \begin{cases} 30, \text{ if } x < \frac{1}{30} \\ \frac{\lceil 5x \rceil}{x} - \epsilon, \text{ if } \frac{1}{30} \le x \le 1 \end{cases}$$ where $$\epsilon > 0$$ is a small positive number, e.g. $$\epsilon = 0.000001$$ (more on that below: strictly speaking this function has no minimum when $$x \ge \frac{1}{30}$$)

### Explanation:

This function is a great example of why gradient descent does not always work. Take $$x = 1/2$$ and look at the graph of $$y$$: Notice that the gradient (derivative) is always negative when it exists, and hence if we take a gradient descent step (following the direction of the negative gradient), we will take a step to the right. With a fixed step size and restricted to the domain $$[5, 30]$$, gradient descent will converge to $$b = 30$$, which is obviously not the minimizer for this particular $$x$$. A "smoothed" version of the function like you were considering will still be "sawtooth-shaped" and have multiple basins of attraction for gradient descent, making convergence highly dependent on the parameters (step-size, stopping criteria, smoothing method, etc.) and initialization.

Luckily, this problem can be solved analytically without the need for numerical optimization techniques. Take $$x$$ to be a fixed value in $$[0, 1]$$ and let

$$y_x (b) = \frac{\lfloor x b \rfloor + 0.5}{b}.$$

If $$0 \le x < \frac{1}{30}$$, then $$0 \le xb < 1$$ and hence $$\lfloor x b \rfloor = 0$$ and $$y_x (b) = \frac{0.5}{b}$$ which is a decreasing function, and is hence minimized at the right endpoint of the domain, i.e. $$b = 30$$.

When $$x \ge 1/30$$, there will be at least one discontinuity. Intuitively, when $$\lfloor x b \rfloor$$ has a fixed value, $$y_x$$ is decreasing, so $$y_x$$ is decreasing almost everywhere. The only places where $$y_x$$ increases is at the discontinuities, thus the minimizer (by minimizer I mean the $$b$$-value that obtains the minimum) must be at one of the right endpoints of the intervals where $$y_x$$ is continuous. Unfortunately, this function is actually right-continuous, so the apparent local minima on the graph are not actually on the graph (if I could graph open circles at the local minimum points I would: the filled-in circles would be at the local maximum points). For this reason, the minimum does not exist, but we can still find an approximate minimum. The $$b$$-coordinates of the would-be minima (meaning a limit point of the graph that is not included in the graph but would be a minimum if included) are where $$x b$$ is a positive integer $$n \in \mathbb{N}$$, hence $$b = n/x$$. This is now a sequence, and let's denote it by $$b_n$$ and denote the sequence of would-be local minimum $$y$$-values by $$y_n$$. Note that $$y_n$$ is increasing because $$y_n = \frac{\lfloor x (n/x) \rfloor_- + 0.5}{n/x} = \frac{n - 1 + 0.5}{n/x} = x - \frac{x}{2n}$$ where $$\lfloor \cdot \rfloor_-$$ is the left-continuous analogue of the floor function. Thus the would-be minimum value is the left-most would-be local minimum value in the domain $$[5, 30]$$. Therefore, the would-be minimizer is given by $$n/x$$ where $$n$$ is the smallest number such that $$n/x \ge 5$$, which implies that $$n = \lceil 5x \rceil$$, and hence $$b = \frac{\lceil 5x \rceil}{x}$$.

As a quick sanity check, using the value $$x = 1/2$$ to match the graph above, this comes out to be $$b = \frac{\lceil 5/2 \rceil}{1/2} = 6$$ which exactly matches the would-be minimizer we see in the graph. Remember that this is a would-be (meaning a limit point of the graph that is not included in the graph) minimum. If you forget and plug in $$b = 6$$, you will get $$y_{1/2} (6) = 3.5/6 = 0.58\overline{3}$$ which is actually the maximum of the function. Instead, pick $$\epsilon = 0.000001$$ or something like that and plug in $$y_{1/2} (5.999999) = 2.5/6 = 0.41\overline{6}$$ which corresponds to the would-be minimum of the graph above.