Flashlight: Enabling Innovation in Tools for Machine Learning

1 INTRODUCTION . The recent rise of deep learning-based techniques has been accompanied and sustained by the wide availability of dedicated frameworks such as TensorFlow ( Abadi et al. , 2016 ) and PyTorch ( Paszke et al. , 2019 ) . These frameworks have enabled the democratization of machine learning research by providing extensive collections of high level primitives to support common use cases . Lowering the barrier to entry for end users has boosted the popularity of both neural networks and the frameworks in which they are implemented . However , in order to support what are now vast ecosystems and a diverse user base , framework size and complexity have increased dramatically over time . As a result , deep , groundbreaking framework research has become extremely onerous and time consuming , precluding rapid innovation . Given these barriers , major deep learning frameworks have become stuck in their existing operating modes . Innovation in this area remains as important as ever . Indeed , framework innovation accelerates machine learning ( ML ) and artificial intelligence ( AI ) research . Frameworks that are easier to use reduce the engineering burden on researchers , and frameworks that are higher-performance decrease the time required to iterate on experimental work and validate hypotheses . Even more critically , tooling plays a fundamental role in deciding which ideas succeed or fail . For example , LeCun et al . ( 1989 ) pioneered the use of convolutional neural networks ( CNNs ) ( Fukushima & Miyake , 1982 ) trained using backpropagation for computer vision tasks in the late 1980s , which was subsequently applied to handwriting recognition . However , widespread success for CNNs was achieved two decades later when Krizhevsky et al . ( 2012a ) leveraged the CUDA programming model to take advantage of graphics processing units ( GPUs ) to train a much deeper model ( AlexNet ) . While deep learning frameworks have been optimized to leverage existing hardware paradigms for common neural network architectures , they often fail to deliver similar efficiencies on designs that diverge from the mainstream . For example , Barham & Isard ( 2019 ) explain how the design of these frameworks results in poor hardware utilization for a novel type of neural network , known as a capsule network ( Hinton et al. , 2018 ) , that leverages new components such as squashing operations and routing by agreement . More generally , what are now unconventional approaches to modern problems in machine learning require highly-specialized additions to popular frameworks . As a result of narrowly-optimized systems , research beyond deep learning may be discounted due to purported computational infeasibility given modern frameworks ’ capabilities . Furthermore , the waning of Moore ’ s law ( Theis & Wong , 2017 ) coupled with the ever-growing computational demands of deep learning are prompting several shifts in hardware . Massive-scale distributed computing is now required to train leading models — a process that established frameworks remain unable to handle truly automatically . In parallel , multiple specialized hardware products are now available to better support deep learning applications : Nvidia ’ s TensorCores ( Markidis et al. , 2018 ) , Google ’ s TPUs ( Jouppi et al. , 2017 ) , Graphcore ’ s IPUs ( Jia et al. , 2019 ) , Apple ’ s Neural Engine1 , and others have been developed to improve total float-pointing operations ( FLOPs ) , cost per-FLOP , or energy consumption . Additionally , numerous efforts are underway to move away from conventional von Neumann computing architectures in which memory and processing units are physically separated , either by storing data closer to compute units or by switching to in-memory computing altogether . While tooling innovation is alive and well given these incentives for progress , working within large , well-established frameworks has become more and more challenging as framework size and scope grows . As a result , many recent innovations have required the development of ad-hoc tools . For example , efforts in machine learning-driven compilation of neural networks are largely built on top of Halide ( Adams et al. , 2019 ; Steiner et al. , 2021 ) and TVM ( Chen et al. , 2018 ; Zheng et al. , 2020 ) ; FlexFlow ( Jia et al. , 2018 ; 2020 ) underpins recent work aimed at improving the use of distributed computing to accelerate the training of large neural networks ; and PET ( Wang et al. , 2021 ) provides a framework that enables graph-level neural network optimizations . With ad-hoc approaches , researchers are required to start from scratch for new directions or adapt their ideas to fit into the scaffolding these frameworks provide — resulting in significant technical burdens . To sustain framework innovation , we introduce Flashlight , an open source minimalist ML library designed to support research in machine learning frameworks , facilitate rapid iteration on ideas , reduce the engineering burden on researchers , and remove the need for new tools . Flashlight includes : • A modular , component-based architecture that makes every aspect of the implementation fully customizable with simple internal APIs . • A compact yet highly-performant reference implementation of each component . • A comprehensive set of benchmarks representative of the state-of-the-art in machine learning on which to evaluate alternative implementations . 2 RELATED WORK . Numerous frameworks have been implemented in recent years to support machine learning , including Lush ( Bottou & LeCun , 2002 ) , Theano ( Bergstra et al. , 2010 ) , Torch ( Collobert et al. , 2011 ) , Caffe ( Jia et al. , 2014 ) , MXNet ( Chen et al. , 2015 ) , deeplearning4j ( Team , 2016 ) , TensorFlow ( Abadi et al. , 2016 ) , Flux ( Innes , 2018 ) , Jax ( Bradbury et al. , 2018 ) , PyTorch ( Paszke et al. , 2019 ) , Chainer ( Tokui et al. , 2019 ) , and PaddlePaddle ( Ma et al. , 2019 ) . These frameworks offer programming models designed around multidimensional arrays ( TENSORS ) , modeled as first-class objects and supported by a comprehensive set of mathematical primitives ( or operators ) to manipulate them . To provide the computing power required by deep learning-based methods , most natively support hardware accelerators such as general-purpose GPUs or custom-designed ASICs such as TPUs . Generally , framework implementations follow one of two computational models : • In the deferred execution model , the neural network to be trained is first encoded as a dataflow graph which can be optimized for a specific set of target hardware devices . The neural network is then executed in a distinct second phase . Since the dataflow graph represents the entire computation , both local and global optimizations can be applied , making the subsequent execution very efficient . However , only programs that can be represented as 1https : //nr.apple.com/dE9q1p9M7t dataflow graphs can be processed with this approach , thus limiting flexibility . Frameworks such as Theano , TensorFlow2 , Caffe , or MXNet fall into this category . • In the eager model , an interpreter ( such as Python ) is extended with the high level kernelbased operations needed to train a neural network . These operations are executed immediately when called , though this precludes many optimizations . By weaving neural network-related operations into a Turing complete programming language , this approach is extremely flexible . Furthermore , the imperative nature of the underlying programming language allows for fine-grained control over the execution order and memory utilization , which enables more specific user-driven optimization . Frameworks such as Torch , PyTorch , or Chainer exemplify this approach . 3 PRINCIPLES . The aforementioned frameworks are designed and implemented to best-serve their user bases — namely , machine learning researchers and practitioners . They rely on large , internally complex codebases to provide comprehensive solutions , as is further discussed in Section 5 . In contrast , Flashlight targets an audience of researchers interested in experimenting with new designs and implementations of machine learning tools or broader computational or modeling paradigms . To foster this type of innovation , Flashlight balances simplicity and nimbleness with the need to provide enough functionality to support real use cases . Internal and external simplicity is the key design principle of Flashlight ; the ability to dramatically modify software and drive it in new directions is inversely correlated with codebase size and complexity ( Gill & Kemerer , 1990 ) . More specifically : • Flashlight is built on a shallow stack of idiomatic , modular , and customizable abstractions . Framework components interact through small , well-defined , stable APIs , which expose most internal aspects of its implementation . This ensures that every component of Flashlight can be modified or replaced with new custom implementations , even e.g . its memory manager and tensor implementation . To support the exploration of a wide array of alternative approaches , Flashlight interfaces are flexible and unopiniated by design . This is in contrast to other frameworks , which impose stricter implementation requirements based on tight design constraints for their computation models and support requirements across hardware , downstream frameworks , and other ecosystem members . • Flashlight provides deliberately-compact default implementations of its APIs . This reduces out-of-the-gate engineering burden and the need for modifications , and enables fast compilation and rapid iteration when experimenting . Furthermore , to mitigate premature optimization , Flashlight deliberately abstains from adding small efficiency improvements if they conflict with the goals of keeping the codebase simple and APIs clean . • Flashlight is a research-first framework , and is not intended for out of the box production use . To keep codebase size small , it forgoes features such as model servers for deployment and integration with cluster management tools . Flashlight is a viable solution for machine learning research , shipping with a comprehensive set of benchmarks and research setups for state-of-the-art neural network architectures such as convolutional neural networks ( CNNs ) ( Krizhevsky et al. , 2012b ) and Transformers ( Vaswani et al. , 2017 ) , as well as task-specific models such as ViT ( Dosovitskiy et al. , 2020 ) , DETR ( Carion et al. , 2020 ) , or BERT ( Devlin et al. , 2018 ) . The speech recognition system wav2letter ( Pratap et al. , 2019 ) , is also built entirely on Flashlight . Benchmarks built on these state-of-the-art models make Flashlight a turn key solution for system researchers who want to quickly evaluate their design and implementation choices without needing to build test benches from the ground-up . More importantly , Flashlight makes possible end-to-end benchmarking on real models rather than microbenchmarks or small-scale tests . 2TensorFlow 2.0 adds support for eager execution semantics as well . 4 DESIGN . Flashlight ’ s design is centered around internal APIs for framework components which form the building blocks for domain-specific ML packages and applications — this structure is outlined in Figure 1 . Flashlight is implemented as a C++ library and follows a Tensor-based programming methodology , with neural network building blocks that derive from a MODULE interface , communicate by exchanging Tensor data , and are composed functionally or imperatively to form complete neural network architectures . Tensor programming in Flashlight is fundamentally dynamic , but given that C++ is a compiled language , code describing models in Flashlight is compiled . This approach promotes type safety , foregoes the runtime overheads associated with interpreters , and , unlike eager-based approaches , enables global optimizations where possible . 4.1 OPEN FOUNDATIONAL INTERFACES . Flashlight is built on top of three open foundational APIs , each addressing design and implementation challenges faced by machine and deep learning tools : a Tensor interface , a memory management subsystem , and a distributed computing interface . These APIs are backed by reference implementations that enable Flashlight to efficiently target CPUs , GPUs , and other accelerators . These include code generation and dedicated kernels for Intel , AMD , OpenCL , and CUDA devices , and leverage libraries such as cuDNN ( Chetlur et al. , 2014 ) , MKL ( Intel , 2020a ) , oneDNN ( Intel , 2020b ) , ArrayFire ( Malcolm et al. , 2012 ) , and MiOpen ( Khan et al. , 2019 ) .

Summary: The paper proposed a minimal design API (or mostly API?) of machine learning framework called Flashlight. The key argument of the paper is that Flashlight is modular and agile. The Flashlight captured the key aspects of machine learning frameworks: Tensor and Operation, Memory Management, and Distributed. There are some evaluations to argue the benefit of Flashlight.