Optimizing sequential computation

Question

The following question came up in reading group discussion of paper -- imagine taking two vectors, multiplying each by a random matrix, then taking tanh pointwise. If you repeat this process forever, does the angle between vectors eventually converge to 0 or Pi?

Mathematica seems like a good way to test this theory computationally, but my approach wasn't fast enough. IE, below is histogram of cosine similarities for 2500 pairs of vectors after applying 200 steps....but I'd like to get this for 20000 steps in reasonable time. Are there some tricks to use to speed this up significantly? (can Mathematica use GPUs?)

(* create properly scaled random matrix *)

randXavier[{rows_, cols_}] := 
  RandomVariate[
   NormalDistribution[0, Sqrt[2/(rows + cols)]], {rows, cols}];
(* normalize each row of matrix to have norm 1 *)

rowNormalize[mat_] := Module[{rowNorms},
   rowNorms = Norm /@ mat;
   DiagonalMatrix[1/rowNorms] . mat
   ];
(* take uppper triangular part of matrix as a vector *)

upperTriangular[mat_] := 
  Statistics`Library`UpperTriangularMatrixToVector[mat];

d = 200; (* depth *)
w = 1024; (* layer dimension *)
b = 50; (* batch \
size *)
mat := randXavier[{w, w}]
x = RandomVariate[NormalDistribution[], {b, w}];
a = Nest[Tanh[# . mat] &, x, d];
a = rowNormalize[a];
gram = a . Transpose[a];
plot1 = Histogram[upperTriangular[gram]]

Answer 1

Have you happened to try FunctionLayer? It provides quite an automation to port normal code to GPU. For this particular question, one can do:

1. Define the computing inside FunctionLayer and wrap it with a trivial NetGraph:

net = Module[{w = 20}
 , NetGraph[<|
     "mat" -> RandomArrayLayer[NormalDistribution[0, 1/Sqrt[w]], "Output" -> {w, w}]
     , "f" -> FunctionLayer@Function[{mat, x}
                , {Normalize[Tanh[mat . x[[1]]]], Normalize[Tanh[mat . x[[2]]]]}
              ]
     , "dot" -> FunctionLayer@Function[x, x[[1]] . x[[2]]]
   |>
  , Flatten@{
      MapThread[#1 -> NetPort["f", #2] &, {{"mat", NetPort@"XPairs"}, {"mat", "x"}}]
      , "f" -> NetPort["newXPairs"], "f" -> "dot" -> NetPort["metric"]
    }, "XPairs" -> {2, w}]
 ]

2. Using net in batch mode, collecting result with iteration functions like NestList, FoldPairList, etc. Note net runs on CPU by default, uncommenting TargetDevice->"GPU",WorkingPrecision->"Real64" if one wants to run it on GPU.

res= Module[{batch = 10, depth = 200, w = 20}
  , FoldPairList[
       RightComposition[
          net[#1(*,TargetDevice->"GPU",WorkingPrecision->"Real64"*)] &
        , {#metric, #newXPairs} &
        ]
       , Association["XPairs" -> RandomVariate[NormalDistribution[], {batch, 2, w}]]
       , ConstantArray[1, depth]
   ]
  ];

3. Check the result:

res // Transpose // ListPlot

Answer 2

Your approach can be made a little faster (roughly 2x) by pre-calculating the random matrices, and using Fold over that list, rather than Nesting over each freshly generated entry. Another simplification is to use Normalize /@ a instead of your rowNormalize; however, this is used only once so it brings no significant speedup.

d = 300;  (*depth*)
w = 1024; (*layer dimension*)
b = 50;   (*batch size*)

x = RandomVariate[NormalDistribution[], {b, w}];
allMats = RandomVariate[NormalDistribution[0, Sqrt[2/(w + w)]], {d, w, w}];

a = Fold[
      Tanh[#1.#2] &,
      x,
      allMats
    ];

a = Normalize /@ a;
gram = a.Transpose[a];

Histogram[Statistics`Library`UpperTriangularMatrixToVector[gram]]

---

If it is acceptable to reduce the layer dimensions (e.g. to 512), then not only does the calculation become faster, but convergence can be obtained in fewer steps as well:

d = 1024; (*depth*)
w = 512;  (*layer dimension*)

(* rest of the code is same *)