t-SNE in plain javascript
Goal Walk through each of the steps needed to perform t-SNE. We will execute each step explicitely in plain javascript, no libraries allowed. At the end there will be a bunch of sliders that allow you to play with the algorithm and watch it attempt to visualize some simulated high-dimensional data. Calculate pairwise distances. The first step of t-SNE is to calculate the pairwise distances between all of the points in your data. Luckily we only need to do this a single time. This is because t-sne only cares about these distances and they don't change ever change. This only-once property is great in that after you have done the distance calculations, the actual data doesn't matter. So really the runtime of the algorithm is bounded not by the dimensionality of your input data but by how many samples you have. The nieve way of doing this pairwise calculation would be a simple outer loop scanning over each observation in turn and then within that outer loop an inner loop that calculates the distance between it and every other point. This works but is wasteful. Let's think about these distances. The distance between point i and point j is the same as the distance between j and point i so really we should only calculate a single one, in addition we don't really need to check between i and i as we know it's zero. To see how much time this can save us we can use the formula for calculating all permutations(order matters) of two numbers and all combinations(order doesn't matter) of two numbers. With total datapoints or observations, we can calculate the total number of permutations:
viewof n = slider2('Set number of observations in your dataset', 'n', 0, 100, 50)
tex` \begin{aligned} n^2 = ${n}^2 = \boxed{${squared(n)} \text{ pairs}} \end{aligned}`
Now if we're efficient about it and only look at the combinations of the data we can cut this down. The combination equation is a bit more complicated...
tex` \begin{aligned} \frac{n!}{2!(n - 2)!} =& \frac{${n}!}{2\cdot${n - 2}!} \\ & \frac{ ... ${n - 3} \cdot ${n - 2} \cdot ${n - 1} \cdot ${n}}{2\cdot(...${n - 3} \cdot ${n - 2})} \\ & \frac{${n - 1} \cdot ${n}}{2} \\ & \boxed{${n*(n - 1)/2}\text{ pairs}} \end{aligned}`
So we can see that by taking advantage of this combination vs permutation distinction we are saving ourselves __ __ total distance calculations. Kinda nice, especially when the data gets larger. Now that we've gotten some superfluous math out of the way let's actually implement this in code. A couple things to note: The function calc_distance simply takes two arrays of some dimension and calculates their euclidean distance(sometimes called the L2 norm) from eachother. The actual code to implement it is at the bottom of this notebook if you're curious. We are filling an array that has every simple permutation of the data in it, but if you note the second for loop's bounds we are only calculating the combinations. We simple fill in two spots in our distances array for each distance. This is because we later need the array to be the size of all the distances due to the quirk of non-symmetry in the conditional probability step. (We'll get to it, don't worry.)
The first thing we are going to do is make a helper function that will take an i and j coordinate and point us to the corresponding position in a javascript array. This will make our lives a lot easier as it can be tricky to think about mapping two dimensional coordinates to a 1-d array
function mat_ind(i,j){ //Get index in a 1-d array from a 2d matrix reference point. (Assumes our matrix is square) //This function will auto update when we change the number of points so we don't need to pass //that as a parameter thanks to the beauty of reactive programing! return i*num_points + j }
Okay, now we can get into the pairwise distances calculation itself.
pairwise_distances = { // initialize some variables we use inside our loops let x_i, x_j, distance; // make_array is another helper function that just generates a new empty array of a given size. const pw_dists = make_array(squared(num_points)); for(let i = 0; i < num_points; i++){ x_i = data[i].location; // deal with the point itself here. pw_dists[mat_ind(i,i)] = {source: i, target: i, distance:0}; for(let j = i + 1; j < num_points; j++){ x_j = data[j].location; distance = calc_distance(x_i, x_j); pw_dists[mat_ind(i,j)] = {source: i, target: j, distance}; pw_dists[mat_ind(j,i)] = {source: j, target: i, distance}; } } return pw_dists }
If you expand the returned array from above you will see we have a series of objects containing source: the index in our original data corresponding to the first point, target: the index for our second point, and distance: the euclidean distance between those two points. Pairwise probability calculations Now that we have the distances we can get to the fun stuff. We need to calculate the conditional probabilities between points in their original high dimensional space. Going from distance to probability This takes the probability seeing point j given a multi-dimensional gaussian distribution centered around point i. This is a bit to wrap your head around so let's look at it in two dimensions. In the following plot we are showing the calculation of the conditional probability of point i given a gaussian centered around point j. Play around with the position of point i and the standard deviation of the distribution to see the relationship between euclidean distances between points and their conditional probabilies.
{ const pdf = linspace(-5, 5, 250).map(x => ({x, px:stats.dnorm(x, 0, orig_sd)})) const svg = d3.select(DOM.svg(width, height)); const x = d3.scaleLinear() .domain(d3.extent(pdf, d => d.x)) .range([margin.left, width - margin.right]); const y = d3.scaleLinear() .domain(d3.extent(pdf, d => d.px)) .range([height - margin.bottom, margin.top]); const line = d3.line() .x(d => x(d.x)) .y(d => y(d.px)) const p_ij = stats.dnorm(new_point, 0, orig_sd) // line showing point up to pdf svg.append('circle') .attr('cx', x(0)) .attr('cy', y(0)) .attr('r', 5) .attr('fill', 'steelblue'); svg.append('circle') .attr('cx', x(new_point)) .attr('cy', y(0)) .attr('r', 5) .attr('fill', 'orangered'); svg.append("line") .attr("fill", "none") .attr("stroke", "orangered") .attr("stroke-width", 1.5) .attr("stroke-linejoin", "round") .attr("stroke-linecap", "round") .attr("x1", x(new_point)) .attr("x2", x(new_point)) .attr("y1", y(p_ij)) .attr("y2", y(0)) svg.append("text") .html( `P(<tspan fill = 'orangered'>i</tspan>|<tspan fill = 'steelblue'>j</tspan>) = ${p_ij.toPrecision(3)}` ) .attr('font-size', 20) .attr("x", x(new_point) + 3) .attr("y", y(p_ij) + 10) svg.append("path") .datum(pdf) .attr("fill", "none") .attr("stroke", "darkgrey") .attr("stroke-width", 1.5) .attr("stroke-linejoin", "round") .attr("stroke-linecap", "round") .attr("d", line); return svg.node(); }
viewof new_point = slider2('Set distance of point j from point i', '', -5, 5, 0.333, 3)
viewof orig_sd = slider2('Set the standard deviation of the distribution', '', 0.1, 3, 1, 3)
Perplexity You may notice that in the above example we can tune the standard deviation, or width, of the distribution that we center over a point to calculate it's respective conditional probabilities. It's tempting to say that in the final algorithm we just set this value as a constant and sit the same distribution over every point, unfortunately that doesn't work very well. The reason is that we may have areas in the space that are very dense with points and others that are very spare. Since we car about structure at multiple scales treating data in both of these regions the same will not give us good results. The way that t-SNE deals with this is by scanning over possible standard deviations for every point's distribution trying to fit to a parameter called 'perplexity'. Perplexity a measure related to the entropy of the system of points. A valuable way of thinking about perplexity is to think of it as the effective number of neighbors for each point. Aka how many points that the distribution over point i capture with a non-trivial probability. Below we plot a point in our data and each other point's distance away from that point. The gaussian probability curve is drawn over these points based upon the current value of beta for this given point. As you change beta with the slider you notice how the distribution gets wider or smaller to essentially 'grab' more or less points that contribute a meaningful amount of probability to it's overall sum of probabilities.
{ const dist_to_prob = calc_gaussian_probs(anchor_distances.map(d => d.distance), beta) const svg = d3.select(DOM.svg(width, height)); const distance_max = d3.max(anchor_distances, d => d.distance); const entropy = calc_entropy(dist_to_prob.map(d => d.prob)); const x = d3.scaleLinear() .domain([0, distance_max]) .range([margin.left, width - margin.right]); // calculate the probability curve given the current beta value const prob_dist_curve = linspace(0, distance_max, 300) .map(dist => ({dist, prob:gaussian_prob(dist, beta)})); const y = d3.scaleLinear() .domain(d3.extent(prob_dist_curve, d => d.prob)) .range([height - margin.bottom, margin.top]); const point_opacity = d3.scalePow() .domain(d3.extent(dist_to_prob, d => d.prob)) .range([0,1]) const area = d3.area() .x(d => x(d.dist)) .y0(y(0)) .y1(d => y(d.prob)) svg.append("text") .text(`Entropy: ${entropy.toPrecision(3)}`) .attr('x', width - 10) .attr('y', 40) .attr('text-anchor', 'end') .attr('font-size', 20) svg.append("text") .text(`Perplexity: ${exp(entropy).toPrecision(3)}`) .attr('x', width - 10) .attr('y', 70) .attr('text-anchor', 'end') .attr('font-size', 20) svg.append("text") .text(`Perplexity Goal: ${perplexity.toPrecision(3)}`) .attr('x', width - 10) .attr('y', 100) .attr('text-anchor', 'end') .attr('font-size', 20) svg.append("path") .datum(prob_dist_curve) .attr("fill", "steelblue") .attr("fill-opacity", 0.2) .attr("d", area); svg.append("g") .selectAll("circle") .data(dist_to_prob) .enter().append("circle") .attr("cx", d => x(d.dist)) .attr("cy", y(0)) .attr("stroke", "grey") .attr("stroke-opacity", 0.2) .attr('fill', d => d.dist === 0 ? 'orangered': 'steelblue') .attr("fill-opacity", d => d.dist === 0 ? 1: point_opacity(d.prob)) .attr("r", 5); return svg.node(); }
viewof beta = slider2('Set beta value for the distribution of the point above', '', 0.16, 10, 1, 8)
anchor_distances = { const anchor_point = 5; const anchor_distances = pairwise_distances .filter(d => d.source === anchor_point); return anchor_distances }
In order to find the optimal beta for each point in our dataset we use a searching technique called binary search. What binary search is doing here is goes point by point and searching in a disciplined way for the beta value that gets the perplexity value as close to the supplied perplexity as possible. Just like you did with the slider above. To do this we first need to find our goal entropy for our points. Do do this we simply take the log of our goal perplexity.
target_entropy = log(perplexity) // entropy level we're shooting for given perplexity
Now that we have the entropy we just need a couple of functions to make our lives easier. The first is one that will simply perform the gaussian probability calculation of a point j given its distance from point i and point i's beta value...
tex` p_{j|i} = \exp\{-(\text{distance from } i \text{ to } j) \cdot \beta_i\}. `
function gaussian_prob(distance, beta){ return exp(-distance * beta) }
Next is a function that takes an array of distances from point i and converts them to an array containing both the distance and the probability according to the guassian model.
function calc_gaussian_probs(distances, beta){ return distances .map(dist => ({dist, prob: (dist !== 0 ? gaussian_prob(dist, beta) : 0)})); }
Last, we need a function that can take our array of probabilities and calculate the entropy for point i given its current beta value. To do this we need to normalize the probabilities and then sum them using the formula for entropy...
tex` (\text{entropy for point } i) = \sum_{j \ne i} -p_{j|i}\log(p_{j|i}) `
function calc_entropy(probs){ const sum_probs = probs.reduce((sum, d) => sum + d, 0); return probs.reduce( (sum, prob) => { const prob_normed = prob/sum_probs; return sum - (prob_normed > 1e-7 ? prob_normed*log(prob_normed): 0) }, 0 ) }
Lastly, we set up a function for doing the binary search for the optimal beta given a point. I won't go too far into this other than to leave the code here. What matters is that the algorithm bounces around the beta value in a disciplined way in an effort to make the entropy match our goal.
function beta_binary_search(i, pw_dists, entropy_tol = 1e-3, max_attempts = 40){ // grab the distances for point i const dists_from_i = pw_dists .reduce( (dists, d) => d.source === i ? [...dists, d.distance] : dists, [] ); let beta_min = 0, // beta lower bound beta_max = Infinity, // beta upper bound beta = 1, // initial beta value searching_for_beta = true, attempts = 0, probs_given_beta, current_entropy; // number of iterations through while looking while(searching_for_beta){ attempts++ // increment out attempts counter // calculate the probabilities and entropy given the current beta. probs_given_beta = calc_gaussian_probs(dists_from_i, beta).map(d => d.prob); current_entropy = calc_entropy(probs_given_beta); // check where our entropy lies in context to the target if(current_entropy > target_entropy){ // distribution is wide. beta_min = beta; // move our min bound up beta = beta_max === Infinity ? beta*2 : (beta+beta_max)/2; } else { // distribution is too narrow. beta_max = beta; // move our max bound down beta = beta_min === -Infinity ? beta/2 : (beta+beta_min)/2; } // check some termination conditions const too_many_attempts = attempts > max_attempts; const close_to_target = Math.abs(current_entropy - target_entropy) < entropy_tol; searching_for_beta = !(too_many_attempts || close_to_target); } return probs_given_beta }
Now we calculate the non-symmetric probabilities. Aka probability of i given j.
P_i_given_j = { const P_i_given_j = make_array(num_points*num_points); for(let i = 0; i < num_points; i++){ // run binary search to find the best beta. let p_row = beta_binary_search(i, pairwise_distances) for(let j=0;j<num_points;j++) { P_i_given_j[i*num_points+j] = p_row[j]; } } return P_i_given_j }
Now that we have the non-symmetric probabilities we need to make them symmetric by averaging both of the conditional probabilities p(i|j) and p(j|i).
P = { const P = make_array(num_points*num_points); for(let i=0; i<num_points; i++) { for(let j=0; j<num_points; j++) { P[i*num_points+j] = Math.max( (P_i_given_j[i*num_points+j] + P_i_given_j[j*num_points+i])/(num_points*2), 1e-10 // substitute 0s with very small number to improve stability due to rounding errors. ); } } return P }
The mapping calculations We have now done all we need to do with the high dimensional data. Since at each iteration the high dimensional probabilities don't move we only have to do everything above once. From now on we are simply going to be taking a random initial mapping of all our datapoints into two(or three) dimensions and nudging those mappings around until they best match the distribution of probabilities we generated in the above sections. Initializing random mapping positions First we will initialize a random starting position for the mappings. We will do this purely randomly with a uniform random number generator, but it would be interesting to compare the performance of initializations using different distributions or patterns (such as the phylotaxisis initialization patterns commonly used in network diagrams.
Next, we make a function that calculates the conditional probabilities in the low-dimensional mapping.
function find_Q(Y_curr){ // takes the a given Y mapping and produces the unnormalized q values for it. const Q = make_array(num_points*num_points); let sum_q = 0, q_ij, // given combination's q value dim_sum; // dimension sum for(let i=0; i<num_points; i++){ for(let j=0; j<num_points; j++){ //scan over the mapping dimensions of i and j and sum their squared differences dim_sum = make_array(map_dims) .reduce((sum, _, dim) => squared(Y_curr[i][dim] - Y_curr[j][dim])); q_ij = 1/(1 + dim_sum); // set the q value for this combination Q[i*num_points+j] = q_ij; // assign it to the big Q array Q[j*num_points+i] = q_ij; sum_q += 2*q_ij; // accumulate for later normalization } } // normalize Q values return Q.map(q => max(q/sum_q, 1e-99)); }
function calc_cost_grad(Y, iters){ let cost_acc = 0, pre_mult, // this is poorly named; grad_sum, early_exag = iters < early_exag_itts? early_exag_amount: 1; const Q = find_Q(Y), grad_acc = make_array(num_points); for(let i=0; i<num_points; i++){ grad_sum = make_array(map_dims).map(()=>0); // initialize array of zeros for gradient accumulation for(let j=0; j<num_points; j++){ cost_acc += -P[i*num_points+j] * log(Q[i*num_points+j]); // this may be a mistake pre_mult = 4*(early_exag*P[i*num_points+j] - Q[i*num_points+j])*Q[i*num_points+j] grad_sum = grad_sum.map((grad, dim) => grad + (pre_mult * (Y[i][dim] - Y[j][dim]))) } grad_acc[i] = grad_sum; } return {cost: cost_acc, grad: grad_acc} }
calc_cost_grad(initialize_points({n: num_points, dim: map_dims}))
Now that we have the cost and gradient at this step we can recalculate our y mapping by moving in the direction of the gradient until we lower our cost as much as possible.
Y_costs = { let iters = 0 const costs = []; // initial random y positions along with a vector to store the last y gradient step and the gains. const Y_new = initialize_points({n: num_points, dim: map_dims}), Y_step = initialize_points({n:num_points, dim: map_dims, fill:0}), Y_gains = initialize_points({n: num_points, dim: map_dims, fill: 1}); let average_y = make_array(map_dims); // calculate the cost and gradient for current y locations while(iters < num_itterations){ iters++ average_y = average_y.map(() => 0) // zero out y average for new iteration let {cost, grad} = calc_cost_grad(Y_new, iters), momentum = momentums[iters > momentum_switch ? 1: 0], grad_id, step_id, gain_id, gain_id_new, step_id_new; costs.push(cost); for(let i=0; i<num_points; i++){ for(let d=0; d<map_dims; d++){ // pull out the gradient, last step and gain for point i dimension d grad_id = grad[i][d]; step_id = Y_step[i][d]; gain_id = Y_gains[i][d]; // compute the gain component and clamp above some constant gain_id_new = max( sign(gain_id) === sign(step_id) ? gain_id * gain_scales_same: gain_id + gain_scales_diff, gain_clamp ) Y_gains[i][d] = gain_id_new; // compute the update amount step_id_new = momentum*step_id - learning_rate*gain_id_new*grad_id; Y_step[i][d] = step_id_new // update the y value Y_new[i][d] += step_id_new; // accumulate the mean so we can center average_y[d] += Y_new[i][d]; } } // center Y for(let i=0; i<num_points; i++){ for(let d=0; d<map_dims; d++){ Y_new[i][d] -= average_y[d]/num_points } } //yield Promises.delay(100,Y_new) if(iters % 5) yield {Y:Y_new, costs} } }
Y = Y_costs.Y
costs = Y_costs.costs
{ const svg = d3.select(DOM.svg(width, height)); const dim1_domain = d3.extent(Y, (y) => y[0]); const dim2_domain = d3.extent(Y, (y) => y[1]); const dim1 = d3.scaleLinear().domain(dim1_domain).range([margin.left, width - margin.right]); const dim2 = d3.scaleLinear().domain(dim2_domain).range([margin.top, height - margin.bottom]); const group_color = d3.scaleOrdinal() .domain(make_array(num_clusters).map((d,i)=>i)) .range(d3_color.schemeCategory10); const vis_data = data.map((d,i) => ({...d, map: Y[i]})); svg.append("g") .attr("stroke", "#000") .attr("stroke-opacity", 0.2) .selectAll("circle") .data(vis_data) .enter().append("circle") .attr("cx", d => dim1(d.map[0])) .attr("cy", d => dim2(d.map[1])) .attr("fill", d => group_color(d.group)) .attr("r", 5); return svg.node(); }
{ const svg = d3.select(DOM.svg(width, height)); const x = d3.scaleLinear().domain([0, costs.length]).range([margin.left, width - margin.right]); const y = d3.scaleLinear().domain(d3.extent(costs)).range([height - margin.bottom, margin.top]); svg.append("g") .attr("stroke", "#000") .attr("stroke-opacity", 0.2) .selectAll("circle") .data(costs) .enter().append("circle") .attr("cx", (d,i) => x(i)) .attr("cy", y) .attr("fill", 'orangered') .attr("r", 2); return svg.node(); }
viewof num_itterations = slider({min: 1, max: 10000, value: 1000})
viewof perplexity = slider({min: 1, max: 70, value: 30})
viewof learning_rate = slider({min: 0, max: 20, value: 10})
viewof early_exag_amount = slider({min: 1, max: 10, value: 4})
viewof early_exag_itts = slider({min: 0, max: num_itterations, value: 250})
viewof gain_scales_same = slider({min: 0, max: 1, value: 0.8})
viewof gain_scales_diff = slider({min: 0, max: 1, value: 0.2})
viewof gain_clamp = slider({min: 0, max: 0.1, value: 0.01})
viewof momentum_switch = slider({min: 0, max: 1000, value: 250})
momentums = [0.5, 0.8]
map_dims = 2
Logic for generating data is here, It's not interesting so I haven't pinned it.
num_dimensions = 10
num_clusters = 4
num_samples = 25
num_points = data.length
cluster_spread = 15
viewof random_noise_sd = html`<input type=range min="0" max="10" value="2">`
random_noise_sd
data = { const make_centroid = (ndim) => { const low = -cluster_spread/2; const high = cluster_spread/2; const point_gen = d3.randomUniform(low, high); return () => [...(new Array(ndim))].map(point_gen); } const gen_centroids = (n, ndim) => { const center_maker = make_centroid(ndim); return [...(new Array(n))].map(center_maker) } const centroids = gen_centroids(num_clusters, num_dimensions) const sample_from_centroid = (nsamples, centroid, centroid_id) => { const rnorm = d3.randomNormal(0, random_noise_sd); const get_sample = () => ({group:centroid_id, location : centroid.map(d => d + rnorm())}); return [...(new Array(nsamples))].map(get_sample); } function gen_data_from_centroids(centroids, nsamples){ return centroids.reduce( (data, d, i) => [...data, ...sample_from_centroid(nsamples, d, i)], []) .map((d,i) => (Object.assign({id: i},d,{}))) // add in node id in addition to the centroid one. } return gen_data_from_centroids(centroids, num_samples) }
function calc_distance(x1, x2){ // calculates the euclidean distance between two vectors. Aka the L2 Norm const sum_of_square_diffs = x1.map( (x1_i, i) => Math.pow(x1_i - x2[i], 2) // get squared difference each dimension ) .reduce((sum, d) => sum + d, 0); // sum the squared differences to return a single number return Math.sqrt(sum_of_square_diffs) }
exp = Math.exp
squared = (val) => Math.pow(val, 2)
max = Math.max
log = Math.log
sign = (x) => x > 0 ? 1: x < 0 ? -1: 0
function initialize_points({n, dim = 2, fill}){ const fill_function = typeof fill == 'undefined' ? d3.randomUniform(-1,1) : () => fill; const rand_point = () => make_array(dim).map(fill_function); return make_array(n).map(rand_point) }
make_array = (n) => [...(new Array(n))]
function linspace(low, high, n = 100){ const width = (high - low), half = width/2; return make_array(100).map((d,i) => low + ((i/100)*width)) }
linspace(0,10, 25)
make_array(10).length
(new Array(10)).length
margin = ({top: 20, right: 30, bottom: 30, left: 40})
height = 300
d3_color = require('d3-scale-chromatic')
function slider2(msg, var_name, min, max, start, precision = 1){ const diff_prec = precision !== 1 if(diff_prec){ min = min*precision; max = max*precision; start = start*precision; } const div = html` ${msg} <input type=range min=${min} max=${max} value = ${start}> ${var_name} = <span>${start/precision}</span> `; const range = div.querySelector("[type=range]"); const value_disp = div.querySelector('span'); div.value = value_disp.innerHTML = diff_prec ? start/precision: start; range.addEventListener("input", () => { const newValue = diff_prec ? range.valueAsNumber/precision: range.valueAsNumber; div.value = newValue; value_disp.innerHTML = diff_prec ? newValue.toPrecision(precision): newValue; }); return div; }
{ let diff_prec = true, min = 0, max = 1, start = 0.8, step = 0.1; if(diff_prec){ min = min/step; max = max/step; start = start/step; } const div = html` ${'hithere'} <input type=range min=${min} max=${max} value = ${start}> ${'hi'} = <span>${start*step}</span> `; const range = div.querySelector("[type=range]"); const value_disp = div.querySelector('span'); div.value = value_disp.innerHTML = diff_prec ? start*step: start; range.addEventListener("input", () => div.value = value_disp.innerHTML = diff_prec ? range.valueAsNumber*step: range.valueAsNumber); return div; }
import {slider} from "@jashkenas/inputs"
d3 = require('d3')
stats = require('statdists')