Show that the weighted average of a vector is convex when the weights are provided by the softmax function

Question

I am working on a problem in the design of computer chips. One of the papers (TSV-aware analytical placement for 3D IC designs, DOI:https://doi.org/10.1145/2024724.2024875) introduces a function to approximate the vector-maximum function: $ \mathop{\rm WA} \colon \mathbb{R}^n \to \mathbb{R}$ defined by

$$ \mathop{\rm WA}(\mathbf{x}) = \frac{\sum_{k=1}^n x_k \exp(x_k)}{\sum_{k=1}^n \exp(x_k)} $$

The name WA comes from the fact that this is a weighted average of $\mathbf{x}$, where the weights are determined by the softmax function $\sigma(\mathbf{x})$. This gives the shorter expression $ \mathop{\rm WA}(\mathbf{x}) = \sigma(\mathbf{x})^T \mathbf{x}$. How can I show that this function is convex?

What I have tried:

1. Find a proof in the literature
2. Find a proof that the Hessian is positive semi-definite
3. Find a direct proof based on the definition of convexity

Approach 1: In the paper linked above the authors claim "It can be shown that the WA wirelength model is strictly convex and continuously differentiable by differentiating [...] twice" but they don't do this explicitly. Other papers I stumbled upon only reference the result in this paper without any proof.

Approach 2: I did the work and differentiated the function to determine the gradient $\nabla \mathop{\rm WA}$ and the Hessian $\nabla^2 \mathop{\rm WA}$: \begin{align*} \nabla \mathop{\rm WA} (\mathbf{x}) &= \sigma(\mathbf{x}) \odot \left( \mathbf{1} + \mathbf{x} - \mathop{\rm WA}(\mathbf{x}) \cdot \mathbf{1} \right) \\ \nabla^2 \mathop{\rm WA} (\mathbf{x}) &= \mathop{\rm diag} \left( \sigma(\mathbf{x}) \odot \left( 2 \cdot \mathbf{1} + \mathbf{x} - \mathop{\rm WA}(\mathbf{x}) \cdot \mathbf{1} \right) \right) \\ &\phantom{=} - \sigma(\mathbf{x}) \sigma(\mathbf{x})^T \odot \left( 2 \cdot \mathbf{1} \mathbf{1}^T + \mathbf{1} \mathbf{x}^T + \mathbf{x} \mathbf{1}^T - 2 \mathop{\rm WA}(\mathbf{x}) \cdot \mathbf{1} \mathbf{1}^T \right) \end{align*} Here, $\odot$ is the element-wise multiplication, $\mathbf{1}$ is the all-ones vector in $\mathbb{R}^n$ and $\mathop{\rm diag}(\mathbf{y})$ is the matrix with the elements of $\mathbf{y}$ on the diagonal and zeros elsewhere. Now to prove convexity through the Hessian I would have to show that $\mathbf{v}^T (\nabla^2 \mathop{\rm WA}(\mathbf{x})) \mathbf{v} \geq 0$ for all $ \mathbf{v}, \mathbf{x} \in \mathbb{R}^n$. Expanding this gives

$$ \sum_{k=1}^n v_k^2 \sigma(\mathbf{x})_k (2 + x_k - \mathop{\rm WA}(\mathbf{x})) - \sum_{i=1}^n \sum_{j=1}^n v_i \sigma(\mathbf{x})_i v_j \sigma(\mathbf{x})_j (2 + x_i + x_j - 2 \mathop{\rm WA}(\mathbf{x})) \geq 0 $$

and I have not found a way to show this. I can only eliminate both 2s at the start of each bracket with the help of the Cauchy-Schwarz inequality.

Approach 3: I did not get far with this, I can't report any useful progress.

Remark about strict convexity: It seems to me that the authors of the paper have made a mistake when claiming that the function is strictly convex. $\mathop{\rm WA}(t \cdot \mathbf{1}) = t$ for all $ t \in \mathbb{R}$ and so the function is linear on a line and cannot be strictly convex.

Remark: I might have made mistakes at any point of the way. I might have overlooked important literature and I might have found the wrong derivatives. Any help here is appreciated.

Answer 1

The function is not convex For $n=2$ with function arguments $x$ and $y$, the second derivative to $x$ is: $$\frac{e^{x+y}\left[e^y(x-y+2)+e^x(-x+y+2)\right]}{(e^x+e^y)^3}$$ which is nonnegative iff $$e^y(x-y+2)+e^x(y-x+2) \geq 0,$$ but it is negative for $x=2$ and $y=-2$. Indeed, if you plot the function with $y$ fixed at $-2$, you can see that the function is not convex.

A commonly used convex approximation for maximum function is the log-sum-exp function.

Answer 2

After LinAlg put me on the right track by showing that the function is not convex I investigated its properties some more and now I am even able to show that the function is not convex for any $n \geq 2$:

Let $\mathbf{x} = (a, 0, \ldots, 0) \in \mathbb{R}^n$ then $\sigma(\mathbf{x})_1 = \frac{e^a}{e^a + n-1}$, $\mathop{\mathrm{WA}}(\mathbf{x}) = \sigma(\mathbf{x})_1 a$ and \begin{align} e_1^T (\nabla^2 \mathop{\mathrm{WA}}) e_1 &= \sigma(\mathbf{x})_1 (2 + a - \mathop{\mathrm{WA}}(\mathbf{x})) - \sigma(\mathbf{x})_1^2 (2+2a-2\mathop{\mathrm{WA}}(\mathbf{x})) \\ &= \sigma(\mathbf{x})_1 (2 + a - \sigma(\mathbf{x})_1a) - \sigma(\mathbf{x})_1^2 (2+2a-2\sigma(\mathbf{x})_1a) \\ &= \sigma(\mathbf{x})_1 (2 - 2 \sigma(\mathbf{x})_1 + (1 - 3\sigma(\mathbf{x})_1 + 2\sigma(\mathbf{x})_1^2)a ) \end{align} For the sign the positive factor $\sigma(\mathbf{x})_1$ does not matter so we only consider the second factor and see $$ \lim_{a \to -\infty} 2 - \underbrace{2 \sigma(\mathbf{x})_1}_{\to 0} + \underbrace{(1 - 3\sigma(\mathbf{x})_1 + 2\sigma(\mathbf{x})_1^2)}_{\to 1} \underbrace{a}_{\to -\infty} = - \infty $$

This proves that there is an $\mathbf{x}$ such that $\nabla^2 \mathop{\mathrm{WA}}(\mathbf{x})$ is not positive semidefinite so $\mathop{\mathrm{WA}}$ is not convex.

Edit:

I want to add another proof: Let $\mathbf{x} = (a, 0, \ldots, 0) \in \mathbb{R}^n$ then $\sigma(\mathbf{x})_1 = \frac{e^a}{e^a + n-1}$, $\mathop{\mathrm{WA}}(\mathbf{x}) = \sigma(\mathbf{x})_1 a$ and \begin{align} e_1^T (\nabla^2 \mathop{\mathrm{WA}}) e_1 &= \sigma(\mathbf{x})_1 (2 + a - \mathop{\mathrm{WA}}(\mathbf{x})) - \sigma(\mathbf{x})_1^2 (2+2a-2\mathop{\mathrm{WA}}(\mathbf{x})) \\ &= \sigma(\mathbf{x})_1 (2 + a - \sigma(\mathbf{x})_1a) - \sigma(\mathbf{x})_1^2 (2+2a-2\sigma(\mathbf{x})_1a) \\ &= \sigma(\mathbf{x})_1 ((2+a) - (3a+2)\sigma(\mathbf{x})_1 + (2a) \sigma(\mathbf{x})_1^2) \end{align}

Now we take a look at the roots of $(2+a) - (3a+2)b + (2a) b^2$ when $b$ is the variable. For every $a \notin \{0, 2\}$ there are exactly two roots: $1$ and $1/a + 1/2$. If $a > 2$ (and thereby $a > 0$) then the quadratic function is negative inside of $[1/a + 1/2, 1]$. If $a < -2$ (and thereby $a < 0$) then the quadratic is negative outside of $[1/a + 1/2, 1]$. When we replace $b$ by $\sigma(\mathbf{x})_1$ we get conditions for $a$ that ensure that the Hessian $\nabla^2 \mathop{\mathrm{WA}}$ is not positive semidefinite: \begin{cases} \frac{e^a}{e^a + n-1} > \frac{1}{a} + \frac{1}{2} & \text{and } a > 2 \\ \frac{e^a}{e^a + n-1} < \frac{1}{a} + \frac{1}{2} & \text{and } a < -2 \end{cases} Taking the reciprocal on both sides gives us \begin{cases} 1 + \frac{n-1}{e^a} < \frac{2a}{2+a} & \text{and } a > 2 \\ 1 + \frac{n-1}{e^a} > \frac{2a}{2+a} & \text{and } a < -2 \end{cases} Because $$ \lim_{a \to \infty} 1 + \frac{n-1}{e^a} = 1, \quad \lim_{a \to -\infty} 1 + \frac{n-1}{e^a} = \infty \quad \text{and} \quad \lim_{a \pm \infty} \frac{2a}{2+a} = 2$$ one of those conditions holds for all values of $a$ with large enough absolute value. $\square$

This proof is also able to show that the function $\mathop{\mathrm{WA}}(\mathbf{x}) + \mathop{\mathrm{WA}}(-\mathbf{x})$ is not convex.