convexNN_poly_act.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     0
    ]
   },
   "outputs": [],
   "source": [
    "# import packages \n",
    "import numpy as np\n",
    "import time\n",
    "import cvxpy as cp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "# paper link: https://arxiv.org/pdf/2101.02429.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "code_folding": [
     2,
     10,
     23,
     34,
     81,
     135
    ]
   },
   "outputs": [],
   "source": [
    "# poly act scalar case\n",
    "# main function\n",
    "def convexNN_poly_act_solver(A, y, beta, scs_max_iters, loss_type, a, b, c):\n",
    "    d = A.shape[1]\n",
    "    X_V = compute_X_V_base_FC(A, d, verbose)\n",
    "    X_KTX_K, X_KTy, y_normsq, X_V_scaled = scale_X_K(X_V, y, a, b, c, d)\n",
    "    Z, Z_prime, noncvx_cost = solve_cvx(X_KTX_K, X_KTy, y_normsq, d, beta, scs_max_iters, loss_type, X_V_scaled, y)\n",
    "    return Z, Z_prime, noncvx_cost\n",
    "\n",
    "# helper functions\n",
    "def compute_X_V_base_FC(A, d, verbose=False):\n",
    "    n = A.shape[0]\n",
    "\n",
    "    X_V = np.zeros((n, d**2+d+1))\n",
    "    for i in range(n):\n",
    "        if i % 100 == 0 and verbose:\n",
    "            print(i, end=\", \")\n",
    "        \n",
    "        x_i = A[i:i+1,:].T\n",
    "        X_V[i, 0:d**2] = np.matmul(x_i, x_i.T).reshape((d**2))\n",
    "        X_V[i, d**2:d**2+d] = x_i.reshape((d))\n",
    "        X_V[i, d**2+d] = 1\n",
    "    return X_V\n",
    "def scale_X_K(X_K, y, a, b, c, ff): # scale, here, refers to multiplying the columns by a,b,c\n",
    "    X_K_scaled = X_K.copy()\n",
    "    X_K_scaled[:, 0:ff**2] = a * X_K_scaled[:, 0:ff**2]\n",
    "    X_K_scaled[:, ff**2:ff**2+ff] = b * X_K_scaled[:, ff**2:ff**2+ff]\n",
    "    X_K_scaled[:, ff**2+ff] = c\n",
    "    \n",
    "    X_KTX_K = np.matmul(X_K_scaled.T, X_K_scaled)\n",
    "    X_KTy = np.matmul(X_K_scaled.T, y)\n",
    "    y_normsq = np.sum(y**2)\n",
    "    \n",
    "    return X_KTX_K, X_KTy, y_normsq, X_K_scaled\n",
    "def solve_cvx(X_KTX_K, X_KTy, y_normsq, ff, beta, SCS_max_iters, loss_name, X_V=None, y=None):\n",
    "    # poly act scalar output\n",
    "    Z1 = cp.Variable((ff, ff), symmetric=True)\n",
    "    Z2 = cp.Variable((ff, 1))\n",
    "    Z4 = cp.Variable((1,1))\n",
    "\n",
    "    Z1_prime = cp.Variable((ff, ff), symmetric=True)\n",
    "    Z2_prime = cp.Variable((ff,1))\n",
    "    Z4_prime = cp.Variable((1,1))\n",
    "\n",
    "\n",
    "    zz = cp.vstack((cp.reshape((Z1-Z1_prime), (ff**2,1)), (Z2-Z2_prime), (Z4-Z4_prime)))\n",
    "    \n",
    "    \n",
    "    yhat = X_V@zz\n",
    "    if loss_name == \"squared_loss\":\n",
    "        objective = 0.5*cp.sum_squares(yhat - y) + beta*(Z4 + Z4_prime)\n",
    "    elif loss_name == \"l1_loss\":\n",
    "        objective = cp.sum(cp.abs(yhat - y)) + beta*(Z4 + Z4_prime)\n",
    "    elif loss_name == \"huber\":\n",
    "        objective = cp.sum(cp.huber(yhat-y)) + beta*(Z4 + Z4_prime)\n",
    "\n",
    "\n",
    "    Z = cp.vstack((cp.hstack((Z1, Z2)), cp.hstack((Z2.T, Z4))))\n",
    "    Z_prime = cp.vstack((cp.hstack((Z1_prime, Z2_prime)), cp.hstack((Z2_prime.T, Z4_prime))))\n",
    "    \n",
    "    constraints = [Z >> 0] + [Z_prime >> 0]\n",
    "\n",
    "    constraints += [cp.trace(Z1) == Z4]\n",
    "    constraints += [cp.trace(Z1_prime) == Z4_prime]\n",
    "    \n",
    "    prob = cp.Problem(cp.Minimize(objective), constraints)\n",
    "    start_time = time.time()\n",
    "    print(\"started..\")\n",
    "    #prob.solve(solver=cp.SCS, warm_start=False, max_iters=SCS_max_iters)\n",
    "    prob.solve(solver=cp.MOSEK)\n",
    "    end_time = time.time()\n",
    "    print(\"time elapsed: \" + str(end_time - start_time))\n",
    "\n",
    "    # Print result.\n",
    "    print(prob.status)\n",
    "    print(\"The optimal value is\", prob.value)\n",
    "    print(\"The optimal value is\", objective.value)\n",
    "    \n",
    "    return Z, Z_prime, objective.value\n",
    "\n",
    "# neural decomposition function\n",
    "def neural_decomposition(Z_decomp, tolerance=10**(-9)):\n",
    "    # decomposes Z_decomp as a sum of r (where r=rank(Z_decomp)) rank-1 matrices \\sum_{j=1}^r y_jy_j^T where\n",
    "    # y_j^TGy_j = 0 for all j=1,...,r\n",
    "    # based on the alg given in the proof of lemma 2.4 of the paper 'A Survey of the S-Lemma'\n",
    "    G = np.identity(Z_decomp.shape[0])\n",
    "    G[-1,-1] = -1\n",
    "\n",
    "    # step 0\n",
    "    evals, evecs = np.linalg.eigh(Z_decomp)\n",
    "    # some eigvals are negative due to numerical issues, tolerance masking deals with that\n",
    "    ind_pos_evals = (evals > tolerance)\n",
    "    p_i_all = evecs[:,ind_pos_evals] * np.sqrt(evals[ind_pos_evals])\n",
    "\n",
    "    outputs_y = np.zeros(p_i_all.shape)\n",
    "\n",
    "    for i in range(outputs_y.shape[1]-1):\n",
    "        # step 1\n",
    "        p_1 = p_i_all[:,0:1]\n",
    "        p_1Gp_1 = np.matmul(p_1.T, np.matmul(G, p_1))\n",
    "\n",
    "        if p_1Gp_1 == 0:\n",
    "            y = p_1.copy()\n",
    "\n",
    "            # update\n",
    "            p_i_all = np.delete(p_i_all, 0, 1) # delete the first column\n",
    "        else:\n",
    "            for j in range(1, p_i_all.shape[1]):\n",
    "                p_j = p_i_all[:,j:j+1]\n",
    "                p_jGp_j = np.matmul(p_j.T, np.matmul(G, p_j))\n",
    "                if p_1Gp_1 * p_jGp_j < 0:\n",
    "                    break\n",
    "\n",
    "            # step 2\n",
    "            p_1Gp_j = np.matmul(p_1.T, np.matmul(G, p_j))\n",
    "            discriminant = 4*p_1Gp_j**2 - 4*p_1Gp_1*p_jGp_j\n",
    "            alpha = (-2*p_1Gp_j + np.sqrt(discriminant)) / (2*p_jGp_j)\n",
    "            y = (p_1 + alpha*p_j) / np.sqrt(1+alpha**2)\n",
    "\n",
    "            # update\n",
    "            p_i_all = np.delete(p_i_all, j, 1) # delete the jth column\n",
    "            p_i_all = np.delete(p_i_all, 0, 1) # delete the first column\n",
    "\n",
    "            u = (p_j - alpha*p_1) / np.sqrt(1+alpha**2)\n",
    "            p_i_all = np.concatenate((p_i_all, u), axis=1) # insert u to the list of p_i's\n",
    "\n",
    "        # save y\n",
    "        outputs_y[:,i:i+1] = y.copy()\n",
    "\n",
    "    # save the remaining column\n",
    "    outputs_y[:, -1:] = p_i_all.copy()\n",
    "    \n",
    "    return outputs_y\n",
    "\n",
    "# forward prop\n",
    "def forward_prop_polyact(X, first_layer_weights, second_layer_weights, a, b, c):\n",
    "    Xu = np.matmul(X, first_layer_weights)\n",
    "    Xu_act = a*(Xu)**2 + b*(Xu) + c\n",
    "    output = np.matmul(Xu_act, second_layer_weights)\n",
    "    return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# parameter selection\n",
    "a = 0.09 # polynomial activation coefficients: a, b, c\n",
    "b = 0.5\n",
    "c = 0.47\n",
    "beta = 10**(-1) # regularization parameter\n",
    "verbose = True\n",
    "scs_max_iters = 50000 # maximum number of iterations for the convex solver\n",
    "tol = 10**(-6) # tolerance parameter\n",
    "\n",
    "loss_type = \"squared_loss\" # pick from: \"squared_loss\", \"huber\", \"l1_loss\" (other convex losses can be implemented)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generate random data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate random data \n",
    "n, d = 100, 20\n",
    "m_pnt = 5 # the number of planted neurons\n",
    "\n",
    "np.random.seed(0)\n",
    "A = np.random.normal(0,1,(n,d))**4\n",
    "\n",
    "u_truth = np.random.normal(0,1,(d,m_pnt))\n",
    "alpha_truth = np.random.normal(0,1,(m_pnt,))\n",
    "noise = 1 * np.random.normal(0,0.1,(n,1))\n",
    "y = np.sum((a*np.matmul(A, u_truth)**2 + b*np.matmul(A, u_truth) + c) * alpha_truth, axis=1, keepdims=True) + noise\n",
    "\n",
    "\n",
    "print(\"A.shape = {}, y.shape = {}\".format(A.shape, y.shape))\n",
    "print(\"num of planted neurons = {}\".format(m_pnt))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Solve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "start_time = time.time()\n",
    "# call to the main solver function\n",
    "Z, Z_prime, noncvx_cost = convexNN_poly_act_solver(A, y, beta, scs_max_iters, loss_type, a, b, c)\n",
    "end_time = time.time()\n",
    "\n",
    "time_elapsed_cvx = end_time - start_time\n",
    "print(\"total time: \" + str(time_elapsed_cvx))\n",
    "\n",
    "num_neurons_cvx = np.sum(np.linalg.eig(Z.value)[0] > tol) + np.sum(np.linalg.eig(Z_prime.value)[0] > tol)\n",
    "print(\"The number of neurons is\", num_neurons_cvx)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Extract weights from the SDP solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# extract the neural network weights from the cvxpy solution\n",
    "tolerance = tol # if the tolerance is too low, 'special_decomposition' might throw an error (NaN error), \n",
    "# in that case choosing a larger tolerance might help\n",
    "decomp = neural_decomposition(Z.value, tolerance)\n",
    "decomp_prime = neural_decomposition(Z_prime.value, tolerance)\n",
    "\n",
    "first_layer_weights = np.concatenate((decomp[:-1, :], decomp_prime[:-1, :]), axis=1)\n",
    "first_layer_weights = first_layer_weights / np.sqrt(np.sum(first_layer_weights**2, axis=0))\n",
    "signs_second_layer = np.sign(np.concatenate((decomp[-1:, :], decomp_prime[-1:, :]), axis=1).T)\n",
    "first_layer_weights = first_layer_weights * signs_second_layer[:,0]\n",
    "second_layer_weights = np.concatenate((decomp[-1:, :]**2, -decomp_prime[-1:, :]**2), axis=1).T\n",
    "\n",
    "print(\"shape of first layer weights: \", first_layer_weights.shape)\n",
    "print(\"shape of second layer weights: \", second_layer_weights.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}