{ "version": "https://jsonfeed.org/version/1.1", "title": "Ramith Hettiarachchi — Scholarly Posts", "description": "Curated research notes, scholarly explainers, and independent findings.", "user_comment": "This is the blog's curated full-text scholarly feed.", "home_page_url": "https://new.ramith.fyi/", "feed_url": "https://new.ramith.fyi/feed.json", "language": "en", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi/" } ], "items": [ { "id": "https://new.ramith.fyi/posts/2024-02-12-an-open-science-journey/", "url": "https://new.ramith.fyi/posts/2024-02-12-an-open-science-journey/", "title": "An inspiring open science journey to remember 💙", "content_html": "

An inspiring open science journey to remember 💙

By Ramith Hettiarachchi ⋅ Published

Aya (Model & Dataset) is going to be released in 6 hours from now. Meanwhile, I thought of writing my thoughts on my journey and the things I have learned while collaborating with so many people across the world who had a common set of values and goals. I’ve been part of this effort since June 2023, and time flew so fast that I forgot how I even joined this in the first place. So this is a reflection for me to look back again at this day, from sometime in the future.

—

We all know that data is a cornerstone of AI advancements. But apart from English, most languages in the world have negligible representation on the Internet. Can we change this with the power of community collaboration? This was what Aya was all about. It was a massive global effort to uplift many under-resourced languages in the context of the current natural language advancement landscape.

Before elaborating on how I joined this project and the specifics of it, I want to think out loud about open science efforts. During the duration of this project, I saw,

All this unfolding on Discord was so inspiring to watch. And to finally see a multilingual large language model (LLM), proficient in 101 languages – and more importantly, seeing Aya-101 [1] proficient in my native language, Sinhala (සිංහල) too was such a joy 🎉❤️.

—

I asked these two questions from Aya in Sinhala: 1) වෙහෙසකර අභියෝගයක් සාර්ථකව ජය ගත් පසු කුමක්ද කරන්න ඕනේ? 2) ලංකාවෙ නිදහස් දිනය කවද්ද?
While there is a long way to go for Sinhala LLMs I’m honestly suprised that it follows Sinhala instructions quite nicely!

\"Aya-101
\"Another
\"Sinhala
(You can try the Aya-101 model here – https://huggingface.co/spaces/Tonic/Aya)

—

The Start

If I recall correctly, during my final year of undergrad, I joined C4AI (Cohere For AI\"Typeset) discord server. I saw so many wonderful initiatives by this community to help students/researchers who wanted to learn about AI and ultimately contribute to research. I didn’t follow these initiatives closely due to other time commitments. But every once in a while, I would check out the cool initiatives of this community.

\"TypesetCohere for AI is a non-profit driven by researchers all over the world with the mission of solving complex machine-learning problems.

—

Filling out a Google Form (May 2023)

\"A

On May 8th, I was getting ready for a morning run and saw this message on Discord.

\"Aya
\"Discord
(Suprisingly I had taken a screenshot of the Google Form (left picture) on that day, lucky for that, I can include it in this blog post!)

—

When filling something up, I always think whether it is something I can do given my other commitments\"Typeset 😅

I knew I didn’t have much time to dedicate, but the message above seemed reasonable so I filled it out. It seemed like only knowing the language was enough. Words likemultilingualandunderrepresented languagesmade it so tempting to fill out the Google form.

\"Typeset During undergrad, there are many occasions when I took too many things to my plate..

During high-school I used to make apps for Sinhala transliteration-link . So I have a softspot for tool that advance the usage of my native language.

—

An Email ✉️

June, 2023 was a great month, I was finishing up some experiments for my ICML 2023 workshop paper, and writing the paper. And then I received this email..

\"Email

it was sort of unexpected.. felt serious, but then again I felt like it’s something I’m capable of doing during my free time (a weekly meeting, spreading info about this project sounded good and doable 🤷🏻‍♂️).

So I replied that I’d want to help out to represent my native language! \"Typeset.

\"Typeset After all I was missing volunteering after completing my undergrad too… so felt that this would be nice..

Journey as a Language Ambassador

I attended the first meeting where I got introduced to Project Aya. The presentation was incredible… I thought to myself, “a great research group with a clear vision of what they want to achieve in 2023..”. There was so much positivity and hope within this community, and I loved that energy!

Somewhat of a rough start for Sinhala 🤕

Now it was my turn to contribute to my language.  I remember going to the Aya annotation platform\"Typeset, and filling out my details (such as languages I speak).

I was surprised to see Sinhala prompts and completions 😃, but soon I realized that they are machine translations which most often did not make much sense 🤕 (see left side of the below picture for some context).

\"Typeset Where human contributions are collected.

\"Sinhala

So, we gradually started to refine those. Soon with the help of many Sinhala contributors, we got a massive pool of prompts and completions that we kept on refining.

Realizing that we need to speed up

From June to the end of September, we gained a lot of contributions, but soon, we realized our pace was too slow 😬. So what we did was to visualize \"Typeset our goal and work towards it. Furthermore, we spread out the message through many contributors (kudos to Jalina, Chamod, Nawoda, and Chanuka)

\"Typeset Our goal was to reach 5000 re-annotations and 4000 original annotations

\"Discord
(a discord message from the Aya Server)

So we accelerated our pace! 🇱🇰

For the next remaining 3 months, all Sinhala contributors helped immensely to reach our goals – we even surpassed our original goals!

The pictures below capture some of the initiatives we took to gather more contributors..

\"Sinhala
\"Sinhala
\"Sinhala
\"Sinhala

—

\"Commemorative
\"Another
(a 1x1 inch chip fabricated with the names of Aya Sinhala Contributors as a memory)

References

  1. “Aya Model: An Instruction Finetuned Open-Access Multilingual Language Model.”
  2. “Aya Dataset: An Open-Access Collection for Multilingual Instruction Tuning.”
", "date_published": "2024-02-12T00:00:00Z", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi" } ], "language": "en", "tags": [ "scholarly" ] }, { "id": "https://new.ramith.fyi/posts/2023-06-24-differentiable-search-evolutionary-trees/", "url": "https://new.ramith.fyi/posts/2023-06-24-differentiable-search-evolutionary-trees/", "title": "Differentiable Search of Evolutionary Trees from Leaves — ICML 2023 Workshop Paper", "content_html": "

Differentiable Search of Evolutionary Trees from Leaves — ICML 2023 Workshop Paper

By Ramith Hettiarachchi ⋅ Published

ICML 2023 workshop paper

Differentiable Search of Evolutionary Trees from Leaves

Ramith Hettiarachchi, Avi Z. Swartz, Sergey Ovchinnikov

ICML 2023 workshops: Differentiable Almost Everything (DiffAE) and Sampling and Optimization in Discrete Space (SODS)

This page collects the workshop paper and its supporting resources. The paper contains the complete method, experiments, and claims.

Read paper ↗ View poster ↗ Browse code ↗ External explainer ↗

💡
In one sentence: We relax the discrete choices behind a phylogenetic tree so that its topology and unknown ancestral sequences can be improved together using gradients.

Related paper

Hettiarachchi, R., Swartz, A. Z., and Ovchinnikov, S. “Differentiable Search of Evolutionary Trees from Leaves.” ICML 2023 Workshops: DiffAE and SODS. DOI: 10.1101/2023.07.23.550206

", "date_published": "2023-06-24T00:00:00Z", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi" } ], "language": "en", "tags": [ "scholarly", "workshop paper", "phylogenetics", "differentiable programming" ] }, { "id": "https://new.ramith.fyi/posts/2023-02-10-esm-2-evolutionary-scale-prediction-of-atomic-level-protein-structure-with-a-language-model/", "url": "https://new.ramith.fyi/posts/2023-02-10-esm-2-evolutionary-scale-prediction-of-atomic-level-protein-structure-with-a-language-model/", "title": "ESM-2 (evolutionary-scale prediction of atomic level protein structure with a language model)", "content_html": "

ESM-2 (evolutionary-scale prediction of atomic level protein structure with a language model)

By Ramith Hettiarachchi ⋅ Published

✏️
I write these paper “summaries” for me to clearly understand the paper by summarizing the paper and synthesizing the literature. I hope they might be helpful to some readers too. If you have feedback, please write to hello@ramith.fyi.

Highlights of the ESM-2 Paper

💡

1. Introduction

1.1 — Structure and function are hidden in sequences

Biological properties of proteins influence which positions in a sequence can undergo mutations. From these observations, we can identify evolutionary patterns such as coevolution and conservation of amino acids. These patterns can help us infer properties of protein function and structure.[Note 4]

Usually, we align sequences before drawing conclusions about function and structure. This intermediate representation, known as a multiple sequence alignment (MSA), has a high time complexity because we must first search for related sequences and then align them.[Note 5]

What if we can get rid of this intermediate representation? That’s one aspect this paper accomplishes.

1.2 — Large language models (LLMs)

Historically, language models were pretrained using objectives such as predicting the next word in a sentence. Devlin et al.‘s BERT showed that masking some input tokens and predicting them—the masked-language-model objective, or MLM—is an effective pretraining strategy.[Note 6]

1.3 — Contributions

Inspired by this widely adopted strategy, the authors hypothesise that filling in missing amino acids might produce representations rich enough to infer structure. They therefore scale protein language models from 8 million parameters up to 15 billion parameters. Doing so reveals the following:

Because this approach improves speed by one to two orders of magnitude and does not need an MSA, the authors expand structure prediction to the much larger and more diverse space of metagenomic proteins. In summary, they:

Open ESM Metagenomic Atlas explorer in a new tab.

2. Method

2.1 — How does structure emerge from a language model trained on sequences?

The ESM-2 language model is trained with approximately 65 million unique sequences.[Note 8] Because the MLM objective asks the model to predict missing amino acids using their neighboring context, the model must learn interdependencies between amino acids. Previous work—[2] and [3]—showed that Transformer models trained with MLM on protein sequences develop attention patterns that correspond to residue-residue contact maps.[Note 9]

After training the language model, the authors use the approach from [3] to compute contact maps from attention patterns. A logistic regression identifies contacts as follows.

\"Annotated
Attention maps from a masked protein language model are symmetrized and passed to logistic regression to predict residue-residue contacts. Source figure: Rao et al. (2020) annotations added for this article.

2.2 — What about atomic-level structure? (Enter ESMFold)

The authors extract contact maps from attention patterns, but predicting atomic spatial coordinates requires an equivariant Transformer. They use the structure module introduced in AlphaFold. This module projects atomic spatial coordinates from the language model’s internal representation. The complete architecture is called ESMFold.

Steps in ESMFold

  1. Process the sequence through ESM-2.
  2. Pass the representation learned by ESM-2 through a series of folding blocks. Each block sequentially updates a sequence representation and a pairwise representation.
  3. Pass the result to the structure module.
  4. Repeat with three recycling steps. View annotated code.
\"ESMFold
ESMFold architecture: ESM-2 representations pass through a folding trunk and an AlphaFold-derived structure module, with recycling. Source figure: Lin et al. (2022) annotations added for this article.

The supporting figures connect the structure-module diagram to implementation details. Select any image to open the full-resolution version.

\"Annotated
AlphaFold’s structure module and invariant point attention. Source figures: Jumper et al. (2021).
\"Annotated
Mapping the folding trunk to its triangular self-attention block.
\"Annotated
Folding-block iteration, the structure-module call, and recycling.

Training: To train the structure model to obtain spatial coordinates, the authors use experimentally determined structures from the Protein Data Bank—approximately 25,000 clusters covering around 325,000 structures. This is augmented with 12 million structures predicted by AlphaFold2.[Note 10]

Evaluation: 194 CAMEO proteins and 51 CASP14 proteins.

This language model based approach vastly simplifies the usual SOTA structure prediction process by eliminating the need for the following \"Typeset,

For example, AlphaFold requires access to these resources.

3. Results

3.1 — How well does it predict structures?

As mentioned before, they evaluate performance on CAMEO and CASP14 proteins and check how well the structure was predicted using the TM-Score.

In predicting the structure just by single sequences, ESMFold achieves very good performance compared to AlphaFold and RoseTTAFold.

\"ESMFold,
Single-sequence and full-pipeline TM-score comparisons on CAMEO and CASP14; point color encodes perplexity. Source: Figure 2B in Lin et al. (2022) annotations added for this article.

3.2 — How important is the language model in the pipeline?

The key question that arises is how important is the representation learnt by the LM for the task of structure prediction . To quantify this we need several metrics.

First, we need to characterize how well the language model understands protein sequences. This is where perplexity comes in. We already have the TM-score to measure how closely a predicted structure matches the ground truth.

Thus, the graph to the right in Fig. 2B shows that,

✏️

What’s perplexity?

Perplexity describes how well a language model can predict a protein sequence. It ranges from 1 for a perfect model to 20 for random predictions; intuitively, it is the number of amino acids the model is choosing between at each prediction.

How can we achieve lower perplexity?

We now know that a better language model representation—with lower perplexity—leads to better structure prediction. How can we obtain a better representation? 🤔 Is scaling all you need?

To answer this question, authors explore the effect of scaling and look at what happens to the following :

\"Scatter
Effect of scaling ESM-2 on long-range contact precision; point color shows the change in perplexity. Source: Figure 1D in Lin et al. (2022) annotations added for this article.

The authors plot how long-range precision at L changes when moving from a smaller model on the x-axis to a larger model on the y-axis. Points above the diagonal suggest that scaling improves long-range precision at L for some proteins.

Is scaling the answer?

It is not that simple. Although P@L increases with scale for some proteins, the number of evolutionarily related sequences tells another story. Language models struggle when less relevant training data is available for a query. This is intuitive—more study, better results—but does it reflect memorization rather than understanding?

\"Long-range
Long-range precision at L as a function of the number of related sequences for several ESM-2 model sizes. Source: Lin et al. (2022) annotations added for this article.

The authors could have used a different color scheme: white points indicate no change in perplexity, while a point on the diagonal indicates no improvement in long-range precision at L.[Note 11]

3.3 — What about prediction speed?

\"Log-scale
Inference time versus sequence length. ESMFold is fastest for shorter sequences, while its pair representation makes long sequences more expensive. MSA search time for the other methods is not included.

3.4 — Comparison with other protein language models

\"Table
Validation perplexity, long-range contact precision, and structure-prediction results across ESM-2 sizes and comparison protein language models. Source: Lin et al. (2022) annotations added for this article.

4. Conclusion

It’s remarkable that these authors scale protein language models and it has resulted in learning structure hidden through databases of sequences, and thus we do not need to depend onto the MSA.

Is it because, the model has learnt to obtain the signal which we previously obtained through MSAs? What can we tell about the performance of sequences that had less number of evolutionary sequences in training data? why does it still struggle to obtain decent performance. It would be very interesting to analyze these directions.

Thanks for reading this, hope you found it useful. If you have any suggestions/ comments please share below.

References

  1. “Evolutionary-scale prediction of atomic level protein structure with a language model.”
  2. “BERTology Meets Biology: Interpreting Attention in Protein Language Models.”
  3. “Transformer protein language models are unsupervised structure learners.”
  4. “Eric Wallace: discussion of protein language models.”

Notes

  1. 1 The largest language model for proteins at the time.
  2. 2 Of these, 225 million were high-confidence predictions. See the ESM Metagenomic Atlas.
  3. 3 Metagenomics describes both sequencing DNA purified directly from a natural environment and the research field studying microbial communities in their natural state (Godzik, 2011).
  4. 4 In essence, information about protein structure and function is hidden in sequences.
  5. 5 The search process in AlphaFold and RoseTTAFold pipelines can take more than ten minutes.
  6. 6 Unlike left-to-right pretraining, the MLM objective lets the representation combine both left and right context, enabling a deeply bidirectional Transformer (paraphrased from the BERT paper).
  7. 7 This took two weeks on 2,000 GPUs. 🤯
  8. 8 The sequences were sampled with even weighting across approximately 43 million UniRef50 training clusters.
  9. 9 I plan to write a separate article about this work.
  10. 10 During training, predicted structures are sampled 75% of the time and real structures 25% of the time.
  11. 11 It would be useful if the plot made the proportion of proteins with improved performance directly visible, although that might make it too cluttered.
  12. 12 This is a 6× speedup compared with AlphaFold2.
", "date_published": "2023-02-10T00:00:00Z", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi" } ], "language": "en", "tags": [ "scholarly" ] }, { "id": "https://new.ramith.fyi/posts/2022-01-26-deep-residual-learning-for-image-recognition/", "url": "https://new.ramith.fyi/posts/2022-01-26-deep-residual-learning-for-image-recognition/", "title": "Deep Residual Learning for Image Recognition", "content_html": "

Deep Residual Learning for Image Recognition

By Ramith Hettiarachchi ⋅ Published

Highlights of ResNet Paper

Summary of He et al. [1].

Introduction

From the ImageNet Classification results in 2014-2015, it was evident that having deeper networks helps to learn greater levels of features. As shown in the figure below, we can see that VGG-Net and GoogLeNet have reduced the top-5 error rate further by having deeper networks.

So if we just stack more and more layers, does that help? Turns out it doesn’t. One reason for this is the vanishing gradient problem which was studied by Sepp Hochreiter in 1991 and discussed over the years [2]. This problem makes it difficult for a network to converge from the start. This issue has been addressed through various initialization methods and through batch normalization [3].

Even when a network starts converging, in deeper networks, researchers have found that there is a degradation of accuracy. Particularly, when we start increasing the depth of a model the accuracy gets saturated, and then it degrades rapidly . As  evident from the learning curves, this is not due to overfitting. Check the training-error curves in the left figure as depth increases in increments of 12.

—

[Note 1]

[Note 2]

—

The authors of the ResNet paper argue that, even if we increase the depth, theoretically there should be solution which gives the same accuracy. So it’s basically the shallow network + layers with identity transform.💡 Why can’t the added layers become an identity mapping 🤷🏻 However, the problem seems that the optimizers cannot reach that solution. So can we do a trick and get there easily? That’s what authors hypothesize.

Methodology - Deep Residual Learning

Fitting a residual mapping

\"Typeset - Mapping that needs to be fit by few stacked layers

\"Typeset - input to the first of those layers

Let’s say we need to approximate the function \"Typeset by some set of layers of a neural network. The authors propose that, rather than learning \"Typeset, the few layers approximate a residual function (\"Typeset). 🤔

Let’s denote this new function by \"Typeset. So, now we have rewritten original the function we need to approximate as, \"Typeset.

Ok, so what benefit does this give?  🤷🏻

By reformulating, we saw was that, \"Typeset was split into an addition of a function \"Typeset with the input. In the degradation problem that we saw earlier, the issue was that learning the identity function was hard. However with this residual learning reformulation, it should be easy for the optimizer to drive the weights of the layers such that \"Typeset becomes a zero mapping. Otherwise, deeper networks should have learnt the identity function in their added layers, giving similar accuracies. In this way, we are left with \"Typeset which is the identity mapping.

—

[Note 3]

References

  1. “Deep Residual Learning for Image Recognition.”
  2. “The fundamental deep learning problem.”
  3. “Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift.”

Notes

  1. Figure 1: Credit : Kaiming He, Xiangyu Zhang, Shaoqing Ren, & Jian Sun. “Deep Residual Learning for Image Recognition”. CVPR 2016 [Slides, Video]
  2. Figure 2: Credit : Kaiming He, Xiangyu Zhang, Shaoqing Ren, & Jian Sun. “Deep Residual Learning for Image Recognition”. CVPR 2016 [Slides, Video]
  3. Figure 3: Residual learning block. Source: He et al., “Deep Residual Learning for Image Recognition” [1].
", "date_published": "2022-01-26T00:00:00Z", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi" } ], "language": "en", "tags": [ "scholarly" ] }, { "id": "https://new.ramith.fyi/posts/2022-01-19-imagenet-classification-with-deep-convolutional-neural-networks/", "url": "https://new.ramith.fyi/posts/2022-01-19-imagenet-classification-with-deep-convolutional-neural-networks/", "title": "ImageNet classification with deep convolutional neural networks (AlexNet)", "content_html": "

ImageNet classification with deep convolutional neural networks (AlexNet)

By Ramith Hettiarachchi ⋅ Published

The AlexNet paper by Krizhevsky et al. [1] was published in 2012. It is a highly influential paper in computer vision which showed that deep networks, along with efficient utilization of GPUs for training, can help build better models. They achieved top-1 and top-5 test set error rates of 37.5% and 17.0%, compared with 47.1% and 28.2% for the previous best results in ILSVRC 2010. These are reductions of 9.6 and 11.2 percentage points, respectively.

Introduction

Before 2008, the computer vision researchers mostly evaluated their methods on datasets with tens of thousands of images. Thanks to the ImageNet dataset [2], which was published in 2009, researchers got the opportunity to move from small scale datasets such as MNIST, CIFAR-10/100 to a much larger dataset which has over 15 million labeled images from more than 22,000 categories. The AlexNet paper evaluates their method for ILSVRC-2010 & 2012 datasets.

The ImageNet Large-Scale Visual Recognition Challenge (ILSVRC) [3], which started in 2010 focuses on a subset of ImageNet which has 1.2 million training images, 50,000 validation images, and 150,000 testing images.

Architecture

Key characteristics of the AlexNet architecture : 8 learned layers (5 Conv, 3 Fully-connected)
The architecture has 60 million parameters.

[Note 1]

1. ReLU nonlinearity

Inspired by Nair and Hinton’s work [4], authors have utilized the ReLU non-linearity and have found that it makes the training process several times faster.

[Note 2]

2. Multiple GPU Training

Authors employ a cross-GPU parallelisation approach to train the AlexNet and the communication between GPUs happen only when it is required by the certain layers. They utilize 2x GTX 580 GPUs. Authors mention that this parallelism architecture is similar to “Columnar” CNN by Ciresan et al. [5].

[Note 3]

3. Local Response Normalization

Krizhevsky et al. has found that their local normalization methods helps achieve better generalization. For this procedure, they sum over \"Typeset adjacent kernel maps at the same spatial location as shown in the equation below. This normalization has been applied in certain layers of AlexNet. \"Typeset = total number of kernels
\"Typeset = response-normalized neuron activity

[Note 4]

4. Overlapping Pooling

By overlapping pooling scheme has reduced top-1 and top-5 error rates by 0.4% and 0.3%. Furthermore, the authors have observed that models with overlapping pooling are slightly more difficult to overfit.

Reducing Overfitting

Authors have utilized 1) data augmentations , 2) dropout to reduce overfitting. Data augmentation has been done in two ways.

1) Generating image translations and horizontal reflections

2) Altering RGB channel intensities in training images

Authors cite their early work on dropout [6] and mention that it has reduced overfitting substantially while roughly doubling the number of training iterations needed to converge.

References

  1. Krizhevsky et al. (2012). “ImageNet Classification with Deep Convolutional Neural Networks.”
  2. Deng et al. (2009). “ImageNet: A Large-Scale Hierarchical Image Database.”
  3. Russakovsky et al. (2014). “ImageNet Large Scale Visual Recognition Challenge.”
  4. Nair and Hinton (2010). “Rectified Linear Units Improve Restricted Boltzmann Machines.”
  5. Ciresan et al. (2012). “Multi-column Deep Neural Networks for Image Classification.”
  6. Hinton et al. (2012). “Improving neural networks by preventing co-adaptation of feature detectors.”

Notes

  1. Figure 1:
  2. Figure 2:
  3. Figure 3: Just to get an idea on how GPUs now and then (2012) compare with each other. Source : GadgetVersus
  4. Figure 4: Local reponse normalization has reduced top-1 and top-5 error rates by 1.4% and 1.2%.
", "date_published": "2022-01-19T00:00:00Z", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi" } ], "language": "en", "tags": [ "scholarly" ] }, { "id": "https://new.ramith.fyi/posts/2022-01-05-jax-001-a-closer-look-at-its-background/", "url": "https://new.ramith.fyi/posts/2022-01-05-jax-001-a-closer-look-at-its-background/", "title": "JAX 001 - A closer look at its background", "content_html": "

JAX 001 - A closer look at its background

By Ramith Hettiarachchi ⋅ Published

I recently wanted to get into parallel programming so that I can optimize one of the optics simulation projects I was working on (it was built without much focus on distributed training). So with all the buzz going around with JAX and some cool applications I found (listed below), I thought to give it a try.

Why JAX?

With all the established libraries such as PyTorch and Tensorflow, why is there a requirement for this new library in the first place? Let’s find out.

We all know that increasing FLOPS (floating point operations per second) is a huge deal in machine learning to train models efficiently. JAX aims to help this goal by enabling researchers to write python programs which are automatically compiled and scaled to utilize accelerators (GPUs/TPUs). Often it ishard to write optimized code in python to leverage the potential of hardware accelerators. JAX aims to keep a balance between research-friendly programming experience vs hardware acceleration.

To do so, JAX aims to accelerate pure-and-statically-composed (PSC) subroutines. It uses a just-in-time (JIT) compiler that traces these routines. To learn more about PSC subroutines, see the NeurIPS 2020: JAX Ecosystem Meetup video. The compilation process monitors the program once; the name JAX was described as “Just After eXecution” in the JAX paper [8].

Key features of JAX

JAX addresses several limitations in numpy. Therefore, it presents,

(autodiff, JIT compilation, vectorization, parallelization)

> JAX is Autograd and XLA, brought together for high-performance numerical computing and machine learning research. It provides composable transformations of Python+NumPy programs: differentiate, vectorize, parallelize, Just-In-Time compile to GPU/TPU, and more.

[Note 1]

Getting familiar with related concepts and history

There were a few loosly defined terms in my mind when I first read through the documentation. So I thought to look into those and have an idea about them.

Autograd - Autograd is an example of an automatic differentiation libraryreleased in 2014 . It is a lightweight tool to automatically differentiate native python and numpy code. But it doesn’t focus on hardware accelerators such GPU/TPUs. Therefore, since 2018, the main developers of Autograd are focusing on JAX which has more functionalities. Dougal, one of the main authors of autograd,calls JAX as the second generation of the project which started with Autograd.

“XLA - (Accelerated Linear Algebra) is a compiler-based linear algebra execution engine. It is the backend that powers machine learning frameworks such as TensorFlow and JAX at Google, on a variety of devices including CPUs, GPUs, and TPUs.” - (toread more )

JIT - (Just in time) compilation: This enables compilation of a JAX Python function so it can be executed efficiently in XLA.

Digging into these concepts made me realize that while the libraries we use makes it really easy to build end-to-end differentiable models, under the hood, these tools have been years-long effort by amazing teams with cool design choices.

In the next writeup I’m going to share my experience with device parallelization (pmap) and automatic vectorization (vmap) aspects.

—

References

  1. Jumper et al. (2021). “Highly accurate protein structure prediction with AlphaFold.”
  2. Schoenholz and Cubuk (2020). “JAX MD: A Framework for Differentiable Physics.”
  3. Freeman et al. (2021). “Brax -- A Differentiable Physics Engine for Large Scale Rigid Body Simulation.”
  4. JAXChem contributors “JAXChem: JAX-based chemical modeling software.”
  5. Kochkov et al. (2021). “Machine learning–accelerated computational fluid dynamics.”
  6. JAX-Cosmo contributors “JAX-Cosmo: A differentiable cosmology library in JAX.”
  7. “JAX documentation.”
  8. Frostig, Johnson, and Leary (2018). “Compiling machine learning programs via high-level tracing.”
  9. “JAX: accelerated ML research — video.”
  10. “JAX: accelerated ML research — slides.”
  11. “JAX at DeepMind — video.”
  12. “JAX at DeepMind — slides.”
  13. “Lecture 6: Automatic Differentiation.”
  14. “Cloud TPU Research Program.”

Notes

  1. Figure 1: Credit: {mattjj, frostig, leary, dougalm, phawkins, skyewm, jekbradbury, necula} @google.com. Full set of slides: [10].
", "date_published": "2022-01-05T00:00:00Z", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi" } ], "language": "en", "tags": [ "scholarly" ] }, { "id": "https://new.ramith.fyi/posts/2021-01-09-setting-up-raspberry-pi-4-with-ubuntu-20-04-ros-intel-realsense/", "url": "https://new.ramith.fyi/posts/2021-01-09-setting-up-raspberry-pi-4-with-ubuntu-20-04-ros-intel-realsense/", "title": "Getting the installation Right! Raspberry Pi 4 with Ubuntu 20.04 + ROS Noetic + Intel RealSense", "content_html": "

Getting the installation Right! Raspberry Pi 4 with Ubuntu 20.04 + ROS Noetic + Intel RealSense

By Ramith Hettiarachchi ⋅ Published

Historical guide: these steps describe my January 2021 setup with Ubuntu 20.04 and ROS Noetic.

Here I am after lot of iterations trying to install Ubuntu on Raspberry Pi 4(RPi4). Initially I made certain mistakes when selecting the correct ubuntu version & Desktop.

The steps below are which I followed to finalize my Rpi Installation with Ubuntu.

Ubuntu on RPi - Confusions?

My goal was to have a stable ROS installation on ubuntu on RPI4 . When considering the dependencies, I realized that RPI officially supports Ubuntu 20.04 & 20.10. And as of now, the ROS version that matches with one of those ubuntu versions (Ubuntu 20.04) is ROS Noetic.

Now in ubuntu’s official download page (For RPi), they have listed the following versions. However, Ubuntu 20.04 Desktop version is not listed here. Therefore, the goal should be to install Ubuntu Server 20.04 and then install a desktop of our choice.

[Note 1]

How this guide is organized

  1. Installing Ubuntu Server 20.04.1
  1. Installing the Desktop for Ubuntu Server
  2. Trying out screen sharing
  1. Installing ROS Noetic
  2. Installing Realsense libraries for Ubuntu 20.04

1. Installing Ubuntu Server 20.04.1

To install ubuntu server follow the official guide. When editing the file in microSD, double check whether you have edited it properly with correct wifi credentials

https://ubuntu.com/tutorials/how-to-install-ubuntu-on-your-raspberry-pi#1-overview

2. Installing the Desktop for Ubuntu Server

This is a tricky part. After searching for “how to install desktop on Ubuntu Server”, I tried the command ❌ sudo apt-get install ubuntu-desktop, which led to many errors in the desktop environment, such as Wi-Fi connections not displaying.

Therefore, we will be using a tool calledDesktopify . After you have installed ubuntu server, follow these steps to install the Desktop

git clone https://github.com/wimpysworld/desktopify.git\ncd desktopify\nsudo ./desktopify --de ubuntu

This will take some time to install. This script will take care of all the dependencies you need, and will install a clean desktop environment

[Note 2]

3. Trying out screensharing

Please follow the guidehere to enable screensharing so that you can view the RPi’s desktop remotely from your Laptop.

NOTE: To connect from your computer you will first need a monitor which is plugged to the RPI4. If you want to access the screen remotely even when a monitor is not attached , Check the very end of this post.

[Note 3]

4. Installing ROS Noetic

You can follow these steps, to install ROS Noetic

http://wiki.ros.org/noetic/Installation/Ubuntu

5. Installing Realsense libraries

My friend Shalutha suggested these steps to install the required dependencies of Realsense.

RealSense installation instructions on ROS Answers.

You can copy paste these steps to an .sh file and execute it

nano realsense.sh\n#copy paste the steps listed in the link to this file\nchmod +x realsense.sh\n./realsense.sh

[Note 4]

Congratulations 🥳

Now you have a stable ubuntu desktop 20.04 installation!

/opt/realsense/bin/realsense-viewer

[Note 5]

Extras

1. Logging in with an SSH key

Follow the SSH key setup tutorial linked in the original guide.

2. Accessing remote screen when a monitor is not attached.

sudo apt-get install xserver-xorg-video-dummy
sudo nano /etc/X11/xorg.conf\n

Copy the following configuration into that file:

Section "Device"\n    Identifier "Configured Video Device"\n    Driver "dummy"\nEndSection\n\nSection "Monitor"\n    Identifier "Configured Monitor"\n    HorizSync 31.5-48.5\n    VertRefresh 50-70\nEndSection\n\nSection "Screen"\n    Identifier "Default Screen"\n    Monitor "Configured Monitor"\n    Device "Configured Video Device"\n    DefaultDepth 24\n    SubSection "Display"\n        Depth 24\n        Modes "1024x800"\n    EndSubSection\nEndSection

Now reboot the RPisudo reboot. Then you can connect from your laptop even if there is no monitor attached to the RPi.

(source:https://askubuntu.com/questions/453109/add-fake-display-when-no-monitor-is-plugged-in )

Final outcome of setting up 🎉

Watch the original setup demo on YouTube.

Notes

  1. Figure 1: source: https://ubuntu.com/download/raspberry-pi
  2. Figure 2: After Desktopifying Ubuntu Server 20.04.1
  3. Figure 3:
  4. Figure 4: (It will take some time to install)
  5. Figure 5: (Screenshot taken from the plugged monitor)
", "date_published": "2021-01-09T00:00:00Z", "authors": [ { "name": "Ramith Hettiarachchi", "url": "https://new.ramith.fyi" } ], "language": "en", "tags": [ "scholarly" ] } ] }