new auto-dif design
Bob Carpenter
more traditional designs. The goal is to get rid
of the virtual function calls during derivative
propagation. And no, just using function pointers doesn't
help -- that's where Matt started. We need to get
rid of the whole concept.
The idea would be to instead store the expression graph
qua graph on a stack-like structure.
Each subexpression needs to store:
- value
- adjoint
- list of subexpressions
- list of partials with respect to subexpressions
There are many ways to do this.
I've been thinking of storing it pretty literally
as a graph, with nodes for expressions represented by three
parallel arrays:
double[] values;
double[] adjs;
int[] first_edge_idx; // index into array of edges
The edges then store the partials and the value of
the dtrs in another pair of parallel arrays:
int[] expr_index; // index into array of nodes
double[] partial;
For instance, expression 7 would have value
values[7], adjoint adjs[7], and dtr expressions
ranging from first_edge_idx[7] (inclusive) to
first_edge_idx[8] (exclusive) {this is a standard
trick for packing because we know the subexpressions
are contiguous, so we don't need both a start and
an end pointer, we can just have a start pointer and
use the next expression's start pointer as an exclusive
end pointer).
Then the partials need to get stored as a double value
and added to the stack.
Derivative propagation is now super easy, as it just
goes down the node stack, and for each expression,
iterates through its subexpression adding the expression's
adjoint times the partial w.r.t. the subexpression to the
subexpression's adjoint.
The equivalent of var now holds an index into the stack
of nodes, making it simple to compute value() and adj()
There is no longer a class for vari. That could
make var shorter if we can go with int (max out at 2G or
so subexpressions).
In fact, we could probably get rid of the values[] in
the node stack if we pushed the value into var itself.
It makes copying var heavier (not so good), but could save
lots of storage (very good).
- Bob
Ben Goodrich
There are many ways to do this.
Motion: Use an Eigen Array with a column for the partial and a column for the adjoint so the derivative is something like foo.rowwise.prod().sum().
Question: How would the proposed design generalize to Hessians? In the old design, I thought it was approximately just a matter of making the derivatives into vars and differentiating them again.
But, in general, I would say implement first (in a branch), ask questions later. If you do it and it turns out to be 1% faster, then I am basically sold.
Ben
Bob Carpenter
On 3/26/12 8:39 PM, Ben Goodrich wrote:
> On Monday, March 26, 2012 6:47:08 PM UTC-4, Bob Carpenter wrote:
>
> There are many ways to do this.
>
>
> Motion: Use an Eigen Array with a column for the partial and a column for the adjoint so the derivative is something
> like foo.rowwise.prod().sum().
That would certainly lead to clean derivative code.
Unfortunately, it would take up way more space.
The method I was suggesting is packed pretty tightly
into memory. Not as tightly as the current approach,
but close.
Check out the implementations of the Eigen objects.
I'm suggesting going a level lower representationally
than Eigen or std::vector.
> Question: How would the proposed design generalize to Hessians? In the old design, I thought it was approximately just a
> matter of making the derivatives into vars and differentiating them again.
Doing reverse-mode auto-dif twice works conceptually.
You essentially auto-dif the evaluation of the first-pass
dif to get the second derivatives.
But given that forward-mode auto-dif is so much
ligher memory wise, it's better to compute Hessians
row-wise and compute forward auto-dif for one dimension
of output and reverse-mode auto-dif that. Forward-mode
doesn't need any allocation and is much easier to template.
In fact, templated forward-mode is much easier than
reverse mode. I implemented the basics a while ago.
Matt's thinking about doing something a bit different, I think,
which would instead of computing a Hessian, would compute
a Hessian-gradient product. That I haven't looked into closely.
> But, in general, I would say implement first (in a branch), ask questions later. If you do it and it turns out to be 1%
> faster, then I am basically sold.
Definitely the way to go for any kind of major change.
I'm pretty sure what I'm proposing will be more than 1%
faster. It does the same arithmetic but without
any virtual function calls -- just a tight loop. It
does read in a bit more memory, which will slow it down.
Right now, we're being lazy in computing the partials
as we need them during derivative propagation. That
saves the space that the new proposal uses to store all
the partials in the forward evaluation. The benefit of
storing them all is that there's no thinking (virtual
function calls) in the reverse step.
It's partly a scalability vs. efficiency call. Though
we don't know exactly what either of those two will look
like in real use cases.
- Bob
Matt Hoffman
generalization of the scheme that you're suggesting for computing the
Hessians row-wise. The difference is that instead of projecting the
gradient onto an indicator vector you're projecting it into an
arbitrary vector. But the cost is essentially the same—in finite
differences terminology, you're just allowing a bigger space of
perturbations.
Matt
Marcus Brubaker
The basic idea of the new design is rather like
more traditional designs. The goal is to get rid
of the virtual function calls during derivative
propagation. And no, just using function pointers doesn't
help -- that's where Matt started. We need to get
rid of the whole concept.
Right, the performance hit with virtual functions comes from a function call that the compiler can't optimize because it's defined at run time. Virtual functions are basically implemented with function pointers, so that doesn't surprise me. There's also the added memory hit with virtual functions for the vtable pointer. I don't actually think your scheme below will increase memory usage that much overall since most expressions only have one or two subexpressions but you're saving at least one 64bit pointers worth of memory per expression. So maybe a 50% increase in memory.
For reference, I'm running a model right now which has ~40k variables and top reports it using around 1G of memory. I would consider that to be pretty reasonable and even doubling it for a large boost in performance would definitely be worth it.
It might be worth making the type of index a compile-time option since an int is likely plenty for small models, but could come up short for really big ones.
In fact, we could probably get rid of the values[] in
the node stack if we pushed the value into var itself.
It makes copying var heavier (not so good), but could save
lots of storage (very good).
I'm not sure how this saves memory, as you still need to store a double for every var. It's just been moved from the node stack to the var. Or did I miss something?
Cheers,
Marcus
Bob Carpenter
On 3/27/12 12:25 PM, Marcus Brubaker wrote:
> On Mon, Mar 26, 2012 at 6:47 PM, Bob Carpenter <ca...@alias-i.com
> <mailto:ca...@alias-i.com>> wrote:
>
> The basic idea of the new design is rather like
> more traditional designs. The goal is to get rid
> of the virtual function calls during derivative
> propagation. And no, just using function pointers doesn't
> help -- that's where Matt started. We need to get
> rid of the whole concept.
>
>
> Right, the performance hit with virtual functions comes from a function
> call that the compiler can't optimize because it's defined at run time.
> Virtual functions are basically implemented with function pointers, so
> that doesn't surprise me. There's also the added memory hit with
> virtual functions for the vtable pointer.
The vtable pointer in the case of vari, since it has
only one virtual function, is exactly the size of a function
pointer. It's just much more convenient to use.
> I don't actually think your
> scheme below will increase memory usage that much overall since most
> expressions only have one or two subexpressions but you're saving at
> least one 64bit pointers worth of memory per expression. So maybe a 50%
> increase in memory.
Not quite. Here's the memory rundown for an
expression with N subexpressions or arguments
under both schemes.
Current Scheme:
var
8 pointer to vari
-----
8
vari
8 vtable pointer to chain()
8 value
8 adjoint
8N pointers to all N subexpressions
--------
24 + 8N
New Scheme (indexed by long / int)
var
8 / 4 ptr to expression table
-------
8 / 4
expression (node) table
8 value
8 adjoint
8 / 4 pointer to first subexpression
partial (edge) table
8N value
8N / 4N pointers to var subexpressions table
----------
24 + 16N / 20 + 12N
Breaking that down by expression arity:
Now New(8) New(4)
0-ary: 24 24 20
1-ary: 32 40 32
2-ary: 40 56 44
3-ary: 48 72 56
N-ary: 24+8N 24+16N 20+12N
So it's more expensive in memory. Not enough
to really worry about for basic ops. For
instance, the normal distribution is 3-ary.
Three caveats.
1. The new scheme only needs to store subexpressions
that are vars. This could lead to lots of savings
in some cases.
2. Think of dot product of two K-dimensional matrices
as an example. Once we optimize vector expressions
by giving them auto-dif classes, this'll be one class
that stores pointers to the two vectors themselves.
This means the old scheme will only use 40 bytes.
The new scheme will store them all. Even if K=20,
we'll need over 300 bytes even if one of the vectors
is a double rather than a var.
3. The implicit end-of-subexpressions coding I'm
using means that memory needs to be allocated so
that the node arrays and edge arrays are just one big
array. That'll require either pre-allocating big,
or copying during resize. Our current scheme has slightly
better properties for this.
Thanks for making me reason through all of this!
>
> For reference, I'm running a model right now which has ~40k variables
> and top reports it using around 1G of memory. I would consider that to
> be pretty reasonable and even doubling it for a large boost in
> performance would definitely be worth it.
What kind of operations are you using? It sounds like its
needing lots of subexpressions.
Given how much space 40K variables are taking, I'm not thinking
we're going to hit 2G any time soon, even if chip density keeps
doubling every 18 months.
> In fact, we could probably get rid of the values[] in
> the node stack if we pushed the value into var itself.
> It makes copying var heavier (not so good), but could save
> lots of storage (very good).
>
>
> I'm not sure how this saves memory, as you still need to store a double
> for every var. It's just been moved from the node stack to the var. Or
> did I miss something?
The difference is that the node stack is kept in memory. It's
malloced and persists until the end of the gradient calcs. A var
is allocated on the function call stack, so will go away as soon
as we're done with the local block (or sooner if we use an optimizing
compiler that reuses its space).
- Bob
Marcus Brubaker
On 3/27/12 12:25 PM, Marcus Brubaker wrote:
On Mon, Mar 26, 2012 at 6:47 PM, Bob Carpenter <ca...@alias-i.com<mailto:ca...@alias-i.com>> wrote:
The basic idea of the new design is rather like
more traditional designs. The goal is to get rid
of the virtual function calls during derivative
propagation. And no, just using function pointers doesn't
help -- that's where Matt started. We need to get
rid of the whole concept.
Right, the performance hit with virtual functions comes from a function
call that the compiler can't optimize because it's defined at run time.
Virtual functions are basically implemented with function pointers, so
that doesn't surprise me. There's also the added memory hit with
virtual functions for the vtable pointer.
The vtable pointer in the case of vari, since it has
only one virtual function, is exactly the size of a function
pointer. It's just much more convenient to use.
Agreed, the function pointers are the sad, lonely C step-child of virtual functions but they're effectively equivalent :)
It's likely to be a very small difference in practice I expect as the vast majority of operations are 0, 1 or 2 ary. Plus, for the few 3+ ary functions (e.g., normal_log, matrix operations, etc) the savings over expressing them in terms of basic operations seem likely to hugely outweigh the small added cost of expressing higher ary functions.
Three caveats.
1. The new scheme only needs to store subexpressions
that are vars. This could lead to lots of savings
in some cases.
2. Think of dot product of two K-dimensional matrices
as an example. Once we optimize vector expressions
by giving them auto-dif classes, this'll be one class
that stores pointers to the two vectors themselves.
This means the old scheme will only use 40 bytes.
The new scheme will store them all. Even if K=20,
we'll need over 300 bytes even if one of the vectors
is a double rather than a var.
3. The implicit end-of-subexpressions coding I'm
using means that memory needs to be allocated so
that the node arrays and edge arrays are just one big
array. That'll require either pre-allocating big,
or copying during resize. Our current scheme has slightly
better properties for this.
You could use a deque to avoid the large contiguous memory requirement. Other than allocation failure for large blocks, I wouldn't worry about it too much though so long as the memory didn't need to be reallocated for each gradient calculation. My profiling has shown that memory allocation is an almost negligible cost.
Thanks for making me reason through all of this!
Hah, glad for my ignorance to be of service!
What kind of operations are you using? It sounds like its
For reference, I'm running a model right now which has ~40k variables
and top reports it using around 1G of memory. I would consider that to
be pretty reasonable and even doubling it for a large boost in
performance would definitely be worth it.
needing lots of subexpressions.
There are a number of cholesky factorizations, matrix solves (in multi_normal_log) and matrix multiplies.
The difference is that the node stack is kept in memory. It's
In fact, we could probably get rid of the values[] in
the node stack if we pushed the value into var itself.
It makes copying var heavier (not so good), but could save
lots of storage (very good).
I'm not sure how this saves memory, as you still need to store a double
for every var. It's just been moved from the node stack to the var. Or
did I miss something?
malloced and persists until the end of the gradient calcs. A var
is allocated on the function call stack, so will go away as soon
as we're done with the local block (or sooner if we use an optimizing
compiler that reuses its space).
Ah, I see. It's a tough call which would be better.
Cheers,
Marcus
Bob Carpenter
On 3/27/12 3:09 PM, Marcus Brubaker wrote: